hnix-store-remote-0.8.0.0: src/System/Nix/Store/Remote/Server.hs
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE Rank2Types #-}
module System.Nix.Store.Remote.Server
( runProxyDaemon
, WorkerHelper
)
where
import Algebra.PartialOrd (leq)
import Control.Concurrent.Classy.Async
import Control.Monad (join, unless, void, when)
import Control.Monad.Conc.Class (MonadConc)
import Control.Monad.Except (MonadError, throwError)
import Control.Monad.Trans (lift)
import Control.Monad.IO.Class (MonadIO, liftIO)
import Data.Default.Class (Default(def))
import Data.Foldable (traverse_)
import Data.IORef (IORef, atomicModifyIORef, newIORef)
import Data.Set qualified
import Data.Text (Text)
import Data.Void (Void, absurd)
import Data.ByteString (ByteString)
import Data.Word (Word32, Word64)
import Network.Socket (Socket, accept, close, listen, maxListenQueue)
import System.Nix.Nar (NarSource)
import System.Nix.Store.Remote.Client (Run, doReq)
import System.Nix.Store.Remote.Serializer (LoggerSError, mapErrorS, storeRequest, validPathInfo, workerMagic, protoFeatures, protoVersion, int, logger, text, trustedFlag)
import System.Nix.Store.Remote.Socket
import System.Nix.Store.Remote.Types.StoreRequest as R
import System.Nix.Store.Remote.Types.StoreReply
import System.Nix.Store.Remote.Types.ProtoVersion (ProtoVersion(..), minVersionNumber)
import System.Nix.Store.Remote.Types.Logger (ErrorInfo, Logger(..))
import System.Nix.Store.Remote.MonadStore (MonadRemoteStore(..), WorkerError(..), WorkerException(..), RemoteStoreError(..), RemoteStoreT, runRemoteStoreT)
import System.Nix.Store.Remote.Types.Handshake (ServerHandshakeInput(..), ServerHandshakeOutput(..))
import System.Nix.Store.Remote.Types.WorkerMagic (WorkerMagic(..))
import Data.Some qualified
import Data.Text.IO qualified
import Data.Bits (shiftL)
import Data.ByteString qualified
import Network.Socket.ByteString qualified
type WorkerHelper m
= forall a
. ( Show a
, StoreReply a
)
=> RemoteStoreT m a
-> Run m a
chatty :: Bool
chatty = False
dbg :: MonadIO m => Text -> m ()
dbg = when chatty . liftIO . Data.Text.IO.putStrLn
-- | Run an emulated nix daemon on given socket address.
-- The deamon will close when the continuation returns.
runProxyDaemon
:: forall m a
. ( MonadIO m
, MonadConc m
)
=> WorkerHelper m
-> RemoteStoreT m ()
-> Socket
-> m a
-> m a
runProxyDaemon workerHelper postGreet lsock k = do
liftIO $ listen lsock maxListenQueue
dbg "listening"
let listener :: m Void
listener = do
(sock, _) <- liftIO $ accept lsock
dbg "accepting"
-- TODO: this, but without the space leak
fmap fst
$ concurrently listener
$ processConnection workerHelper postGreet sock
either absurd id <$> race listener k
-- | "main loop" of the daemon for a single connection.
--
-- this function should take care to not throw errors from client connections.
processConnection
:: forall m
. MonadIO m
=> WorkerHelper m
-> RemoteStoreT m ()
-> Socket
-> m ()
processConnection workerHelper postGreet sock = do
~() <- void $ runRemoteStoreT sock $ do
ServerHandshakeOutput{..}
<- greet
ServerHandshakeInput
{ serverHandshakeInputNixVersion = "nixVersion (hnix-store-remote)"
, serverHandshakeInputOurVersion = def
, serverHandshakeInputTrust = Nothing
}
setProtoVersion serverHandshakeOutputLeastCommonVersion
tunnelLogger <- liftIO $ newTunnelLogger
-- Send startup error messages to the client.
startWork tunnelLogger
-- TODO: do we need auth at all? probably?
-- If we can't accept clientVersion, then throw an error *here* (not above).
--authHook(*store);
stopWork tunnelLogger
-- so we can set store dir
postGreet
let perform
:: ( Show a
, StoreReply a
)
=> StoreRequest a
-> RemoteStoreT m ()
perform req = do
special <- case req of
AddToStore {} -> do
let proxyNarSource :: NarSource IO
proxyNarSource sink = readFramedSource sock sink
pure $ setNarSource proxyNarSource
AddToStoreScanning {} -> do
let proxyNarSource :: NarSource IO
proxyNarSource sink = readFramedSource sock sink
pure $ setNarSource proxyNarSource
_ -> pure $ pure ()
res <-
bracketLogger
tunnelLogger
$ lift
$ workerHelper
$ special >> doReq req
case fst res of
Left e -> throwError e
Right reply -> do
storeDir <- getStoreDir
pv <- getProtoVersion
case req of
-- AddToStore returns ValidPathInfo on the wire,
-- not just StorePath
AddToStore {} -> do
-- Query the metadata so we can send full ValidPathInfo
metadata <- lift $ workerHelper $ doReq (QueryPathInfo reply)
case fst metadata of
Left e -> throwError e
Right (Just meta) ->
sockPutS
(mapErrorS RemoteStoreError_SerializerPut
$ validPathInfo storeDir)
(reply, meta)
Right Nothing ->
throwError $ RemoteStoreError_Fixme "AddToStore: path not found after adding"
AddToStoreScanning {} -> do
metadata <- lift $ workerHelper $ doReq (QueryPathInfo reply)
case fst metadata of
Left e -> throwError e
Right (Just meta) ->
sockPutS
(mapErrorS RemoteStoreError_SerializerPut
$ validPathInfo storeDir)
(reply, meta)
Right Nothing ->
throwError $ RemoteStoreError_Fixme "AddToStoreScanning: path not found after adding"
_ ->
sockPutS
(mapErrorS
RemoteStoreError_SerializerReply
$ getReplyS storeDir pv
)
reply
-- Process client requests.
let loop = do
storeDir <- getStoreDir
pv <- getProtoVersion
someReq <-
sockGetS
$ mapErrorS
RemoteStoreError_SerializerRequest
(storeRequest storeDir pv)
-- have to be explicit here
-- because otherwise GHC can't conjure Show a, StoreReply a
-- out of thin air
() <- Data.Some.withSome someReq $ \case
r@AddToStore {} -> perform r
r@AddToStoreNar {} -> perform r
r@AddToStoreScanning {} -> perform r
r@AddTextToStore {} -> perform r
r@AddSignatures {} -> perform r
r@AddTempRoot {} -> perform r
r@AddIndirectRoot {} -> perform r
r@BuildDerivation {} -> perform r
r@BuildPaths {} -> perform r
r@CollectGarbage {} -> perform r
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
r@QueryPathInfo {} -> perform r
r@QueryReferrers {} -> perform r
r@QueryValidDerivers {} -> perform r
r@QueryDerivationOutputs {} -> perform r
r@QueryDerivationOutputNames {} -> perform r
r@QueryPathFromHashPart {} -> perform r
r@QueryMissing {} -> perform r
r@OptimiseStore {} -> perform r
r@SubmitOutput {} -> perform r
r@SyncWithGC {} -> perform r
r@VerifyStore {} -> perform r
loop
loop
dbg "daemon connection done"
liftIO $ close sock
where
-- Exchange the greeting.
greet
:: MonadIO m
=> ServerHandshakeInput
-> RemoteStoreT m ServerHandshakeOutput
greet ServerHandshakeInput{..} = do
magic <-
sockGetS
$ mapErrorS
RemoteStoreError_SerializerHandshake
workerMagic
when (magic /= WorkerMagic_One)
$ throwError
$ RemoteStoreError_WorkerException
WorkerException_ProtocolMismatch
sockPutS
(mapErrorS
RemoteStoreError_SerializerHandshake
workerMagic
)
WorkerMagic_Two
sockPutS protoVersion serverHandshakeInputOurVersion
clientVersion <- sockGetS protoVersion
let leastCommonVersion = minVersionNumber clientVersion serverHandshakeInputOurVersion
when (not (ProtoVersion 1 37 mempty `leq` clientVersion))
$ throwError
$ RemoteStoreError_WorkerException
WorkerException_ClientVersionTooOld
-- Feature exchange (>= 1.38)
negotiatedFeatures <- if ProtoVersion 1 38 mempty `leq` leastCommonVersion
then do
clientFeatures <- sockGetS
$ mapErrorS RemoteStoreError_SerializerGet protoFeatures
sockPutS
(mapErrorS RemoteStoreError_SerializerPut protoFeatures)
(protoVersion_features serverHandshakeInputOurVersion)
pure $ Data.Set.intersection clientFeatures (protoVersion_features serverHandshakeInputOurVersion)
else pure mempty
let leastCommonVersionWithFeatures = leastCommonVersion { protoVersion_features = negotiatedFeatures }
-- postHandshake: read affinity (obsolete), reserveSpace (obsolete)
x :: Word32 <- sockGetS int
when (x /= 0) $ do
-- Obsolete CPU affinity.
_ :: Word32 <- sockGetS int
pure ()
_ :: Word32 <- sockGetS int -- obsolete reserveSpace
sockPutS
(mapErrorS RemoteStoreError_SerializerPut text)
serverHandshakeInputNixVersion
sockPutS
(mapErrorS RemoteStoreError_SerializerHandshake trustedFlag)
serverHandshakeInputTrust
pure ServerHandshakeOutput
{ serverHandshakeOutputLeastCommonVersion = leastCommonVersionWithFeatures
, serverHandshakeOutputClientVersion = clientVersion
}
{-# WARNING _unimplemented "not yet implemented" #-}
_unimplemented :: RemoteStoreError
_unimplemented = RemoteStoreError_WorkerException $ WorkerException_Error $ WorkerError_NotYetImplemented
bracketLogger
:: MonadRemoteStore m
=> TunnelLogger
-> m a
-> m a
bracketLogger tunnelLogger m = do
startWork tunnelLogger
a <- m
stopWork tunnelLogger
pure a
---
data TunnelLogger = TunnelLogger
{ _tunnelLogger_state :: IORef TunnelLoggerState
}
data TunnelLoggerState = TunnelLoggerState
{ _tunnelLoggerState_canSendStderr :: Bool
, _tunnelLoggerState_pendingMsgs :: [Logger]
}
newTunnelLogger :: IO TunnelLogger
newTunnelLogger = TunnelLogger <$> newIORef (TunnelLoggerState False [])
enqueueMsg
:: ( MonadRemoteStore m
, MonadError LoggerSError m
)
=> TunnelLogger
-> Logger
-> m ()
enqueueMsg x l = updateLogger x $ \st@(TunnelLoggerState c p) -> case c of
True -> (st, do pv <- getProtoVersion; sockPutS (logger pv) l)
False -> (TunnelLoggerState c (l:p), pure ())
_log
:: ( MonadRemoteStore m
, MonadError LoggerSError m
)
=> TunnelLogger
-> Text
-> m ()
_log l s = enqueueMsg l (Logger_Next s)
startWork
:: MonadRemoteStore m
=> TunnelLogger
-> m ()
startWork x = updateLogger x $ \(TunnelLoggerState _ p) -> (,)
(TunnelLoggerState True []) $ do
pv <- getProtoVersion
let logger' = mapErrorS RemoteStoreError_SerializerLogger (logger pv)
traverse_ (sockPutS logger') $ reverse p
stopWork
:: MonadRemoteStore m
=> TunnelLogger
-> m ()
stopWork x = updateLogger x $ \_ -> (,)
(TunnelLoggerState False []) $ do
pv <- getProtoVersion
sockPutS (mapErrorS RemoteStoreError_SerializerLogger (logger pv)) Logger_Last
-- | Stop sending logging and report an error.
--
-- Returns true if the the session was in a state that allowed the error to be
-- sent.
--
-- Unlike 'stopWork', this function may be called at any time to (try) to end a
-- session with an error.
_stopWorkOnError
:: MonadRemoteStore m
=> TunnelLogger
-> ErrorInfo
-> m Bool
_stopWorkOnError x ex = updateLogger x $ \st ->
case _tunnelLoggerState_canSendStderr st of
False -> (st, pure False)
True -> (,) (TunnelLoggerState False []) $ do
pv <- getProtoVersion
let logger' = mapErrorS RemoteStoreError_SerializerLogger (logger pv)
sockPutS logger' (Logger_Error ex)
pure True
updateLogger
:: MonadRemoteStore m
=> TunnelLogger
-> (TunnelLoggerState -> (TunnelLoggerState, m a))
-> m a
updateLogger x = join . liftIO . atomicModifyIORef (_tunnelLogger_state x)
-- | Read framed source data from a socket.
--
-- The framed protocol sends chunks as: little-endian u64 length
-- followed by that many raw bytes, terminated by a zero-length frame.
readFramedSource
:: Socket
-> (ByteString -> IO ())
-> IO ()
readFramedSource sock sink = go
where
go = do
chunkLen <- recvWord64le sock
unless (chunkLen == 0) $ do
recvExact sock (fromIntegral chunkLen) sink
go
recvWord64le :: Socket -> IO Word64
recvWord64le s = do
bs <- recvExactBS s 8
pure $ sum
[ fromIntegral (Data.ByteString.index bs i) `shiftL` (i * 8)
| i <- [0..7]
]
recvExactBS :: Socket -> Int -> IO ByteString
recvExactBS s n = do
bs <- Network.Socket.ByteString.recv s n
let got = Data.ByteString.length bs
if got == 0 then fail "readFramedSource: unexpected EOF"
else if got == n then pure bs
else do
rest <- recvExactBS s (n - got)
pure $ bs <> rest
recvExact :: Socket -> Int -> (ByteString -> IO ()) -> IO ()
recvExact s remaining f = when (remaining > 0) $ do
let toRead = min 16384 remaining
bs <- Network.Socket.ByteString.recv s toRead
let got = Data.ByteString.length bs
when (got == 0) $ fail "readFramedSource: unexpected EOF"
f bs
recvExact s (remaining - got) f