From 46519800470231157fc05ed79c34a10ea2246301 Mon Sep 17 00:00:00 2001
From: sorki <srk@48.io>
Date: Sun, 10 Dec 2023 16:47:04 +0100
Subject: [PATCH] remote: NixDaemonSpec MITM

---
 hnix-store-remote/hnix-store-remote.cabal   |  1 +
 hnix-store-remote/tests-io/NixDaemonSpec.hs | 37 +++++++++++++++++++--
 2 files changed, 35 insertions(+), 3 deletions(-)

diff --git a/hnix-store-remote/hnix-store-remote.cabal b/hnix-store-remote/hnix-store-remote.cabal
index 831b820e..b4c8c73d 100644
--- a/hnix-store-remote/hnix-store-remote.cabal
+++ b/hnix-store-remote/hnix-store-remote.cabal
@@ -202,6 +202,7 @@ test-suite remote-io
     , hnix-store-remote
     , hnix-store-tests
     , bytestring
+    , concurrency
     , containers
     , crypton
     , directory
diff --git a/hnix-store-remote/tests-io/NixDaemonSpec.hs b/hnix-store-remote/tests-io/NixDaemonSpec.hs
index 171aaf6b..13cac764 100644
--- a/hnix-store-remote/tests-io/NixDaemonSpec.hs
+++ b/hnix-store-remote/tests-io/NixDaemonSpec.hs
@@ -5,8 +5,9 @@ module NixDaemonSpec
   , spec
   ) where
 
-import Control.Monad (forM_, unless, void)
+import Control.Monad (forM_, unless, void, (<=<))
 import Control.Monad.Catch (MonadMask)
+import Control.Monad.Conc.Class (MonadConc)
 import Control.Monad.IO.Class (MonadIO, liftIO)
 import Crypto.Hash (SHA256)
 import Data.Some (Some(Some))
@@ -22,6 +23,7 @@ import System.Nix.DerivedPath (DerivedPath(..))
 import System.Nix.StorePath (StoreDir(..), StorePath)
 import System.Nix.StorePath.Metadata (Metadata(..))
 import System.Nix.Store.Remote
+import System.Nix.Store.Remote.Server (WorkerHelper)
 import System.Process (CreateProcess(..), ProcessHandle)
 import qualified Control.Concurrent
 import qualified Control.Exception
@@ -186,12 +188,40 @@ withNixDaemon
   -> IO a
 withNixDaemon action =
   withNixDaemon' $ \_tmpPath storeDir storeConn ->
-    action $ \a ->
+    action $ \(mstore :: RemoteStoreT m a) ->
       runStoreConnection storeConn
         ( setStoreDir storeDir
-          >> a
+          >> mstore
         )
 
+withManInTheMiddleNixDaemon
+  :: forall m a
+   . ( MonadIO m
+     , MonadMask m
+     , MonadConc m
+     )
+  => ((RemoteStoreT m a -> Run m a) -> IO a)
+  -> IO a
+withManInTheMiddleNixDaemon action =
+  withNixDaemon' $ \tmpPath storeDir storeConn ->
+  let
+    sockFp2 = tmpPath </> "var/nix/daemon-socket/socket2"
+    storeConn2 = StoreConnection_Socket $ StoreSocketPath sockFp2
+
+    handler :: WorkerHelper m
+    handler = either (error . show) pure
+      <=< fmap fst
+      . runStoreConnection storeConn
+      . (setStoreDir storeDir >>)
+      . doReq
+
+  in action $ \(mstore :: RemoteStoreT m a) ->
+    runDaemonConnection handler storeConn2
+      $ runStoreConnection storeConn2
+          ( setStoreDir storeDir
+            >> mstore
+          )
+
 checks
   :: ( Show a
      , Show b
@@ -289,6 +319,7 @@ spec :: Spec
 spec = do
   describe "Remote store protocol" $ do
     describe "Direct" $ makeProtoSpec withNixDaemon
+    describe "MITM" $ makeProtoSpec withManInTheMiddleNixDaemon
 
 makeProtoSpec
   :: (ActionWith