remote: add AddToStoreNar client

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

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.

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
+        processOutput
+        processReply
 
-      _ -> pure ()
+      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
 
-    processOutput
-    sockGetS
-      (mapErrorS RemoteStoreError_SerializerReply
-        $ getReplyS @a
-      )
+      _ -> 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

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 }

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)

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

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)
+

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
 

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'

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

hnix-store-remote/tests-io/SampleNar.hs

@@ -0,0 +1,88 @@
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE NamedFieldPuns #-}
+
+module SampleNar
+
+( SampleNar(..)
+, buildDataSource
+, sampleNar0
+)
+
+where
+
+import Crypto.Hash ( Digest, SHA256, hashInit, hashUpdate, hashFinalize )
+import Data.ByteString (ByteString)
+import Data.Dependent.Sum (DSum((:=>)))
+import Data.Time.Calendar.OrdinalDate (fromOrdinalDate)
+import Data.Time.Clock (UTCTime(..))
+import System.Nix.Hash (HashAlgo(HashAlgo_SHA256))
+import System.Nix.StorePath (StorePath, parsePath)
+import System.Nix.StorePath.Metadata (Metadata(..), StorePathTrust(..))
+
+import Control.Monad.ST
+import Data.Default.Class
+import Data.STRef
+import Data.Word
+
+import qualified Data.ByteString
+import qualified System.Nix.Nar
+
+-- | Sample data for an AddToStoreNar operation
+data SampleNar
+  = SampleNar
+  { sampleNar_fileData :: !ByteString -- ^ Contents of the file to be put in the store
+  , sampleNar_narData :: !ByteString -- ^ The file data serialized as a NAR
+  , sampleNar_storePath :: !StorePath
+  , sampleNar_metadata :: !(Metadata StorePath)
+  }
+
+sampleNar0 :: IO SampleNar
+sampleNar0 = do
+  let sampleNar_fileData = "hello"
+  sampleNar_narData <- bytesToNar sampleNar_fileData
+  let sampleNar_metadata = Metadata
+        { metadataDeriverPath = Just $ forceParsePath "/nix/store/g2mxdrkwr1hck4y5479dww7m56d1x81v-hello-2.12.1.drv"
+        , metadataNarHash = sha256 sampleNar_narData
+        , metadataReferences = mempty
+        , metadataRegistrationTime = UTCTime (fromOrdinalDate 1980 1) 0
+        , metadataNarBytes = Just ((fromIntegral . Data.ByteString.length) sampleNar_narData)
+        , metadataTrust = BuiltElsewhere
+        , metadataSigs = mempty
+        , metadataContentAddress = Nothing
+        }
+      sampleNar_storePath = forceParsePath "/nix/store/00000lj3clbkc0aqvjjzfa6slp4zdvlj-hello-2.12.1"
+  pure SampleNar{..}
+
+buildDataSource :: SampleNar -> IO(Word64 -> IO (Maybe ByteString))
+buildDataSource SampleNar{sampleNar_narData} = dataSourceFromByteString sampleNar_narData
+
+dataSourceFromByteString :: ByteString -> IO (Word64 -> IO(Maybe ByteString))
+dataSourceFromByteString bs = do
+  posRef <- stToIO $ newSTRef (0::Word64)
+  let len = fromIntegral $ Data.ByteString.length bs
+      outFn chunkSize = do
+        pos <- stToIO $ readSTRef posRef
+        if pos >= len then pure Nothing
+        else do
+          let bs' = Data.ByteString.drop (fromIntegral pos) bs
+              bs'' = Data.ByteString.take (fromIntegral chunkSize) bs'
+              takenLen = fromIntegral $ Data.ByteString.length bs''
+          stToIO $ modifySTRef posRef (takenLen +)
+          pure $ Just bs''
+  pure outFn
+
+forceParsePath :: ByteString -> StorePath
+forceParsePath path = case parsePath def path of
+  Left err -> error $ mconcat [ "forceParsePath failed: ", show err ]
+  Right x -> x
+sha256 :: ByteString -> DSum HashAlgo Digest
+sha256 bs = HashAlgo_SHA256 :=> hashFinalize (hashUpdate (hashInit @SHA256) bs)
+
+bytesToNar :: ByteString -> IO ByteString
+bytesToNar bytes = do
+  ref <- stToIO $ newSTRef mempty
+  let accumFn chunk = do
+       stToIO $ modifySTRef ref (<> chunk)
+  System.Nix.Nar.dumpString bytes accumFn
+  stToIO $ readSTRef ref
+