{-# LANGUAGE DeriveLift #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE OverloadedStrings #-}

-- |

module Casa.Client
  ( blobsSource
  , SourceConfig(..)
  , blobsSink
  , CasaRepoPrefix
  , parseCasaRepoPrefix
  , thParserCasaRepo
  , PushException(..)
  , PullException(..)
  ) where

import           Casa.Types
import           Control.Monad
import           Control.Monad.Catch
import           Control.Monad.IO.Class
import           Control.Monad.IO.Unlift
import           Control.Monad.Trans.Resource
import qualified Crypto.Hash as Crypto
import           Data.Aeson
import qualified Data.Attoparsec.ByteString as Atto
import qualified Data.ByteArray as Mem
import           Data.ByteString (ByteString)
import qualified Data.ByteString as S
import qualified Data.ByteString.Builder as SB
import           Data.Conduit
import           Data.Conduit.Attoparsec
import           Data.Conduit.ByteString.Builder
import qualified Data.Conduit.List as CL
import           Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as HM
import           Data.List
import           Data.Monoid ((<>))
import           Data.Typeable
import           Language.Haskell.TH
import           Language.Haskell.TH.Lift
import           Network.HTTP.Client.Conduit (requestBodySourceChunked)
import           Network.HTTP.Simple
import           Network.HTTP.Types
import           Network.URI

-- | An exception from blob consuming/sending.
data PullException
  = AttoParseError ParseError
  | BadHttpStatus Status
  | TooManyReturnedKeys Int
  deriving (Int -> PullException -> ShowS
[PullException] -> ShowS
PullException -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PullException] -> ShowS
$cshowList :: [PullException] -> ShowS
show :: PullException -> String
$cshow :: PullException -> String
showsPrec :: Int -> PullException -> ShowS
$cshowsPrec :: Int -> PullException -> ShowS
Show, Typeable)
instance Exception PullException

-- | An exception from blob consuming/sending.
data PushException
  = PushBadHttpStatus Status
  deriving (Int -> PushException -> ShowS
[PushException] -> ShowS
PushException -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [PushException] -> ShowS
$cshowList :: [PushException] -> ShowS
show :: PushException -> String
$cshow :: PushException -> String
showsPrec :: Int -> PushException -> ShowS
$cshowsPrec :: Int -> PushException -> ShowS
Show, Typeable)
instance Exception PushException

-- | The URL prefix for a casa repo.
-- Commonly: @https://casa.fpcomplete.com@
-- Parsers will strip out a trailing slash.
newtype CasaRepoPrefix =
  CasaRepoPrefix String
  deriving (Int -> CasaRepoPrefix -> ShowS
[CasaRepoPrefix] -> ShowS
CasaRepoPrefix -> String
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
showList :: [CasaRepoPrefix] -> ShowS
$cshowList :: [CasaRepoPrefix] -> ShowS
show :: CasaRepoPrefix -> String
$cshow :: CasaRepoPrefix -> String
showsPrec :: Int -> CasaRepoPrefix -> ShowS
$cshowsPrec :: Int -> CasaRepoPrefix -> ShowS
Show, forall t.
(forall (m :: * -> *). Quote m => t -> m Exp)
-> (forall (m :: * -> *). Quote m => t -> Code m t) -> Lift t
forall (m :: * -> *). Quote m => CasaRepoPrefix -> m Exp
forall (m :: * -> *).
Quote m =>
CasaRepoPrefix -> Code m CasaRepoPrefix
liftTyped :: forall (m :: * -> *).
Quote m =>
CasaRepoPrefix -> Code m CasaRepoPrefix
$cliftTyped :: forall (m :: * -> *).
Quote m =>
CasaRepoPrefix -> Code m CasaRepoPrefix
lift :: forall (m :: * -> *). Quote m => CasaRepoPrefix -> m Exp
$clift :: forall (m :: * -> *). Quote m => CasaRepoPrefix -> m Exp
Lift)
instance FromJSON CasaRepoPrefix where
  parseJSON :: Value -> Parser CasaRepoPrefix
parseJSON Value
j = do
    String
s <- forall a. FromJSON a => Value -> Parser a
parseJSON Value
j
    forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either forall (m :: * -> *) a. MonadFail m => String -> m a
fail forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String CasaRepoPrefix
parseCasaRepoPrefix String
s)

-- | TH compile-time parser.
thParserCasaRepo :: String -> Q Exp
thParserCasaRepo :: String -> Q Exp
thParserCasaRepo = forall a c b. (a -> c) -> (b -> c) -> Either a b -> c
either forall a. HasCallStack => String -> a
error forall t (m :: * -> *). (Lift t, Quote m) => t -> m Exp
lift forall b c a. (b -> c) -> (a -> b) -> a -> c
. String -> Either String CasaRepoPrefix
parseCasaRepoPrefix

-- | Parse and normalize a Casa repo prefix.
parseCasaRepoPrefix :: String -> Either String CasaRepoPrefix
parseCasaRepoPrefix :: String -> Either String CasaRepoPrefix
parseCasaRepoPrefix String
s =
  case String -> Maybe URI
parseURI String
s of
    Maybe URI
Nothing ->
      forall a b. a -> Either a b
Left
        String
"Invalid URI for repo. Should be a valid URI e.g. https://casa.fpcomplete.com"
    Just {} -> forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> CasaRepoPrefix
CasaRepoPrefix (ShowS
stripTrailing String
s))
  where
    stripTrailing :: ShowS
stripTrailing = forall a. [a] -> [a]
reverse forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. (a -> Bool) -> [a] -> [a]
dropWhile (forall a. Eq a => a -> a -> Bool
== Char
'/') forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall a. [a] -> [a]
reverse

-- | Used to build request paths.
casaServerVersion :: String
casaServerVersion :: String
casaServerVersion = String
"v1"

-- | Build the URL from a repo prefix.
casaRepoPushUrl :: CasaRepoPrefix -> String
casaRepoPushUrl :: CasaRepoPrefix -> String
casaRepoPushUrl (CasaRepoPrefix String
uri) = String
uri forall a. [a] -> [a] -> [a]
++ String
"/" forall a. [a] -> [a] -> [a]
++ String
casaServerVersion forall a. [a] -> [a] -> [a]
++ String
"/push"

-- | Build the URL from a repo prefix.
casaRepoPullUrl :: CasaRepoPrefix -> String
casaRepoPullUrl :: CasaRepoPrefix -> String
casaRepoPullUrl (CasaRepoPrefix String
uri) = String
uri forall a. [a] -> [a] -> [a]
++ String
"/" forall a. [a] -> [a] -> [a]
++ String
casaServerVersion forall a. [a] -> [a] -> [a]
++ String
"/pull"

-- | A sink to push blobs to the server. Throws 'PushException'.
blobsSink ::
     (MonadIO m, MonadThrow m, MonadUnliftIO m)
  => CasaRepoPrefix
  -> ConduitT () ByteString m ()
  -> m ()
blobsSink :: forall (m :: * -> *).
(MonadIO m, MonadThrow m, MonadUnliftIO m) =>
CasaRepoPrefix -> ConduitT () ByteString m () -> m ()
blobsSink CasaRepoPrefix
casaRepoUrl ConduitT () ByteString m ()
blobs = do
  UnliftIO m
runInIO <- forall (m :: * -> *). MonadUnliftIO m => m (UnliftIO m)
askUnliftIO
  Request
request <- forall {f :: * -> *}. MonadThrow f => UnliftIO m -> f Request
makeRequest UnliftIO m
runInIO
  Response ()
response <- forall (m :: * -> *). MonadIO m => Request -> m (Response ())
httpNoBody Request
request
  case forall a. Response a -> Status
getResponseStatus Response ()
response of
    Status Int
200 ByteString
_ -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    Status
status -> forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (Status -> PushException
PushBadHttpStatus Status
status)
  where
    makeRequest :: UnliftIO m -> f Request
makeRequest (UnliftIO forall a. m a -> IO a
runInIO) =
      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
        (RequestBody -> Request -> Request
setRequestBody
           (ConduitM () ByteString IO () -> RequestBody
requestBodySourceChunked
              (forall (m :: * -> *) (n :: * -> *) i o r.
Monad m =>
(forall a. m a -> n a) -> ConduitT i o m r -> ConduitT i o n r
transPipe forall a. m a -> IO a
runInIO ConduitT () ByteString m ()
blobs forall (m :: * -> *) a b c r.
Monad m =>
ConduitT a b m () -> ConduitT b c m r -> ConduitT a c m r
.|
               forall (m :: * -> *) a b. Monad m => (a -> b) -> ConduitT a b m ()
CL.map
                 (\ByteString
v ->
                    Word64 -> Builder
SB.word64BE (forall a b. (Integral a, Num b) => a -> b
fromIntegral (ByteString -> Int
S.length ByteString
v)) forall a. Semigroup a => a -> a -> a
<> ByteString -> Builder
SB.byteString ByteString
v) forall (m :: * -> *) a b c r.
Monad m =>
ConduitT a b m () -> ConduitT b c m r -> ConduitT a c m r
.|
               forall (m :: * -> *).
PrimMonad m =>
ConduitT Builder ByteString m ()
builderToByteString)) forall b c a. (b -> c) -> (a -> b) -> a -> c
.
         ByteString -> Request -> Request
setRequestMethod ByteString
"POST")
        (forall (m :: * -> *). MonadThrow m => String -> m Request
parseRequest (CasaRepoPrefix -> String
casaRepoPushUrl CasaRepoPrefix
casaRepoUrl))

-- | Configuration for sourcing blobs from the server.
data SourceConfig =
  SourceConfig
    { SourceConfig -> CasaRepoPrefix
sourceConfigUrl :: !CasaRepoPrefix
      -- ^ URL to pull from.
    , SourceConfig -> HashMap BlobKey Int
sourceConfigBlobs :: !(HashMap BlobKey Int)
      -- ^ The blobs to pull.
    , SourceConfig -> Int
sourceConfigMaxBlobsPerRequest :: !Int
      -- ^ Maximum number of blobs per request; we split requests into
      -- chunks of this number.
    }

-- | Make a source of blobs from a URL. Throws 'PullException'.
blobsSource ::
     (MonadThrow m, MonadResource m, MonadIO m)
  => SourceConfig
  -> ConduitT i (BlobKey, ByteString) m ()
blobsSource :: forall (m :: * -> *) i.
(MonadThrow m, MonadResource m, MonadIO m) =>
SourceConfig -> ConduitT i (BlobKey, ByteString) m ()
blobsSource SourceConfig
sourceConfig = do
  Request
skeletonRequest <- ConduitT i (BlobKey, ByteString) m Request
makeSkeletonRequest
  forall {m :: * -> *} {i}.
(MonadResource m, MonadThrow m) =>
Request -> [(BlobKey, Int)] -> ConduitT i ByteString m ()
source Request
skeletonRequest (forall k v. HashMap k v -> [(k, v)]
HM.toList (SourceConfig -> HashMap BlobKey Int
sourceConfigBlobs SourceConfig
sourceConfig)) forall (m :: * -> *) a b c r.
Monad m =>
ConduitT a b m () -> ConduitT b c m r -> ConduitT a c m r
.| ConduitT
  ByteString
  (Either ParseError (PositionRange, (BlobKey, ByteString)))
  m
  ()
conduit forall (m :: * -> *) a b c r.
Monad m =>
ConduitT a b m () -> ConduitT b c m r -> ConduitT a c m r
.|
    forall {m :: * -> *} {t} {a} {o}.
(MonadThrow m, Eq t, Num t) =>
t -> ConduitT (Either ParseError (a, o)) o m ()
consumer (forall k v. HashMap k v -> Int
HM.size (SourceConfig -> HashMap BlobKey Int
sourceConfigBlobs SourceConfig
sourceConfig))
  where
    makeSkeletonRequest :: ConduitT i (BlobKey, ByteString) m Request
makeSkeletonRequest =
      forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap
        (ByteString -> Request -> Request
setRequestMethod ByteString
"POST")
        (forall (m :: * -> *). MonadThrow m => String -> m Request
parseRequest (CasaRepoPrefix -> String
casaRepoPullUrl (SourceConfig -> CasaRepoPrefix
sourceConfigUrl SourceConfig
sourceConfig)))
    source :: Request -> [(BlobKey, Int)] -> ConduitT i ByteString m ()
source Request
skeletonRequest [(BlobKey, Int)]
blobs =
      forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless
        (forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(BlobKey, Int)]
blobs)
        (do forall (m :: * -> *) (n :: * -> *) i o r.
(MonadResource m, MonadIO n) =>
Request
-> (Response (ConduitM i ByteString n ()) -> ConduitM i o m r)
-> ConduitM i o m r
httpSource
              Request
filledRequest
              (\Response (ConduitT i ByteString m ())
response ->
                 case forall a. Response a -> Status
getResponseStatus Response (ConduitT i ByteString m ())
response of
                   Status Int
200 ByteString
_ -> forall a. Response a -> a
getResponseBody Response (ConduitT i ByteString m ())
response
                   Status
status -> forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (Status -> PullException
BadHttpStatus Status
status))
            Request -> [(BlobKey, Int)] -> ConduitT i ByteString m ()
source Request
skeletonRequest [(BlobKey, Int)]
remainingBlobs)
      where
        (Request
filledRequest, [(BlobKey, Int)]
remainingBlobs) =
          SourceConfig
-> [(BlobKey, Int)] -> Request -> (Request, [(BlobKey, Int)])
setRequestBlobs SourceConfig
sourceConfig [(BlobKey, Int)]
blobs Request
skeletonRequest
    conduit :: ConduitT
  ByteString
  (Either ParseError (PositionRange, (BlobKey, ByteString)))
  m
  ()
conduit =
      forall (m :: * -> *) a b.
(Monad m, AttoparsecInput a) =>
Parser a b
-> ConduitT a (Either ParseError (PositionRange, b)) m ()
conduitParserEither (HashMap BlobKey Int -> Parser (BlobKey, ByteString)
blobKeyValueParser (SourceConfig -> HashMap BlobKey Int
sourceConfigBlobs SourceConfig
sourceConfig))
    consumer :: t -> ConduitT (Either ParseError (a, o)) o m ()
consumer t
remaining = do
      Maybe (Either ParseError (a, o))
mkeyValue <- forall (m :: * -> *) i o. Monad m => ConduitT i o m (Maybe i)
await
      case Maybe (Either ParseError (a, o))
mkeyValue of
        Maybe (Either ParseError (a, o))
Nothing -> forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
        Just (Left ParseError
x) -> forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM (ParseError -> PullException
AttoParseError ParseError
x)
        Just (Right (a
_position, o
keyValue)) ->
          if t
remaining forall a. Eq a => a -> a -> Bool
== t
0
            then forall (m :: * -> *) e a. (MonadThrow m, Exception e) => e -> m a
throwM
                   (Int -> PullException
TooManyReturnedKeys
                      (forall k v. HashMap k v -> Int
HM.size (SourceConfig -> HashMap BlobKey Int
sourceConfigBlobs SourceConfig
sourceConfig)))
            else do
              forall (m :: * -> *) o i. Monad m => o -> ConduitT i o m ()
yield o
keyValue
              t -> ConduitT (Either ParseError (a, o)) o m ()
consumer (t
remaining forall a. Num a => a -> a -> a
- t
1)

-- | Fill the body of the request with max blobs per request.
setRequestBlobs ::
     SourceConfig -> [(BlobKey, Int)] -> Request -> (Request, [(BlobKey, Int)])
setRequestBlobs :: SourceConfig
-> [(BlobKey, Int)] -> Request -> (Request, [(BlobKey, Int)])
setRequestBlobs SourceConfig
sourceConfig [(BlobKey, Int)]
blobs Request
skeletonRequest = (Request
request, [(BlobKey, Int)]
remaining)
  where
    request :: Request
request =
      ByteString -> Request -> Request
setRequestBodyLBS
        (Builder -> ByteString
SB.toLazyByteString
           (forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl'
              (\Builder
a (BlobKey
k, Int
v) ->
                 Builder
a forall a. Semigroup a => a -> a -> a
<> (BlobKey -> Builder
blobKeyToBuilder BlobKey
k forall a. Semigroup a => a -> a -> a
<> Word64 -> Builder
SB.word64BE (forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
v)))
              forall a. Monoid a => a
mempty
              [(BlobKey, Int)]
thisBatch))
        Request
skeletonRequest
    ([(BlobKey, Int)]
thisBatch, [(BlobKey, Int)]
remaining) =
      forall a. Int -> [a] -> ([a], [a])
splitAt (SourceConfig -> Int
sourceConfigMaxBlobsPerRequest SourceConfig
sourceConfig) [(BlobKey, Int)]
blobs

-- | Parser for a key/value.
blobKeyValueParser :: HashMap BlobKey Int -> Atto.Parser (BlobKey, ByteString)
blobKeyValueParser :: HashMap BlobKey Int -> Parser (BlobKey, ByteString)
blobKeyValueParser HashMap BlobKey Int
lengths = do
  BlobKey
blobKey <- Parser BlobKey
blobKeyBinaryParser
  case forall k v. (Eq k, Hashable k) => k -> HashMap k v -> Maybe v
HM.lookup BlobKey
blobKey HashMap BlobKey Int
lengths of
    Maybe Int
Nothing -> forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
"Invalid key: " forall a. Semigroup a => a -> a -> a
<> forall a. Show a => a -> String
show BlobKey
blobKey)
    Just Int
len -> do
      ByteString
blob <- (Int -> Parser ByteString
Atto.take Int
len)
      if ByteString -> BlobKey
BlobKey (ByteString -> ByteString
sha256Hash ByteString
blob) forall a. Eq a => a -> a -> Bool
== BlobKey
blobKey
        then forall (f :: * -> *) a. Applicative f => a -> f a
pure (BlobKey
blobKey, ByteString
blob)
        else forall (m :: * -> *) a. MonadFail m => String -> m a
fail (String
"Content does not match SHA256 hash: " forall a. [a] -> [a] -> [a]
++ forall a. Show a => a -> String
show BlobKey
blobKey)

-- | Hash some raw bytes.
sha256Hash :: ByteString -> ByteString
sha256Hash :: ByteString -> ByteString
sha256Hash = forall bin bout.
(ByteArrayAccess bin, ByteArray bout) =>
bin -> bout
Mem.convert forall b c a. (b -> c) -> (a -> b) -> a -> c
. forall ba alg.
(ByteArrayAccess ba, HashAlgorithm alg) =>
alg -> ba -> Digest alg
Crypto.hashWith SHA256
Crypto.SHA256