remote: add AddToStoreNar client

squalus · squalus · commit 5494cc3edd47 · 2023-12-12T19:05:07.000-08:00
diff --git a/hnix-store-remote/hnix-store-remote.cabal b/hnix-store-remote/hnix-store-remote.cabal
@@ -85,6 +85,7 @@ library
     , System.Nix.Store.Remote.Types.GC
     , System.Nix.Store.Remote.Types.Handshake
     , System.Nix.Store.Remote.Types.Logger
+    , System.Nix.Store.Remote.Types.NoReply
     , System.Nix.Store.Remote.Types.ProtoVersion
     , System.Nix.Store.Remote.Types.Query
     , System.Nix.Store.Remote.Types.Query.Missing
@@ -196,6 +197,7 @@ test-suite remote-io
   ghc-options:       -rtsopts -fprof-auto "-with-rtsopts -V0"
   other-modules:
       NixDaemonSpec
+    , SampleNar
   build-depends:
       base >=4.12 && <5
     , hnix-store-core
@@ -206,6 +208,8 @@ test-suite remote-io
     , concurrency
     , containers
     , crypton
+    , data-default-class
+    , dependent-sum
     , directory
     , exceptions
     , filepath
@@ -216,5 +220,6 @@ test-suite remote-io
     , some
     , temporary
     , text
+    , time
     , unix
     , unordered-containers
diff --git a/hnix-store-remote/src/System/Nix/Store/Remote/Client.hs b/hnix-store-remote/src/System/Nix/Store/Remote/Client.hs
@@ -1,5 +1,6 @@
 module System.Nix.Store.Remote.Client
   ( addToStore
+  , addToStoreNar
   , addTextToStore
   , addSignatures
   , addTempRoot
@@ -28,10 +29,12 @@ module System.Nix.Store.Remote.Client
 
 import Control.Monad (void, when)
 import Control.Monad.Except (throwError)
+import Data.ByteString (ByteString)
 import Data.HashSet (HashSet)
 import Data.Map (Map)
 import Data.Set (Set)
 import Data.Some (Some)
+import Data.Word (Word64)
 
 import System.Nix.Build (BuildMode, BuildResult)
 import System.Nix.DerivedPath (DerivedPath)
@@ -73,6 +76,19 @@ addToStore name source method hashAlgo repair = do
   setNarSource source
   doReq (AddToStore name method hashAlgo repair)
 
+addToStoreNar
+  :: MonadRemoteStore m
+  => StorePath
+  -> Metadata StorePath
+  -> RepairMode
+  -> CheckMode
+  -> (Word64 -> IO(Maybe ByteString))
+  -> m ()
+addToStoreNar path metadata repair checkSigs source = do
+  setDataSource source
+  void $ doReq (AddToStoreNar path metadata repair checkSigs)
+  pure ()
+
 -- | Add @StoreText@ to the store
 -- Reference accepts repair but only uses it
 -- to throw error in case of remote talking to nix-daemon.
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
@@ -7,9 +7,13 @@ module System.Nix.Store.Remote.Client.Core
 import Control.Monad (unless, when)
 import Control.Monad.Except (throwError)
 import Control.Monad.IO.Class (MonadIO, liftIO)
+import Data.ByteString (ByteString)
 import Data.DList (DList)
 import Data.Some (Some(Some))
+import Data.Word (Word64)
+import Network.Socket (Socket)
 import System.Nix.Nar (NarSource)
+import System.Nix.StorePath.Metadata (Metadata(..))
 import System.Nix.Store.Remote.Logger (processOutput)
 import System.Nix.Store.Remote.MonadStore
   ( MonadRemoteStore(..)
@@ -28,11 +32,13 @@ import System.Nix.Store.Remote.Serializer
   )
 import System.Nix.Store.Remote.Types.Handshake (ClientHandshakeOutput(..))
 import System.Nix.Store.Remote.Types.Logger (Logger)
+import System.Nix.Store.Remote.Types.NoReply (NoReply(..))
 import System.Nix.Store.Remote.Types.ProtoVersion (ProtoVersion(..))
 import System.Nix.Store.Remote.Types.StoreRequest (StoreRequest(..))
 import System.Nix.Store.Remote.Types.StoreReply (StoreReply(..))
 import System.Nix.Store.Remote.Types.WorkerMagic (WorkerMagic(..))
 
+import qualified Data.ByteString
 import qualified Network.Socket.ByteString
 
 type Run m a = m (Either RemoteStoreError a, DList Logger)
@@ -69,14 +75,58 @@ doReq = \case
           Nothing ->
             throwError
               RemoteStoreError_NoNarSourceProvided
-
-      _ -> pure ()
-
-    processOutput
-    sockGetS
-      (mapErrorS RemoteStoreError_SerializerReply
-        $ getReplyS @a
-      )
+        processOutput
+        processReply
+
+      AddToStoreNar _ meta _ _ -> do
+        let narBytes = maybe 0 id $ metadataNarBytes meta
+        maybeDataSource <- takeDataSource
+        soc <- getStoreSocket
+        case maybeDataSource of
+          Nothing ->
+            if narBytes == 0 then writeFramedSource (const (pure Nothing)) soc 0
+            else throwError RemoteStoreError_NoDataSourceProvided
+          Just dataSource -> do
+            writeFramedSource dataSource soc narBytes
+        processOutput
+        pure NoReply
+
+      _ -> do
+        processOutput
+        processReply
+
+  where
+    processReply = sockGetS
+          (mapErrorS RemoteStoreError_SerializerReply
+            $ getReplyS @a
+          )
+
+writeFramedSource
+  :: forall m
+   . ( MonadIO m
+     , MonadRemoteStore m
+     )
+  => (Word64 -> IO(Maybe ByteString))
+  -> Socket
+  -> Word64
+  -> m ()
+writeFramedSource dataSource soc remainingBytes = do
+  let chunkSize = 16384
+  maybeBytes <- liftIO $ dataSource chunkSize
+  case maybeBytes of
+    Nothing -> do
+      unless (remainingBytes == 0) $ throwError RemoteStoreError_DataSourceExhausted
+      let eof :: Word64 = 0
+      sockPutS int eof
+    Just bytes -> do
+      let bytesInChunk = fromIntegral $ Data.ByteString.length bytes
+      when (bytesInChunk > chunkSize || bytesInChunk > remainingBytes) $ throwError RemoteStoreError_DataSourceReadTooLarge
+      when (bytesInChunk == 0) $ throwError RemoteStoreError_DataSourceZeroLengthRead
+      sockPutS int bytesInChunk
+      liftIO
+        $ Network.Socket.ByteString.sendAll soc bytes
+      let nextRemainingBytes = remainingBytes - bytesInChunk
+      writeFramedSource dataSource soc nextRemainingBytes
 
 greetServer
   :: MonadRemoteStore 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
@@ -77,6 +77,8 @@ data RemoteStoreError
   | RemoteStoreError_LoggerParserFail String ByteString -- when incremental parser returns ((Fail msg leftover) :: Result)
   | RemoteStoreError_NoDataSourceProvided -- remoteStoreStateMDataSource is required but it is Nothing
   | RemoteStoreError_DataSourceExhausted -- remoteStoreStateMDataSource returned Nothing but more data was requested
+  | 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_NoNarSourceProvided
   | RemoteStoreError_OperationFailed
@@ -250,6 +252,15 @@ class ( MonadIO m
    -> m ()
   setDataSource x = lift (setDataSource x)
 
+  takeDataSource :: m (Maybe (Word64 -> IO (Maybe ByteString)))
+  default takeDataSource
+   :: ( MonadTrans t
+      , MonadRemoteStore m'
+      , m ~ t m'
+      )
+   => m (Maybe (Word64 -> IO (Maybe ByteString)))
+  takeDataSource = lift takeDataSource
+
   getDataSource :: m (Maybe (Word64 -> IO (Maybe ByteString)))
   default getDataSource
    :: ( MonadTrans t
@@ -327,6 +338,11 @@ instance MonadIO m => MonadRemoteStore (RemoteStoreT m) where
   getDataSource = RemoteStoreT (gets remoteStoreStateMDataSource)
   clearDataSource = RemoteStoreT $ modify $ \s -> s { remoteStoreStateMDataSource = Nothing }
 
+  takeDataSource = RemoteStoreT $ do
+    x <- remoteStoreStateMDataSource <$> get
+    modify $ \s -> s { remoteStoreStateMDataSource = Nothing }
+    pure x
+
   setDataSink x = RemoteStoreT $ modify $ \s -> s { remoteStoreStateMDataSink = pure x }
   getDataSink = RemoteStoreT (gets remoteStoreStateMDataSink)
   clearDataSink = RemoteStoreT $ modify $ \s -> s { remoteStoreStateMDataSink = Nothing }
diff --git a/hnix-store-remote/src/System/Nix/Store/Remote/Serializer.hs b/hnix-store-remote/src/System/Nix/Store/Remote/Serializer.hs
@@ -85,6 +85,7 @@ module System.Nix.Store.Remote.Serializer
   -- ** Reply
   , ReplySError(..)
   , opSuccess
+  , noop
   -- *** Realisation
   , derivationOutputTyped
   , realisation
@@ -1077,6 +1078,16 @@ storeRequest = Serializer
 
         pure $ Some (AddToStore pathName recursive hashAlgo repair)
 
+      WorkerOp_AddToStoreNar -> mapGetE $ do
+        storePath' <- getS storePath
+        metadata <- getS pathMetadata
+        repair <- getS bool
+        let repairMode = if repair then RepairMode_DoRepair else RepairMode_DontRepair
+        dontCheckSigs <- getS bool
+        let checkSigs = if dontCheckSigs then CheckMode_DontCheck else CheckMode_DoCheck
+
+        pure $ Some (AddToStoreNar storePath' metadata repairMode checkSigs)
+
       WorkerOp_AddTextToStore -> mapGetE $ do
         txt <- getS storeText
         paths <- getS (hashSet storePath)
@@ -1175,7 +1186,6 @@ storeRequest = Serializer
 
       w@WorkerOp_AddBuildLog -> notYet w
       w@WorkerOp_AddMultipleToStore -> notYet w
-      w@WorkerOp_AddToStoreNar -> notYet w
       w@WorkerOp_BuildPathsWithResults -> notYet w
       w@WorkerOp_ClearFailedPaths -> notYet w
       w@WorkerOp_ExportPath -> notYet w
@@ -1207,6 +1217,14 @@ storeRequest = Serializer
         putS bool (recursive == FileIngestionMethod_FileRecursive)
         putS someHashAlgo hashAlgo
 
+      Some (AddToStoreNar storePath' metadata repair checkSigs) -> mapPutE $ do
+        putS workerOp WorkerOp_AddToStoreNar
+
+        putS storePath storePath'
+        putS pathMetadata metadata
+        putS bool $ repair == RepairMode_DoRepair
+        putS bool $ checkSigs == CheckMode_DontCheck
+
       Some (AddTextToStore txt paths _repair) -> mapPutE $ do
         putS workerOp WorkerOp_AddTextToStore
 
@@ -1379,6 +1397,12 @@ opSuccess = Serializer
   , putS = \_ -> mapPutER $ putS bool True
   }
 
+noop :: a -> NixSerializer r ReplySError a
+noop ret = Serializer
+  { getS = pure ret
+  , putS = \_ -> pure ()
+  }
+
 -- *** Realisation
 
 derivationOutputTyped :: NixSerializer r ReplySError (System.Nix.Realisation.DerivationOutput OutputName)
diff --git a/hnix-store-remote/src/System/Nix/Store/Remote/Server.hs b/hnix-store-remote/src/System/Nix/Store/Remote/Server.hs
@@ -170,6 +170,7 @@ processConnection workerHelper postGreet sock = do
           -- out of thin air
           () <- Data.Some.withSome someReq $ \case
             r@AddToStore {} -> perform r
+            r@AddToStoreNar {} -> perform r
             r@AddTextToStore {} -> perform r
             r@AddSignatures {} -> perform r
             r@AddTempRoot {} -> perform r
diff --git a/hnix-store-remote/src/System/Nix/Store/Remote/Types/NoReply.hs b/hnix-store-remote/src/System/Nix/Store/Remote/Types/NoReply.hs
@@ -0,0 +1,8 @@
+module System.Nix.Store.Remote.Types.NoReply
+  ( NoReply(..)
+  ) where
+
+-- | Reply type for the case where the server does not reply
+data NoReply = NoReply
+  deriving (Show)
+
diff --git a/hnix-store-remote/src/System/Nix/Store/Remote/Types/StoreReply.hs b/hnix-store-remote/src/System/Nix/Store/Remote/Types/StoreReply.hs
@@ -8,6 +8,7 @@ import System.Nix.Build (BuildResult)
 import System.Nix.StorePath (StorePath, StorePathName)
 import System.Nix.StorePath.Metadata (Metadata)
 import System.Nix.Store.Remote.Serializer
+import System.Nix.Store.Remote.Types.NoReply (NoReply(..))
 import System.Nix.Store.Remote.Types.SuccessCodeReply (SuccessCodeReply)
 import System.Nix.Store.Remote.Types.GC (GCResult, GCRoot)
 import System.Nix.Store.Remote.Types.Query.Missing (Missing)
@@ -24,6 +25,9 @@ class StoreReply a where
 instance StoreReply SuccessCodeReply where
   getReplyS = opSuccess
 
+instance StoreReply NoReply where
+  getReplyS = noop NoReply
+
 instance StoreReply Bool where
   getReplyS = mapPrimE bool
 
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
@@ -25,6 +25,7 @@ import System.Nix.StorePath (StorePath, StorePathName, StorePathHashPart)
 import System.Nix.StorePath.Metadata (Metadata)
 import System.Nix.Store.Remote.Types.GC (GCOptions, GCResult, GCRoot)
 import System.Nix.Store.Remote.Types.CheckMode (CheckMode)
+import System.Nix.Store.Remote.Types.NoReply (NoReply)
 import System.Nix.Store.Remote.Types.Query.Missing (Missing)
 import System.Nix.Store.Remote.Types.StoreText (StoreText)
 import System.Nix.Store.Remote.Types.SubstituteMode (SubstituteMode)
@@ -39,6 +40,14 @@ data StoreRequest :: Type -> Type where
     -> RepairMode -- ^ Only used by local store backend
     -> StoreRequest StorePath
 
+  -- | Add a NAR with Metadata to the store.
+  AddToStoreNar
+    :: StorePath
+    -> Metadata StorePath
+    -> RepairMode
+    -> CheckMode -- ^ Whether to check signatures
+    -> StoreRequest NoReply
+
   -- | Add text to store.
   --
   -- Reference accepts repair but only uses it
@@ -159,6 +168,7 @@ deriveGShow ''StoreRequest
 
 instance {-# OVERLAPPING #-} Eq (Some StoreRequest) where
   Some (AddToStore a b c d) == Some (AddToStore a' b' c' d') = (a, b, c, d) == (a', b', c', d')
+  Some (AddToStoreNar a b c d) == Some (AddToStoreNar a' b' c' d') = (a, b, c, d) == (a', b', c', d')
   Some (AddTextToStore a b c) == Some (AddTextToStore a' b' c') = (a, b, c) == (a', b', c')
   Some (AddSignatures a b) == Some (AddSignatures a' b') = (a, b) == (a', b')
   Some (AddIndirectRoot a) == Some (AddIndirectRoot a') = a == a'
diff --git a/hnix-store-remote/tests-io/NixDaemonSpec.hs b/hnix-store-remote/tests-io/NixDaemonSpec.hs
@@ -35,6 +35,7 @@ import qualified Data.Map
 import qualified Data.Set
 import qualified Data.Text
 import qualified Data.Text.Encoding
+import qualified SampleNar
 import qualified System.Directory
 import qualified System.Environment
 import qualified System.IO.Temp
@@ -488,3 +489,13 @@ makeProtoSpec f flavor = around f $ do
               }
         gcResultDeletedPaths `shouldBe` Data.HashSet.fromList [path]
         gcResultBytesFreed `shouldBe` 4
+
+  context "addToStoreNar" $ do
+    itRights "adds nar file" $ do
+      unless (flavor == SpecFlavor_MITM) $ do
+        sampleNar@SampleNar.SampleNar{..} <- liftIO SampleNar.sampleNar0
+        dataSource <- liftIO $ SampleNar.buildDataSource sampleNar
+        addToStoreNar sampleNar_storePath sampleNar_metadata RepairMode_DontRepair CheckMode_DontCheck dataSource
+
+        meta <- queryPathInfo sampleNar_storePath
+        (metadataDeriverPath =<< meta) `shouldBe` metadataDeriverPath sampleNar_metadata
diff --git a/hnix-store-remote/tests-io/SampleNar.hs b/hnix-store-remote/tests-io/SampleNar.hs
