Merge pull request #279 from squalus/narfrompath

remote: add NarFromPath client

Author: sorki

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

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

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

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

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

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
 

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

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'

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

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"