Stream object transfers with shared byte throttling

test/Spec.hs

@@ -1,5 +1,6 @@
 module Main (main) where
 
+import           Control.Exception (try)
 import           Control.Monad.IO.Class (liftIO)
 import           Control.Lens ((&), (?~))
 import           Crypto.JOSE (runJOSE)
@@ -12,6 +13,7 @@ import           Crypto.JOSE.JWK
   , genJWK
   )
 import           Crypto.JOSE.JWS (newJWSHeaderProtected)
+import           Crypto.Hash (Digest, SHA256, hash)
 import           Crypto.JWT
   ( Audience (..)
   , ClaimsSet
@@ -23,7 +25,9 @@ import           Crypto.JWT
   )
 import qualified Data.Aeson as A
 import qualified Data.Aeson.KeyMap as KM
+import qualified Data.ByteArray as BA
 import qualified Data.ByteString as BS
+import qualified Data.ByteString.Base16 as B16
 import qualified Data.ByteString.Lazy as BL
 import           Data.Either (isLeft, isRight)
 import           Data.IORef (IORef, modifyIORef', newIORef, readIORef, writeIORef)
@@ -87,10 +91,21 @@ import           Nuthatch.Protocol
   )
 import           Nuthatch.Types
 import           Nuthatch.Session (mintSession)
+import           Nuthatch.Streaming
+  ( UploadStreamError (..)
+  , consumeUploadParts
+  , maximumMultipartBytes
+  , multipartPartBytes
+  )
 import           Nuthatch.Ui (browsePage, browserScript)
 import           Servant
 import           Servant.Server.Experimental.Auth (AuthHandler, mkAuthHandler)
-import           Servant.Types.SourceT (source)
+import           Servant.Types.SourceT
+  ( SourceT
+  , StepT (..)
+  , fromStepT
+  , source
+  )
 import           Test.Tasty (TestTree, defaultMain, testGroup)
 import           Test.Tasty.HUnit ((@?=), assertBool, testCase)
 import           Test.Tasty.Hedgehog (testProperty)
@@ -103,10 +118,61 @@ tests = testGroup "nuthatch-unit"
   [ pathTests
   , identityTests
   , protocolTests
+  , streamingTests
   , oidcTests
   , apiTests
   ]
 
+streamingTests :: TestTree
+streamingTests = testGroup "streaming"
+  [ testCase "upload bodies are inspected and emitted in bounded parts" $ do
+      inspected <- newIORef []
+      emitted <- newIORef []
+      consumeUploadParts
+        6
+        (strictSha256 "abcdef")
+        2
+        4
+        (\piece -> modifyIORef' inspected (<> [piece]))
+        (source ["abcdef"])
+        (\part -> modifyIORef' emitted (<> [part]))
+      readIORef inspected >>= (@?= ["ab", "cd", "ef"])
+      readIORef emitted >>= (@?= ["abcd", "ef"])
+  , testCase "multipart buffering has a fixed part-count bound" $
+      maximumMultipartBytes @?= toInteger multipartPartBytes * 10000
+  , testCase "upload rejects early end, surplus bytes, and identity mismatch" $ do
+      short <- uploadFailure 4 (strictSha256 "body") (source ["bod"])
+      short @?= UploadStreamSizeMismatch
+      long <- uploadFailure 4 (strictSha256 "body") (source ["body!"])
+      long @?= UploadStreamSizeMismatch
+      identity <- uploadFailure 4 (strictSha256 "else") (source ["body"])
+      identity @?= UploadStreamIdentityMismatch
+  , testCase "upload preserves source errors" $ do
+      failure <- uploadFailure
+        4
+        (strictSha256 "body")
+        (fromStepT (Error "request failed"))
+      failure @?= UploadStreamInvalid "request failed"
+  ]
+
+uploadFailure
+  :: Integer
+  -> ContentSha256
+  -> SourceT IO BS.ByteString
+  -> IO UploadStreamError
+uploadFailure size identity body = do
+  result <- try $
+    consumeUploadParts
+      size identity 2 4 (const (pure ())) body (const (pure ()))
+  case result of
+    Left failure -> pure failure
+    Right _ -> fail "expected upload stream rejection"
+
+strictSha256 :: BS.ByteString -> ContentSha256
+strictSha256 bytes =
+  ContentSha256 . TE.decodeUtf8 . B16.encode . BA.convert $
+    (hash bytes :: Digest SHA256)
+
 pathTests :: TestTree
 pathTests = testGroup "canonical paths"
   [ testProperty "valid paths parse and render without rewriting" propPathRoundTrip
@@ -604,7 +670,7 @@ typedTestServer calls =
             }))
   :<|> (\_ _ _ -> counted
         (addHeader "attachment; filename=\"file.txt\""
-          (addHeader "nosniff" (source ["body"]))))
+          (addHeader "nosniff" (addHeader 4 (source ["body"])))))
   :<|> (\_ _ _ operation _ _ _ _ -> counted UploadResponse
         { uploadResponseOperationId = operation
         , uploadResponseResult = MutationSucceeded