From 3d774b81873a7e153bf42e873306c4635a95e6a4 Mon Sep 17 00:00:00 2001
From: sorki <srk@48.io>
Date: Sun, 10 Dec 2023 17:50:02 +0100
Subject: [PATCH] remote/server: add proxyNarSource

---
 .../src/System/Nix/Store/Remote.hs            |  4 +--
 .../src/System/Nix/Store/Remote/Server.hs     | 32 +++++++++++++++++--
 hnix-store-remote/tests-io/NixDaemonSpec.hs   |  1 -
 3 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/hnix-store-remote/src/System/Nix/Store/Remote.hs b/hnix-store-remote/src/System/Nix/Store/Remote.hs
index fab180b..e43e2d7 100644
--- a/hnix-store-remote/src/System/Nix/Store/Remote.hs
+++ b/hnix-store-remote/src/System/Nix/Store/Remote.hs
@@ -108,9 +108,9 @@ justdoit = do
         def
         "/nix/store/yyznqbwam67cmp7zfwk0rkgmi9yqsdsm-hnix-store-core-0.8.0.0"
 
-    handler :: MonadIO m => WorkerHelper m
+    handler :: RemoteStoreT IO a -> IO a
     handler k = do
-      x <- liftIO $ runStore $ doReq k
+      x <- runStore k
       either (error . show) pure (fst x)
 
 runDaemon
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 0a17331..d55d604 100644
--- a/hnix-store-remote/src/System/Nix/Store/Remote/Server.hs
+++ b/hnix-store-remote/src/System/Nix/Store/Remote/Server.hs
@@ -22,6 +22,8 @@ import Data.Word (Word32)
 import qualified Data.Text
 import qualified Data.Text.IO
 import Network.Socket (Socket, accept, close, listen, maxListenQueue)
+import System.Nix.Nar (NarSource)
+import System.Nix.Store.Remote.Client (doReq)
 import System.Nix.Store.Remote.Serializer (LoggerSError, mapErrorS, storeRequest, workerMagic, protoVersion, int, logger, text, trustedFlag)
 import System.Nix.Store.Remote.Socket
 import System.Nix.Store.Remote.Types.StoreRequest as R
@@ -34,13 +36,16 @@ import System.Nix.Store.Remote.Types.WorkerMagic (WorkerMagic(..))
 
 -- wip
 import Data.Some (withSome)
+import qualified System.Timeout
+import qualified Network.Socket.ByteString
 
 type WorkerHelper m
   = forall a
   . ( Show a
     , StoreReply a
     )
-  => StoreRequest a -> m a
+  => RemoteStoreT m a
+  -> m a
 
 chatty :: Bool
 chatty = False
@@ -119,7 +124,30 @@ processConnection workerHelper postGreet sock = do
           => StoreRequest a
           -> RemoteStoreT m ()
         perform req = do
-          resp <- bracketLogger tunnelLogger $ lift $ workerHelper req
+
+          special <- case req of
+            AddToStore {} -> do
+              let proxyNarSource :: NarSource IO
+                  proxyNarSource f =
+                    liftIO
+                      (System.Timeout.timeout
+                         1000000
+                         (Network.Socket.ByteString.recv sock 8)
+                      )
+                    >>= \case
+                      Nothing -> pure ()
+                      Just x -> f x >> proxyNarSource f
+
+              pure $ setNarSource proxyNarSource
+            _ -> pure $ pure ()
+
+          resp <-
+            bracketLogger
+              tunnelLogger
+              $ lift
+              $ workerHelper
+              $ special >> doReq req
+
           sockPutS
             (mapErrorS
                RemoteStoreError_SerializerReply
diff --git a/hnix-store-remote/tests-io/NixDaemonSpec.hs b/hnix-store-remote/tests-io/NixDaemonSpec.hs
index 1daa5bb..3f578a7 100644
--- a/hnix-store-remote/tests-io/NixDaemonSpec.hs
+++ b/hnix-store-remote/tests-io/NixDaemonSpec.hs
@@ -213,7 +213,6 @@ withManInTheMiddleNixDaemon action =
       <=< fmap fst
       . runStoreConnection storeConn
       . (setStoreDir storeDir >>)
-      . doReq
 
   in action $ \(mstore :: RemoteStoreT m a) ->
     runDaemonConnection handler