From 5225bb53df5ec4cf1c4686362ab99860111f440d Mon Sep 17 00:00:00 2001
From: squalus <squalus@squalus.net>
Date: Thu, 14 Dec 2023 00:48:44 -0800
Subject: [PATCH] remote: add NarFromPath client

---
 hnix-store-remote/hnix-store-remote.cabal     |  3 +-
 .../src/System/Nix/Store/Remote/Arbitrary.hs  |  1 +
 .../src/System/Nix/Store/Remote/Client.hs     | 13 +++++++
 .../System/Nix/Store/Remote/Client/Core.hs    | 34 ++++++++++++++++++
 .../src/System/Nix/Store/Remote/MonadStore.hs | 36 +++++++++++++++++++
 .../src/System/Nix/Store/Remote/Serializer.hs |  8 ++++-
 .../src/System/Nix/Store/Remote/Server.hs     |  1 +
 .../Nix/Store/Remote/Types/StoreRequest.hs    |  6 ++++
 hnix-store-remote/tests-io/DataSink.hs        | 26 ++++++++++++++
 hnix-store-remote/tests-io/NixDaemonSpec.hs   | 22 +++++++++++-
 hnix-store-remote/tests-io/SampleNar.hs       |  7 ++--
 11 files changed, 151 insertions(+), 6 deletions(-)
 create mode 100644 hnix-store-remote/tests-io/DataSink.hs

diff --git a/hnix-store-remote/hnix-store-remote.cabal b/hnix-store-remote/hnix-store-remote.cabal
index bf901955..9633b7d8 100644
--- a/hnix-store-remote/hnix-store-remote.cabal
+++ b/hnix-store-remote/hnix-store-remote.cabal
@@ -196,7 +196,8 @@ test-suite remote-io
   -- See https://github.com/redneb/hs-linux-namespaces/issues/3
   ghc-options:       -rtsopts -fprof-auto "-with-rtsopts -V0"
   other-modules:
-      NixDaemonSpec
+      DataSink
+    , NixDaemonSpec
     , SampleNar
   build-depends:
       base >=4.12 && <5
diff --git a/hnix-store-remote/src/System/Nix/Store/Remote/Arbitrary.hs b/hnix-store-remote/src/System/Nix/Store/Remote/Arbitrary.hs
index 0ab81b95..700d502d 100644
--- a/hnix-store-remote/src/System/Nix/Store/Remote/Arbitrary.hs
+++ b/hnix-store-remote/src/System/Nix/Store/Remote/Arbitrary.hs
@@ -106,6 +106,7 @@ instance Arbitrary (Some StoreRequest) where
     , Some . EnsurePath <$> arbitrary
     , pure $ Some FindRoots
     , Some . IsValidPath <$> arbitrary
+    , Some . NarFromPath <$> arbitrary
     , Some <$> (QueryValidPaths <$> arbitrary <*> arbitrary)
     , pure $ Some QueryAllValidPaths
     , Some . QuerySubstitutablePaths <$> arbitrary
diff --git a/hnix-store-remote/src/System/Nix/Store/Remote/Client.hs b/hnix-store-remote/src/System/Nix/Store/Remote/Client.hs
index f1874eac..28914c45 100644
--- a/hnix-store-remote/src/System/Nix/Store/Remote/Client.hs
+++ b/hnix-store-remote/src/System/Nix/Store/Remote/Client.hs
@@ -11,6 +11,7 @@ module System.Nix.Store.Remote.Client
   , ensurePath
   , findRoots
   , isValidPath
+  , narFromPath
   , queryValidPaths
   , queryAllValidPaths
   , querySubstitutablePaths
@@ -181,6 +182,18 @@ isValidPath
   -> m Bool
 isValidPath = doReq . IsValidPath
 
+-- | Download a NAR file.
+narFromPath
+  :: MonadRemoteStore m
+  => StorePath -- ^ Path to generate a NAR for
+  -> Word64 -- ^ Byte length of NAR
+  -> (ByteString -> IO()) -- ^ Data sink where NAR bytes will be written
+  -> m ()
+narFromPath path narSize sink = do
+  setDataSink sink
+  setDataSinkSize narSize
+  void $ doReq (NarFromPath path)
+
 -- | Query valid paths from a set,
 -- optionally try to use substitutes
 queryValidPaths
diff --git a/hnix-store-remote/src/System/Nix/Store/Remote/Client/Core.hs b/hnix-store-remote/src/System/Nix/Store/Remote/Client/Core.hs
index e10b9c3b..5373dd19 100644
--- a/hnix-store-remote/src/System/Nix/Store/Remote/Client/Core.hs
+++ b/hnix-store-remote/src/System/Nix/Store/Remote/Client/Core.hs
@@ -91,6 +91,22 @@ doReq = \case
         processOutput
         pure NoReply
 
+      NarFromPath _ -> do
+        maybeSink <- getDataSink
+        sink <- case maybeSink of
+          Nothing -> throwError RemoteStoreError_NoDataSinkProvided
+          Just sink -> pure sink
+        clearDataSink
+        maybeNarSize <- getDataSinkSize
+        narSize <- case maybeNarSize of
+          Nothing -> throwError RemoteStoreError_NoDataSinkSizeProvided
+          Just narSize -> pure narSize
+        clearDataSinkSize
+        soc <- getStoreSocket
+        processOutput
+        copyToSink sink narSize soc
+        pure NoReply
+
       _ -> do
         processOutput
         processReply
@@ -101,6 +117,24 @@ doReq = \case
             $ getReplyS @a
           )
 
+copyToSink
+  :: forall m
+   . ( MonadIO m
+     , MonadRemoteStore m
+     )
+  => (ByteString -> IO()) --  ^ data sink
+  -> Word64 -- ^ byte length to read
+  -> Socket
+  -> m ()
+copyToSink sink remainingBytes soc =
+  when (remainingBytes > 0) $ do
+    let chunkSize = 16384
+        bytesToRead = min chunkSize remainingBytes
+    bytes <- liftIO $ Network.Socket.ByteString.recv soc (fromIntegral bytesToRead)
+    liftIO $ sink bytes
+    let nextRemainingBytes = remainingBytes - (fromIntegral . Data.ByteString.length) bytes
+    copyToSink sink nextRemainingBytes soc
+
 writeFramedSource
   :: forall m
    . ( MonadIO m
diff --git a/hnix-store-remote/src/System/Nix/Store/Remote/MonadStore.hs b/hnix-store-remote/src/System/Nix/Store/Remote/MonadStore.hs
index 0d4263a2..ab147b04 100644
--- a/hnix-store-remote/src/System/Nix/Store/Remote/MonadStore.hs
+++ b/hnix-store-remote/src/System/Nix/Store/Remote/MonadStore.hs
@@ -47,6 +47,8 @@ data RemoteStoreState = RemoteStoreState {
   , remoteStoreStateMDataSink :: Maybe (ByteString -> IO ())
   -- ^ Sink for @Logger_Write@, called repeatedly by the daemon
   -- to dump us some data. Used by @ExportPath@ operation.
+  , remoteStoreStateMDataSinkSize :: Maybe Word64
+  -- ^ Byte length to be written to the sink, for NarForPath
   , remoteStoreStateMNarSource :: Maybe (NarSource IO)
   }
 
@@ -80,6 +82,7 @@ data RemoteStoreError
   | RemoteStoreError_DataSourceZeroLengthRead -- remoteStoreStateMDataSource returned a zero length ByteString
   | RemoteStoreError_DataSourceReadTooLarge -- remoteStoreStateMDataSource returned a ByteString larger than the chunk size requested or the remaining bytes
   | RemoteStoreError_NoDataSinkProvided -- remoteStoreStateMDataSink is required but it is Nothing
+  | RemoteStoreError_NoDataSinkSizeProvided -- remoteStoreStateMDataSinkSize is required but it is Nothing
   | RemoteStoreError_NoNarSourceProvided
   | RemoteStoreError_OperationFailed
   | RemoteStoreError_ProtocolMismatch
@@ -148,6 +151,7 @@ runRemoteStoreT sock =
       , remoteStoreStateLogs = mempty
       , remoteStoreStateMDataSource = Nothing
       , remoteStoreStateMDataSink = Nothing
+      , remoteStoreStateMDataSinkSize = Nothing
       , remoteStoreStateMNarSource = Nothing
       }
 
@@ -307,6 +311,34 @@ class ( MonadIO m
     => m ()
   clearDataSink = lift clearDataSink
 
+  setDataSinkSize :: Word64 -> m ()
+  default setDataSinkSize
+   :: ( MonadTrans t
+      , MonadRemoteStore m'
+      , m ~ t m'
+      )
+   => Word64
+   -> m ()
+  setDataSinkSize x = lift (setDataSinkSize x)
+
+  getDataSinkSize :: m (Maybe Word64)
+  default getDataSinkSize
+   :: ( MonadTrans t
+      , MonadRemoteStore m'
+      , m ~ t m'
+      )
+   => m (Maybe Word64)
+  getDataSinkSize = lift getDataSinkSize
+
+  clearDataSinkSize :: m ()
+  default clearDataSinkSize
+    :: ( MonadTrans t
+       , MonadRemoteStore m'
+       , m ~ t m'
+       )
+    => m ()
+  clearDataSinkSize = lift clearDataSinkSize
+
 instance MonadRemoteStore m => MonadRemoteStore (StateT s m)
 instance MonadRemoteStore m => MonadRemoteStore (ReaderT r m)
 instance MonadRemoteStore m => MonadRemoteStore (ExceptT RemoteStoreError m)
@@ -347,6 +379,10 @@ instance MonadIO m => MonadRemoteStore (RemoteStoreT m) where
   getDataSink = RemoteStoreT (gets remoteStoreStateMDataSink)
   clearDataSink = RemoteStoreT $ modify $ \s -> s { remoteStoreStateMDataSink = Nothing }
 
+  setDataSinkSize x = RemoteStoreT $ modify $ \s -> s { remoteStoreStateMDataSinkSize = pure x }
+  getDataSinkSize = RemoteStoreT (gets remoteStoreStateMDataSinkSize)
+  clearDataSinkSize = RemoteStoreT $ modify $ \s -> s { remoteStoreStateMDataSinkSize = Nothing }
+
   setNarSource x = RemoteStoreT $ modify $ \s -> s { remoteStoreStateMNarSource = pure x }
   takeNarSource = RemoteStoreT $ do
     x <- remoteStoreStateMNarSource <$> get
diff --git a/hnix-store-remote/src/System/Nix/Store/Remote/Serializer.hs b/hnix-store-remote/src/System/Nix/Store/Remote/Serializer.hs
index df7b4602..b661bf33 100644
--- a/hnix-store-remote/src/System/Nix/Store/Remote/Serializer.hs
+++ b/hnix-store-remote/src/System/Nix/Store/Remote/Serializer.hs
@@ -1135,6 +1135,9 @@ storeRequest = Serializer
       WorkerOp_IsValidPath -> mapGetE $ do
         Some . IsValidPath <$> getS storePath
 
+      WorkerOp_NarFromPath -> mapGetE $ do
+        Some . NarFromPath <$> getS storePath
+
       WorkerOp_QueryValidPaths -> mapGetE $ do
         paths <- getS (hashSet storePath)
         substituteMode <- getS enum
@@ -1191,7 +1194,6 @@ storeRequest = Serializer
       w@WorkerOp_ExportPath -> notYet w
       w@WorkerOp_HasSubstitutes -> notYet w
       w@WorkerOp_ImportPaths -> notYet w
-      w@WorkerOp_NarFromPath -> notYet w
       w@WorkerOp_QueryDerivationOutputMap -> notYet w
       w@WorkerOp_QueryDeriver -> notYet w
       w@WorkerOp_QueryFailedPaths -> notYet w
@@ -1280,6 +1282,10 @@ storeRequest = Serializer
         putS workerOp WorkerOp_IsValidPath
         putS storePath path
 
+      Some (NarFromPath path) -> mapPutE $ do
+        putS workerOp WorkerOp_NarFromPath
+        putS storePath path
+
       Some (QueryValidPaths paths substituteMode) -> mapPutE $ do
         putS workerOp WorkerOp_QueryValidPaths
 
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 fb469e83..47557b73 100644
--- a/hnix-store-remote/src/System/Nix/Store/Remote/Server.hs
+++ b/hnix-store-remote/src/System/Nix/Store/Remote/Server.hs
@@ -181,6 +181,7 @@ processConnection workerHelper postGreet sock = do
             r@EnsurePath {} -> perform r
             r@FindRoots {} -> perform r
             r@IsValidPath {} -> perform r
+            r@NarFromPath {} -> perform r
             r@QueryValidPaths {} -> perform r
             r@QueryAllValidPaths {} -> perform r
             r@QuerySubstitutablePaths {} -> perform r
diff --git a/hnix-store-remote/src/System/Nix/Store/Remote/Types/StoreRequest.hs b/hnix-store-remote/src/System/Nix/Store/Remote/Types/StoreRequest.hs
index 872713df..c4aeea75 100644
--- a/hnix-store-remote/src/System/Nix/Store/Remote/Types/StoreRequest.hs
+++ b/hnix-store-remote/src/System/Nix/Store/Remote/Types/StoreRequest.hs
@@ -104,6 +104,11 @@ data StoreRequest :: Type -> Type where
     :: StorePath
     -> StoreRequest Bool
 
+  -- | Fetch a NAR from the server
+  NarFromPath
+    :: StorePath
+    -> StoreRequest NoReply
+
   -- | Query valid paths from set, optionally try to use substitutes.
   QueryValidPaths
     :: HashSet StorePath
@@ -179,6 +184,7 @@ instance {-# OVERLAPPING #-} Eq (Some StoreRequest) where
   Some (EnsurePath a) == Some (EnsurePath a') = a == a'
   Some (FindRoots) == Some (FindRoots) = True
   Some (IsValidPath a) == Some (IsValidPath a') = a == a'
+  Some (NarFromPath a) == Some (NarFromPath a') = a == a'
   Some (QueryValidPaths a b) == Some (QueryValidPaths a' b') = (a, b) == (a', b')
   Some QueryAllValidPaths == Some QueryAllValidPaths = True
   Some (QuerySubstitutablePaths a) == Some (QuerySubstitutablePaths a') = a == a'
diff --git a/hnix-store-remote/tests-io/DataSink.hs b/hnix-store-remote/tests-io/DataSink.hs
new file mode 100644
index 00000000..fd592e67
--- /dev/null
+++ b/hnix-store-remote/tests-io/DataSink.hs
@@ -0,0 +1,26 @@
+module DataSink
+
+( DataSink(..)
+, dataSinkResult
+, dataSinkWriter
+, newDataSink
+)
+
+where
+
+import Data.ByteString (ByteString)
+
+import Control.Monad.ST
+import Data.STRef
+
+-- | Basic data sink for testing
+newtype DataSink = DataSink (STRef RealWorld ByteString)
+
+newDataSink :: IO DataSink
+newDataSink = DataSink <$> (stToIO . newSTRef) mempty
+
+dataSinkWriter :: DataSink -> (ByteString -> IO())
+dataSinkWriter (DataSink stref) chunk = stToIO (modifySTRef stref (chunk <>))
+
+dataSinkResult :: DataSink -> IO ByteString
+dataSinkResult (DataSink stref) = (stToIO . readSTRef) stref
diff --git a/hnix-store-remote/tests-io/NixDaemonSpec.hs b/hnix-store-remote/tests-io/NixDaemonSpec.hs
index 9b41fbad..c93e2fae 100644
--- a/hnix-store-remote/tests-io/NixDaemonSpec.hs
+++ b/hnix-store-remote/tests-io/NixDaemonSpec.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE NamedFieldPuns #-}
 
 module NixDaemonSpec
   ( enterNamespaces
@@ -35,6 +36,7 @@ import qualified Data.Map
 import qualified Data.Set
 import qualified Data.Text
 import qualified Data.Text.Encoding
+import qualified DataSink
 import qualified SampleNar
 import qualified System.Directory
 import qualified System.Environment
@@ -264,6 +266,9 @@ itLefts
   -> SpecWith (m () -> IO (Either a b, c))
 itLefts name action = it name action Data.Either.isLeft
 
+sampleText :: Text
+sampleText = "test"
+
 withPath
   :: MonadRemoteStore m
   => (StorePath -> m a)
@@ -273,7 +278,7 @@ withPath action = do
     addTextToStore
       (StoreText
         (forceRight $ System.Nix.StorePath.mkStorePathName "hnix-store")
-        "test"
+        sampleText
       )
       mempty
       RepairMode_DontRepair
@@ -341,6 +346,7 @@ makeProtoSpec
   -> SpecFlavor
   -> Spec
 makeProtoSpec f flavor = around f $ do
+
   context "syncWithGC" $
     itRights "syncs with garbage collector" syncWithGC
 
@@ -499,3 +505,17 @@ makeProtoSpec f flavor = around f $ do
 
         meta <- queryPathInfo sampleNar_storePath
         (metadataDeriverPath =<< meta) `shouldBe` metadataDeriverPath sampleNar_metadata
+
+  context "narFromPath" $ do
+    itRights "downloads nar file" $ do
+      unless (flavor == SpecFlavor_MITM) $ do
+        withPath $ \path -> do
+          maybeMetadata <- queryPathInfo path
+          case maybeMetadata of
+            Just Metadata{metadataNarBytes=Just narBytes} -> do
+              dataSink <- liftIO DataSink.newDataSink
+              narFromPath path narBytes (DataSink.dataSinkWriter dataSink)
+              narData <- liftIO $ DataSink.dataSinkResult dataSink
+              expectedNarData <- liftIO $ SampleNar.encodeNar (Data.Text.Encoding.encodeUtf8 sampleText)
+              narData `shouldBe` expectedNarData
+            _ -> expectationFailure "missing metadata or narBytes"
diff --git a/hnix-store-remote/tests-io/SampleNar.hs b/hnix-store-remote/tests-io/SampleNar.hs
index e4c1307d..c6ea2915 100644
--- a/hnix-store-remote/tests-io/SampleNar.hs
+++ b/hnix-store-remote/tests-io/SampleNar.hs
@@ -6,6 +6,7 @@ module SampleNar
 ( SampleNar(..)
 , buildDataSource
 , sampleNar0
+, encodeNar
 )
 
 where
@@ -39,7 +40,7 @@ data SampleNar
 sampleNar0 :: IO SampleNar
 sampleNar0 = do
   let sampleNar_fileData = "hello"
-  sampleNar_narData <- bytesToNar sampleNar_fileData
+  sampleNar_narData <- encodeNar sampleNar_fileData
   let sampleNar_metadata = Metadata
         { metadataDeriverPath = Just $ forceParsePath "/nix/store/g2mxdrkwr1hck4y5479dww7m56d1x81v-hello-2.12.1.drv"
         , metadataNarHash = sha256 sampleNar_narData
@@ -78,8 +79,8 @@ forceParsePath path = case parsePath def path of
 sha256 :: ByteString -> DSum HashAlgo Digest
 sha256 bs = HashAlgo_SHA256 :=> hashFinalize (hashUpdate (hashInit @SHA256) bs)
 
-bytesToNar :: ByteString -> IO ByteString
-bytesToNar bytes = do
+encodeNar :: ByteString -> IO ByteString
+encodeNar bytes = do
   ref <- stToIO $ newSTRef mempty
   let accumFn chunk = do
        stToIO $ modifySTRef ref (<> chunk)