Merge pull request #11 from restaumatic/save-all-requests-from-a-singe-run

Save all requests during proxying and allow compression of cassette

Author: mateuszbujalski

package.yaml

@@ -43,6 +43,7 @@ dependencies:
 - uri-bytestring
 - optparse-applicative
 - zlib
+- zstd
 
 library:
   source-dirs: src

src/Network/VCR.hs

@@ -17,7 +17,7 @@ import Control.Exception (SomeException, bracket)
 import Control.Monad
 import qualified Data.ByteString.Char8 as BS
 import qualified Data.ByteString.Lazy.Char8 as LBS
-import Data.IORef (IORef, newIORef)
+import Data.IORef (IORef, newIORef, readIORef)
 import qualified Data.Text as T
 import qualified Network.HTTP.Client as HC
 import qualified Network.HTTP.Conduit as HC
@@ -50,13 +50,14 @@ import Options.Applicative (
  )
 import System.Environment (getArgs)
 
-import Data.Yaml (decodeFileEither, encodeFile)
+import Data.Yaml (decodeFileEither, encodeFile, ToJSON, ParseException, FromJSON)
 import System.Directory (doesFileExist)
 import System.IO (
     BufferMode (..),
     hSetBuffering,
     stdout,
  )
+import qualified Network.VCR.Compression as Compression
 
 server :: IO ()
 server = execParser opts >>= run
@@ -74,14 +75,14 @@ run options = withServer options $ do
     forever $ threadDelay 1000000000
 
 withServer :: Options -> IO a -> IO a
-withServer options@Options{mode, cassettePath, port} action = do
+withServer options@Options{mode, cassettePath, port, compression} action = do
     putStrLn $ "Starting VCR proxy, mode: " <> show mode <> ", cassette file: " <> cassettePath <> ", listening on port: " <> show port
     case mode of
         Record endpoint -> do
             exists <- doesFileExist cassettePath
-            when (not exists) $ encodeFile cassettePath (emptyCassette $ T.pack endpoint)
+            when (not exists) $ saveCassette compression cassettePath (emptyCassette $ T.pack endpoint)
         _ -> pure ()
-    cas <- decodeFileEither cassettePath
+    cas <- loadCassette compression cassettePath
     case cas of
         Left err -> die $ "Cassette: " <> cassettePath <> " couldn't be decoded or found! " <> (show err)
         Right cassette -> do
@@ -91,17 +92,24 @@ withServer options@Options{mode, cassettePath, port} action = do
                 action
 
 runInternal :: Options -> IORef Cassette -> IO a -> IO a
-runInternal Options{mode, cassettePath, port} cassetteIORef action = do
+runInternal Options{mode, cassettePath, port, compression} cassetteIORef action = do
     -- Set line buffering, because if we use it from a parent process, pipes are full buffered by default
     hSetBuffering stdout LineBuffering
     started <- newEmptyMVar
-    bracket (start started) killThread $ \_ -> do
+    bracket (start started) (stop compression cassettePath cassetteIORef) $ \_ -> do
         takeMVar started
         action
   where
     start started = forkIO $ do
         mgr <- HC.newManager HC.tlsManagerSettings
-        Warp.runSettings (warpSettings started proxySettings) $ middleware mode cassetteIORef cassettePath $ HProxy.httpProxyApp proxySettings mgr
+        Warp.runSettings (warpSettings started proxySettings) $ middleware mode cassetteIORef $ HProxy.httpProxyApp proxySettings mgr
+    stop compression cassettePath casetteIORef threadId = do
+        case mode of
+            Record _ -> do
+                cassette <- readIORef casetteIORef
+                saveCassette compression cassettePath cassette
+            _ -> pure ()
+        killThread threadId
     proxySettings = HProxy.defaultProxySettings{HProxy.proxyPort = port}
 
 warpSettings ::
@@ -124,3 +132,15 @@ defaultExceptionResponse e =
         HT.badGateway502
         [(HT.hContentType, "text/plain; charset=utf-8")]
         $ LBS.fromChunks [BS.pack $ show e]
+
+saveCassette :: ToJSON a => Bool -> FilePath -> a -> IO ()
+saveCassette compression path v =
+    if compression
+    then Compression.save 3 path v
+    else encodeFile path v
+
+loadCassette :: FromJSON a => Bool -> FilePath -> IO (Either ParseException a)
+loadCassette compression path =
+    if compression
+    then Compression.load path
+    else decodeFileEither path

src/Network/VCR/Compression.hs

@@ -0,0 +1,32 @@
+{-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE OverloadedStrings #-}
+
+module Network.VCR.Compression where
+
+import qualified Codec.Compression.Zstd as Zstd
+import qualified Data.ByteString as BS
+import Data.Yaml
+
+{- | Magic number at the start of zstd frame, see <https://github.com/facebook/zstd/blob/dev/doc/zstd_compression_format.md#zstandard-frames>
+ Used to detect compressed payloads.
+-}
+zstdMagicNumber :: BS.ByteString
+zstdMagicNumber = "\x28\xb5\x2f\xfd"
+
+save :: ToJSON a => Int -> FilePath -> a -> IO ()
+save compressLevel path v = do
+  let compressed = Zstd.compress compressLevel $ encode v
+  BS.writeFile path compressed
+
+load :: FromJSON a => FilePath -> IO (Either ParseException a)
+load path = do
+  compressed <- BS.readFile path
+  if zstdMagicNumber `BS.isPrefixOf` compressed
+  then do
+    case Zstd.decompress compressed of
+      Zstd.Error err -> error $ "Invalid zstd compressed value: " <> show err
+      Zstd.Skip -> error "Invalid zstd compressed value (Skip frame)"
+      Zstd.Decompress v -> pure $ decodeEither' v
+  else
+    error $ "Expected " <> show path <> " to start with zstd magic number"

src/Network/VCR/Middleware.hs

@@ -51,23 +51,23 @@ import qualified Data.Aeson.Diff as A
 import qualified Data.Aeson.Encode.Pretty as A
 
 import Data.CaseInsensitive (CI, mk)
-import Data.IORef (atomicModifyIORef)
+import Data.IORef (atomicModifyIORef')
 import qualified System.Exit as XIO
 import System.IO (stderr)
 import qualified System.IO (hPutStrLn)
 import qualified URI.ByteString as URI
 
-middleware :: Mode -> IORef Cassette -> FilePath -> Wai.Middleware
+middleware :: Mode -> IORef Cassette -> Wai.Middleware
 middleware Replay = replayingMiddleware findAnyResponse
 middleware ReplayStrict = replayingMiddleware consumeRequestsInOrder
 middleware Record {endpoint} = recordingMiddleware endpoint
 
 {- | Middleware which only records API calls, the requests are proxied to the `endpoint` and recorded in the
  `filePath` cassette file
 -}
-recordingMiddleware :: String -> IORef Cassette -> FilePath -> Wai.Middleware
-recordingMiddleware endpoint cassetteIORef filePath app req respond = do
-  cassette@Cassette {apiCalls, ignoredHeaders} <- readIORef cassetteIORef
+recordingMiddleware :: String -> IORef Cassette -> Wai.Middleware
+recordingMiddleware endpoint cassetteIORef app req respond = do
+  Cassette {ignoredHeaders} <- readIORef cassetteIORef
   (req', body) <- getRequestBody req
   -- Construct a request that can be sent to the actual remote API, by replacing the host in the request with the endpoint
   -- passed as an argument to the middleware
@@ -80,10 +80,7 @@ recordingMiddleware endpoint cassetteIORef filePath app req respond = do
     (status, headers, reBody) <- getResponseBody response
     savedResponse <- buildResponse reBody response
     -- Store the request, response pair
-    encodeFile filePath $
-      cassette
-        { apiCalls = apiCalls <> [ApiCall {request = savedRequest, response = savedResponse}]
-        }
+    atomicModifyIORef' cassetteIORef $ \cassette -> (cassette { apiCalls = apiCalls cassette <> [ApiCall {request = savedRequest, response = savedResponse}] }, ())
     respond $ Wai.responseLBS status headers (gzipIfNeeded headers reBody)
 
 -- | A policy for obtaining responses given a request.
@@ -114,8 +111,8 @@ consumeRequestsInOrder cassetteIORef savedRequest = do
 {- | Middleware which only replays API calls, if a request is not found in the filePath provided cassette file,
  a 500 error will be thrown
 -}
-replayingMiddleware :: FindResponse -> IORef Cassette -> FilePath -> Wai.Middleware
-replayingMiddleware findResponse cassetteIORef filePath app req respond = do
+replayingMiddleware :: FindResponse -> IORef Cassette -> Wai.Middleware
+replayingMiddleware findResponse cassetteIORef app req respond = do
   cassette@Cassette {apiCalls, ignoredHeaders} <- readIORef cassetteIORef
   b <- Wai.strictRequestBody req
   let receivedRequest = buildRequest ignoredHeaders req b
@@ -209,7 +206,7 @@ getRequestBody req = do
   -- This implementation ensures that each chunk is only returned
   -- once.
   ichunks <- newIORef body
-  let rbody = atomicModifyIORef ichunks $ \chunks ->
+  let rbody = atomicModifyIORef' ichunks $ \chunks ->
         case chunks of
           [] -> ([], BS.empty)
           x : y -> (y, x)

src/Network/VCR/Types.hs

@@ -43,6 +43,7 @@ import Options.Applicative (
   short,
   showDefault,
   strOption,
+  switch,
   value,
  )
 
@@ -79,6 +80,7 @@ data Options = Options
   { cassettePath :: FilePath
   , mode :: Mode
   , port :: Int
+  , compression :: Bool
   }
   deriving (Eq, Show)
 
@@ -91,6 +93,7 @@ parseOptions =
     <$> strOption (long "cassette" <> short 'c' <> metavar "CASSETTE_FILE" <> help "Cassette yaml file for recording/replaying the API interactions")
     <*> parseMode
     <*> option auto (long "port" <> help "Port to listen on" <> showDefault <> value DEFAULT_PORT <> metavar "INT")
+    <*> switch (long "compression" <> help "Should cassette yaml be (de)compressed")
 
 data Body = JSONBody Value | RawBody B.ByteString
   deriving (Show, Eq)

vcr-proxy.cabal

@@ -26,6 +26,7 @@ source-repository head
 library
   exposed-modules:
       Network.VCR
+      Network.VCR.Compression
       Network.VCR.Middleware
       Network.VCR.Types
   other-modules:
@@ -56,6 +57,7 @@ library
     , warp
     , yaml
     , zlib
+    , zstd
   default-language: Haskell2010
 
 executable vcr-proxy
@@ -90,6 +92,7 @@ executable vcr-proxy
     , warp
     , yaml
     , zlib
+    , zstd
   default-language: Haskell2010
 
 test-suite vcr-proxy-test
@@ -125,4 +128,5 @@ test-suite vcr-proxy-test
     , warp
     , yaml
     , zlib
+    , zstd
   default-language: Haskell2010