{-# LANGUAGE OverloadedStrings   #-}
{-# LANGUAGE RecordWildCards     #-}
module System.Nix.Store.Remote.Util where

import           Control.Monad.Except
import           Control.Monad.Reader

import           Data.Either
import           Data.Binary.Get
import           Data.Binary.Put
import           Data.Text                 (Text)
import qualified Data.Text.Encoding        as T
import qualified Data.Text.Lazy            as TL
import qualified Data.Text.Lazy.Encoding   as TL
import           Data.Time
import           Data.Time.Clock.POSIX
import           Data.ByteString           (ByteString)
import qualified Data.ByteString.Char8     as BSC
import qualified Data.ByteString.Lazy      as BSL

import           Network.Socket.ByteString (recv, sendAll)

import           Nix.Derivation

import           System.Nix.Build
import           System.Nix.StorePath
import           System.Nix.Store.Remote.Binary
import           System.Nix.Store.Remote.Types

import qualified Data.HashSet
import qualified Data.Map

genericIncremental :: (MonadIO m) => m (Maybe ByteString) -> Get a -> m a
genericIncremental :: m (Maybe ByteString) -> Get a -> m a
genericIncremental m (Maybe ByteString)
getsome Get a
parser = Decoder a -> m a
forall a. Decoder a -> m a
go Decoder a
decoder
  where
    decoder :: Decoder a
decoder = Get a -> Decoder a
forall a. Get a -> Decoder a
runGetIncremental Get a
parser
    go :: Decoder a -> m a
go (Done ByteString
_leftover ByteOffset
_consumed a
x) = do
      a -> m a
forall (m :: * -> *) a. Monad m => a -> m a
return a
x
    go (Partial Maybe ByteString -> Decoder a
k) = do
      Maybe ByteString
chunk <- m (Maybe ByteString)
getsome
      Decoder a -> m a
go (Maybe ByteString -> Decoder a
k Maybe ByteString
chunk)
    go (Fail ByteString
_leftover ByteOffset
_consumed String
msg) = do
      String -> m a
forall a. HasCallStack => String -> a
error String
msg

getSocketIncremental :: Get a -> MonadStore a
getSocketIncremental :: Get a -> MonadStore a
getSocketIncremental = ExceptT
  String
  (StateT (Maybe ByteString, [Logger]) (ReaderT StoreConfig IO))
  (Maybe ByteString)
-> Get a -> MonadStore a
forall (m :: * -> *) a.
MonadIO m =>
m (Maybe ByteString) -> Get a -> m a
genericIncremental ExceptT
  String
  (StateT (Maybe ByteString, [Logger]) (ReaderT StoreConfig IO))
  (Maybe ByteString)
sockGet8
  where
    sockGet8 :: MonadStore (Maybe BSC.ByteString)
    sockGet8 :: ExceptT
  String
  (StateT (Maybe ByteString, [Logger]) (ReaderT StoreConfig IO))
  (Maybe ByteString)
sockGet8 = do
      Socket
soc <- (StoreConfig -> Socket)
-> ExceptT
     String
     (StateT (Maybe ByteString, [Logger]) (ReaderT StoreConfig IO))
     Socket
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks StoreConfig -> Socket
storeSocket
      IO (Maybe ByteString)
-> ExceptT
     String
     (StateT (Maybe ByteString, [Logger]) (ReaderT StoreConfig IO))
     (Maybe ByteString)
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO (Maybe ByteString)
 -> ExceptT
      String
      (StateT (Maybe ByteString, [Logger]) (ReaderT StoreConfig IO))
      (Maybe ByteString))
-> IO (Maybe ByteString)
-> ExceptT
     String
     (StateT (Maybe ByteString, [Logger]) (ReaderT StoreConfig IO))
     (Maybe ByteString)
forall a b. (a -> b) -> a -> b
$ ByteString -> Maybe ByteString
forall a. a -> Maybe a
Just (ByteString -> Maybe ByteString)
-> IO ByteString -> IO (Maybe ByteString)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Socket -> Int -> IO ByteString
recv Socket
soc Int
8

sockPut :: Put -> MonadStore ()
sockPut :: Put -> MonadStore ()
sockPut Put
p = do
  Socket
soc <- (StoreConfig -> Socket)
-> ExceptT
     String
     (StateT (Maybe ByteString, [Logger]) (ReaderT StoreConfig IO))
     Socket
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks StoreConfig -> Socket
storeSocket
  IO () -> MonadStore ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> MonadStore ()) -> IO () -> MonadStore ()
forall a b. (a -> b) -> a -> b
$ Socket -> ByteString -> IO ()
sendAll Socket
soc (ByteString -> IO ()) -> ByteString -> IO ()
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
BSL.toStrict (ByteString -> ByteString) -> ByteString -> ByteString
forall a b. (a -> b) -> a -> b
$ Put -> ByteString
runPut Put
p

sockGet :: Get a -> MonadStore a
sockGet :: Get a -> MonadStore a
sockGet = Get a -> MonadStore a
forall a. Get a -> MonadStore a
getSocketIncremental

sockGetInt :: Integral a => MonadStore a
sockGetInt :: MonadStore a
sockGetInt = Get a -> MonadStore a
forall a. Get a -> MonadStore a
getSocketIncremental Get a
forall a. Integral a => Get a
getInt

sockGetBool :: MonadStore Bool
sockGetBool :: MonadStore Bool
sockGetBool = (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== (Int
1 :: Int)) (Int -> Bool)
-> ExceptT
     String
     (StateT (Maybe ByteString, [Logger]) (ReaderT StoreConfig IO))
     Int
-> MonadStore Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> ExceptT
  String
  (StateT (Maybe ByteString, [Logger]) (ReaderT StoreConfig IO))
  Int
forall a. Integral a => MonadStore a
sockGetInt

sockGetStr :: MonadStore ByteString
sockGetStr :: MonadStore ByteString
sockGetStr = Get ByteString -> MonadStore ByteString
forall a. Get a -> MonadStore a
getSocketIncremental Get ByteString
getByteStringLen

sockGetStrings :: MonadStore [ByteString]
sockGetStrings :: MonadStore [ByteString]
sockGetStrings = Get [ByteString] -> MonadStore [ByteString]
forall a. Get a -> MonadStore a
getSocketIncremental Get [ByteString]
getByteStrings

sockGetPath :: MonadStore StorePath
sockGetPath :: MonadStore StorePath
sockGetPath = do
  String
sd <- MonadStore String
getStoreDir
  Either String StorePath
pth <- Get (Either String StorePath)
-> MonadStore (Either String StorePath)
forall a. Get a -> MonadStore a
getSocketIncremental (String -> Get (Either String StorePath)
getPath String
sd)
  case Either String StorePath
pth of
    Left String
e -> String -> MonadStore StorePath
forall e (m :: * -> *) a. MonadError e m => e -> m a
throwError String
e
    Right StorePath
x -> StorePath -> MonadStore StorePath
forall (m :: * -> *) a. Monad m => a -> m a
return StorePath
x

sockGetPathMay :: MonadStore (Maybe StorePath)
sockGetPathMay :: MonadStore (Maybe StorePath)
sockGetPathMay = do
  String
sd <- MonadStore String
getStoreDir
  Either String StorePath
pth <- Get (Either String StorePath)
-> MonadStore (Either String StorePath)
forall a. Get a -> MonadStore a
getSocketIncremental (String -> Get (Either String StorePath)
getPath String
sd)
  Maybe StorePath -> MonadStore (Maybe StorePath)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe StorePath -> MonadStore (Maybe StorePath))
-> Maybe StorePath -> MonadStore (Maybe StorePath)
forall a b. (a -> b) -> a -> b
$ case Either String StorePath
pth of
    Left String
_e -> Maybe StorePath
forall a. Maybe a
Nothing
    Right StorePath
x -> StorePath -> Maybe StorePath
forall a. a -> Maybe a
Just StorePath
x

sockGetPaths :: MonadStore StorePathSet
sockGetPaths :: MonadStore StorePathSet
sockGetPaths = do
  String
sd <- MonadStore String
getStoreDir
  Get StorePathSet -> MonadStore StorePathSet
forall a. Get a -> MonadStore a
getSocketIncremental (String -> Get StorePathSet
getPaths String
sd)

bsToText :: ByteString -> Text
bsToText :: ByteString -> Text
bsToText = ByteString -> Text
T.decodeUtf8

textToBS :: Text -> ByteString
textToBS :: Text -> ByteString
textToBS = Text -> ByteString
T.encodeUtf8

bslToText :: BSL.ByteString -> Text
bslToText :: ByteString -> Text
bslToText = Text -> Text
TL.toStrict (Text -> Text) -> (ByteString -> Text) -> ByteString -> Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Text
TL.decodeUtf8

textToBSL :: Text -> BSL.ByteString
textToBSL :: Text -> ByteString
textToBSL = Text -> ByteString
TL.encodeUtf8 (Text -> ByteString) -> (Text -> Text) -> Text -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> Text
TL.fromStrict

putText :: Text -> Put
putText :: Text -> Put
putText = ByteString -> Put
putByteStringLen (ByteString -> Put) -> (Text -> ByteString) -> Text -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Text -> ByteString
textToBSL

putTexts :: [Text] -> Put
putTexts :: [Text] -> Put
putTexts = [ByteString] -> Put
forall (t :: * -> *). Foldable t => t ByteString -> Put
putByteStrings ([ByteString] -> Put) -> ([Text] -> [ByteString]) -> [Text] -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Text -> ByteString) -> [Text] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map Text -> ByteString
textToBSL

getPath :: FilePath -> Get (Either String StorePath)
getPath :: String -> Get (Either String StorePath)
getPath String
sd = String -> ByteString -> Either String StorePath
parsePath String
sd (ByteString -> Either String StorePath)
-> Get ByteString -> Get (Either String StorePath)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get ByteString
getByteStringLen

getPaths :: FilePath -> Get StorePathSet
getPaths :: String -> Get StorePathSet
getPaths String
sd = [StorePath] -> StorePathSet
forall a. (Eq a, Hashable a) => [a] -> HashSet a
Data.HashSet.fromList ([StorePath] -> StorePathSet)
-> ([ByteString] -> [StorePath]) -> [ByteString] -> StorePathSet
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Either String StorePath] -> [StorePath]
forall a b. [Either a b] -> [b]
rights ([Either String StorePath] -> [StorePath])
-> ([ByteString] -> [Either String StorePath])
-> [ByteString]
-> [StorePath]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ByteString -> Either String StorePath)
-> [ByteString] -> [Either String StorePath]
forall a b. (a -> b) -> [a] -> [b]
map (String -> ByteString -> Either String StorePath
parsePath String
sd) ([ByteString] -> StorePathSet)
-> Get [ByteString] -> Get StorePathSet
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get [ByteString]
getByteStrings

putPath :: StorePath -> Put
putPath :: StorePath -> Put
putPath  = ByteString -> Put
putByteStringLen (ByteString -> Put)
-> (StorePath -> ByteString) -> StorePath -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
BSL.fromStrict (ByteString -> ByteString)
-> (StorePath -> ByteString) -> StorePath -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StorePath -> ByteString
storePathToRawFilePath

putPaths :: StorePathSet -> Put
putPaths :: StorePathSet -> Put
putPaths = [ByteString] -> Put
forall (t :: * -> *). Foldable t => t ByteString -> Put
putByteStrings ([ByteString] -> Put)
-> (StorePathSet -> [ByteString]) -> StorePathSet -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. HashSet ByteString -> [ByteString]
forall a. HashSet a -> [a]
Data.HashSet.toList (HashSet ByteString -> [ByteString])
-> (StorePathSet -> HashSet ByteString)
-> StorePathSet
-> [ByteString]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (StorePath -> ByteString) -> StorePathSet -> HashSet ByteString
forall b a.
(Hashable b, Eq b) =>
(a -> b) -> HashSet a -> HashSet b
Data.HashSet.map (ByteString -> ByteString
BSL.fromStrict (ByteString -> ByteString)
-> (StorePath -> ByteString) -> StorePath -> ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. StorePath -> ByteString
storePathToRawFilePath)

putBool :: Bool -> Put
putBool :: Bool -> Put
putBool Bool
True  = Int -> Put
forall a. Integral a => a -> Put
putInt (Int
1 :: Int)
putBool Bool
False = Int -> Put
forall a. Integral a => a -> Put
putInt (Int
0 :: Int)

getBool :: Get Bool
getBool :: Get Bool
getBool = (Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
==Int
1) (Int -> Bool) -> Get Int -> Get Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Get Int
forall a. Integral a => Get a
getInt :: Get Int)

putEnum :: (Enum a) => a -> Put
putEnum :: a -> Put
putEnum = Int -> Put
forall a. Integral a => a -> Put
putInt (Int -> Put) -> (a -> Int) -> a -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. a -> Int
forall a. Enum a => a -> Int
fromEnum

getEnum :: (Enum a) => Get a
getEnum :: Get a
getEnum = Int -> a
forall a. Enum a => Int -> a
toEnum (Int -> a) -> Get Int -> Get a
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get Int
forall a. Integral a => Get a
getInt

putTime :: UTCTime -> Put
putTime :: UTCTime -> Put
putTime = (Int -> Put
forall a. Integral a => a -> Put
putInt :: Int -> Put) (Int -> Put) -> (UTCTime -> Int) -> UTCTime -> Put
forall b c a. (b -> c) -> (a -> b) -> a -> c
. POSIXTime -> Int
forall a b. (RealFrac a, Integral b) => a -> b
round (POSIXTime -> Int) -> (UTCTime -> POSIXTime) -> UTCTime -> Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. UTCTime -> POSIXTime
utcTimeToPOSIXSeconds

getTime :: Get UTCTime
getTime :: Get UTCTime
getTime = POSIXTime -> UTCTime
posixSecondsToUTCTime (POSIXTime -> UTCTime) -> Get POSIXTime -> Get UTCTime
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get POSIXTime
forall a. Enum a => Get a
getEnum

getBuildResult :: Get BuildResult
getBuildResult :: Get BuildResult
getBuildResult = BuildStatus
-> Maybe Text
-> Integer
-> Bool
-> UTCTime
-> UTCTime
-> BuildResult
BuildResult
  (BuildStatus
 -> Maybe Text
 -> Integer
 -> Bool
 -> UTCTime
 -> UTCTime
 -> BuildResult)
-> Get BuildStatus
-> Get
     (Maybe Text
      -> Integer -> Bool -> UTCTime -> UTCTime -> BuildResult)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get BuildStatus
forall a. Enum a => Get a
getEnum
  Get
  (Maybe Text
   -> Integer -> Bool -> UTCTime -> UTCTime -> BuildResult)
-> Get (Maybe Text)
-> Get (Integer -> Bool -> UTCTime -> UTCTime -> BuildResult)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Text -> Maybe Text
forall a. a -> Maybe a
Just (Text -> Maybe Text)
-> (ByteString -> Text) -> ByteString -> Maybe Text
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> Text
bsToText (ByteString -> Maybe Text) -> Get ByteString -> Get (Maybe Text)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Get ByteString
getByteStringLen)
  Get (Integer -> Bool -> UTCTime -> UTCTime -> BuildResult)
-> Get Integer -> Get (Bool -> UTCTime -> UTCTime -> BuildResult)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get Integer
forall a. Integral a => Get a
getInt
  Get (Bool -> UTCTime -> UTCTime -> BuildResult)
-> Get Bool -> Get (UTCTime -> UTCTime -> BuildResult)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get Bool
getBool
  Get (UTCTime -> UTCTime -> BuildResult)
-> Get UTCTime -> Get (UTCTime -> BuildResult)
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get UTCTime
getTime
  Get (UTCTime -> BuildResult) -> Get UTCTime -> Get BuildResult
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Get UTCTime
getTime

putDerivation :: Derivation StorePath Text -> Put
putDerivation :: Derivation StorePath Text -> Put
putDerivation Derivation{Text
Map Text Text
Map Text (DerivationOutput StorePath Text)
Map StorePath (Set Text)
Set StorePath
Vector Text
outputs :: forall fp txt.
Derivation fp txt -> Map txt (DerivationOutput fp txt)
inputDrvs :: forall fp txt. Derivation fp txt -> Map fp (Set txt)
inputSrcs :: forall fp txt. Derivation fp txt -> Set fp
platform :: forall fp txt. Derivation fp txt -> txt
builder :: forall fp txt. Derivation fp txt -> txt
args :: forall fp txt. Derivation fp txt -> Vector txt
env :: forall fp txt. Derivation fp txt -> Map txt txt
env :: Map Text Text
args :: Vector Text
builder :: Text
platform :: Text
inputSrcs :: Set StorePath
inputDrvs :: Map StorePath (Set Text)
outputs :: Map Text (DerivationOutput StorePath Text)
..} = do
  (((Text, DerivationOutput StorePath Text) -> Put)
 -> [(Text, DerivationOutput StorePath Text)] -> Put)
-> [(Text, DerivationOutput StorePath Text)]
-> ((Text, DerivationOutput StorePath Text) -> Put)
-> Put
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Text, DerivationOutput StorePath Text) -> Put)
-> [(Text, DerivationOutput StorePath Text)] -> Put
forall (t :: * -> *) a. Foldable t => (a -> Put) -> t a -> Put
putMany (Map Text (DerivationOutput StorePath Text)
-> [(Text, DerivationOutput StorePath Text)]
forall k a. Map k a -> [(k, a)]
Data.Map.toList Map Text (DerivationOutput StorePath Text)
outputs)
    (((Text, DerivationOutput StorePath Text) -> Put) -> Put)
-> ((Text, DerivationOutput StorePath Text) -> Put) -> Put
forall a b. (a -> b) -> a -> b
$ \(Text
outputName, DerivationOutput{Text
StorePath
path :: forall fp txt. DerivationOutput fp txt -> fp
hashAlgo :: forall fp txt. DerivationOutput fp txt -> txt
hash :: forall fp txt. DerivationOutput fp txt -> txt
hash :: Text
hashAlgo :: Text
path :: StorePath
..}) -> do
      Text -> Put
putText Text
outputName
      StorePath -> Put
putPath StorePath
path
      Text -> Put
putText Text
hashAlgo
      Text -> Put
putText Text
hash

  (StorePath -> Put) -> Set StorePath -> Put
forall (t :: * -> *) a. Foldable t => (a -> Put) -> t a -> Put
putMany StorePath -> Put
putPath Set StorePath
inputSrcs
  Text -> Put
putText Text
platform
  Text -> Put
putText Text
builder
  (Text -> Put) -> Vector Text -> Put
forall (t :: * -> *) a. Foldable t => (a -> Put) -> t a -> Put
putMany Text -> Put
putText Vector Text
args

  (((Text, Text) -> Put) -> [(Text, Text)] -> Put)
-> [(Text, Text)] -> ((Text, Text) -> Put) -> Put
forall a b c. (a -> b -> c) -> b -> a -> c
flip ((Text, Text) -> Put) -> [(Text, Text)] -> Put
forall (t :: * -> *) a. Foldable t => (a -> Put) -> t a -> Put
putMany (Map Text Text -> [(Text, Text)]
forall k a. Map k a -> [(k, a)]
Data.Map.toList Map Text Text
env)
    (((Text, Text) -> Put) -> Put) -> ((Text, Text) -> Put) -> Put
forall a b. (a -> b) -> a -> b
$ \(Text
first, Text
second) -> Text -> Put
putText Text
first Put -> Put -> Put
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Text -> Put
putText Text
second