From 3dc7ab6f0614679057410c6ace375132f486b12e Mon Sep 17 00:00:00 2001
From: sorki <srk@48.io>
Date: Sun, 10 Dec 2023 17:06:44 +0100
Subject: [PATCH] remote: add postGreet so we setStoreDir in Server as well

---
 hnix-store-remote/src/System/Nix/Store/Remote.hs    | 13 ++++++++-----
 .../src/System/Nix/Store/Remote/Server.hs           | 11 ++++++++---
 hnix-store-remote/tests-io/NixDaemonSpec.hs         |  4 +++-
 3 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/hnix-store-remote/src/System/Nix/Store/Remote.hs b/hnix-store-remote/src/System/Nix/Store/Remote.hs
index 6b38841..fab180b 100644
--- a/hnix-store-remote/src/System/Nix/Store/Remote.hs
+++ b/hnix-store-remote/src/System/Nix/Store/Remote.hs
@@ -94,7 +94,7 @@ runStoreSocket sockFamily sockAddr code =
 
 justdoit :: Run IO (Bool, Bool)
 justdoit = do
-  runDaemonConnection handler (StoreConnection_Socket "/tmp/dsock") $
+  runDaemonConnection handler (pure ()) (StoreConnection_Socket "/tmp/dsock") $
     runStoreConnection (StoreConnection_Socket "/tmp/dsock")
       $ do
         a <- isValidPath pth
@@ -124,6 +124,7 @@ runDaemon
 runDaemon workerHelper =
   runDaemonConnection
     workerHelper
+    (pure ())
     def
 
 -- | Run an emulated nix daemon using given @StoreConnection@
@@ -134,14 +135,15 @@ runDaemonConnection
     , MonadConc m
     )
   => WorkerHelper m
+  -> RemoteStoreT m ()
   -> StoreConnection
   -> m a
   -> m a
-runDaemonConnection workerHelper sc k =
+runDaemonConnection workerHelper postGreet sc k =
   connectionToSocket sc
   >>= \case
     Left e -> error $ show e
-    Right (fam, sock) -> runDaemonSocket workerHelper fam sock k
+    Right (fam, sock) -> runDaemonSocket workerHelper postGreet fam sock k
 
 -- | Run an emulated nix daemon using given @StoreConnection@
 -- the deamon will close when the continuation returns.
@@ -151,11 +153,12 @@ runDaemonSocket
     , MonadConc m
     )
   => WorkerHelper m
+  -> RemoteStoreT m ()
   -> Family
   -> SockAddr
   -> m a
   -> m a
-runDaemonSocket workerHelper sockFamily sockAddr k =
+runDaemonSocket workerHelper postGreet sockFamily sockAddr k =
   Control.Monad.Catch.bracket
     (liftIO
       $ Network.Socket.socket
@@ -177,7 +180,7 @@ runDaemonSocket workerHelper sockFamily sockAddr k =
 
     -- set up the listening socket
     liftIO $ Network.Socket.bind lsock sockAddr
-    runProxyDaemon workerHelper lsock k
+    runProxyDaemon workerHelper postGreet lsock k
 
 connectionToSocket
   :: MonadIO m
diff --git a/hnix-store-remote/src/System/Nix/Store/Remote/Server.hs b/hnix-store-remote/src/System/Nix/Store/Remote/Server.hs
index 313f35d..6124075 100644
--- a/hnix-store-remote/src/System/Nix/Store/Remote/Server.hs
+++ b/hnix-store-remote/src/System/Nix/Store/Remote/Server.hs
@@ -50,10 +50,11 @@ runProxyDaemon
     , MonadConc m
     )
   => WorkerHelper m
+  -> RemoteStoreT m ()
   -> Socket
   -> m a
   -> m a
-runProxyDaemon workerHelper lsock k = do
+runProxyDaemon workerHelper postGreet lsock k = do
   liftIO $ listen lsock maxListenQueue
 
   liftIO $ Data.Text.IO.putStrLn "listening"
@@ -66,7 +67,7 @@ runProxyDaemon workerHelper lsock k = do
         -- TODO: this, but without the space leak
         fmap fst
           $ concurrently listener
-          $ processConnection workerHelper sock
+          $ processConnection workerHelper postGreet sock
 
   either absurd id <$> race listener k
 
@@ -77,9 +78,10 @@ processConnection
   :: forall m
   .  MonadIO m
   => WorkerHelper m
+  -> RemoteStoreT m ()
   -> Socket
   -> m ()
-processConnection workerHelper sock = do
+processConnection workerHelper postGreet sock = do
   ~() <- void $ runRemoteStoreT sock $ do
 
     ServerHandshakeOutput{..}
@@ -101,6 +103,9 @@ processConnection workerHelper sock = do
     --authHook(*store);
     stopWork tunnelLogger
 
+    -- so we can set store dir
+    postGreet
+
     let perform
           :: ( Show a
              , StoreReply a
diff --git a/hnix-store-remote/tests-io/NixDaemonSpec.hs b/hnix-store-remote/tests-io/NixDaemonSpec.hs
index 13cac76..1daa5bb 100644
--- a/hnix-store-remote/tests-io/NixDaemonSpec.hs
+++ b/hnix-store-remote/tests-io/NixDaemonSpec.hs
@@ -216,7 +216,9 @@ withManInTheMiddleNixDaemon action =
       . doReq
 
   in action $ \(mstore :: RemoteStoreT m a) ->
-    runDaemonConnection handler storeConn2
+    runDaemonConnection handler
+      (setStoreDir storeDir)
+      storeConn2
       $ runStoreConnection storeConn2
           ( setStoreDir storeDir
             >> mstore