From 4ac6f25229e7ff177ce8aef76e1fd599d9fbfeb4 Mon Sep 17 00:00:00 2001 From: Andrew Martin Date: Fri, 12 Jan 2024 11:32:03 -0500 Subject: [PATCH] Add exchangeDiscardBody --- src/Exchange.hs | 131 ++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 110 insertions(+), 21 deletions(-) diff --git a/src/Exchange.hs b/src/Exchange.hs index cda4d9c..f6f1cf8 100644 --- a/src/Exchange.hs +++ b/src/Exchange.hs @@ -8,6 +8,7 @@ module Exchange ( Exception(..) , HttpException(..) , exchange + , exchangeDiscardBody ) where import Channel (M,Resource,SendException,ReceiveException,send,receive) @@ -86,6 +87,22 @@ instance Show Exception where showsPrec d (Receive e) = showParen (d > 10) (showString "Receive " . Channel.showsPrecReceiveException 11 e) +-- | Variant of @exchange@ that discards the response body. This can be +-- used safely even when the size of the response body is greater than +-- the amount of memory available. +-- +-- This is intended as a resident-memory optimization for situations where +-- the caller ignores the response body. +exchangeDiscardBody :: + Resource + -> Bodied Request -- http request line and headers + -> M (Either Exception Response) +exchangeDiscardBody ctx req = do + let enc = Request.bodiedToChunks req + send ctx enc >>= \case + Left err -> pure (Left (Send err)) + Right () -> receiveResponseDiscardBody ctx + -- | Send an HTTP request and await a response. This function takes -- responsibility for encoding the request and decoding the response. -- It deals with the @Transfer-Encoding@ of the response and supports @@ -98,30 +115,50 @@ exchange ctx req = do let enc = Request.bodiedToChunks req send ctx enc >>= \case Left err -> pure (Left (Send err)) - Right () -> receiveResponse ctx + Right () -> receiveResponsePreserveBody ctx -receiveResponse :: +-- Returns response. Also returns leftovers that belong to the body. +receiveHeaders :: + Resource + -> M (Either Exception (Response, Bytes)) +receiveHeaders !ctx = go mempty + where + go :: Bytes -> M (Either Exception (Response, Bytes)) + go !oldOutput = receive ctx >>= \case + Left err -> pure (Left (Receive err)) + Right newOutput -> case Bytes.length newOutput of + 0 -> pure (Left (Http (E.HeadersEndOfInput oldOutput))) + _ -> do + let output = oldOutput <> newOutput + case splitEndOfHeaders output of + Nothing -> if Bytes.length output > 16000 + then pure (Left (Http E.HeadersTooLarge)) + else go output + Just (pre,post) -> case Response.decode 128 pre of + Nothing -> pure (Left (Http E.HeadersMalformed)) + Just resp -> pure (Right (resp, post)) + +receiveResponsePreserveBody :: Resource -> M (Either Exception (Bodied Response)) -receiveResponse !ctx = do - let go !oldOutput = receive ctx >>= \case - Left err -> pure (Left (Receive err)) - Right newOutput -> case Bytes.length newOutput of - 0 -> pure (Left (Http (E.HeadersEndOfInput oldOutput))) - _ -> do - let output = oldOutput <> newOutput - case splitEndOfHeaders output of - Nothing -> if Bytes.length output > 16000 - then pure (Left (Http E.HeadersTooLarge)) - else go output - Just (pre,post) -> case Response.decode 128 pre of - Nothing -> pure (Left (Http E.HeadersMalformed)) - Just resp@Response.Response{headers} -> case lookupTransferEncoding headers of - Left err -> pure (Left (Http err)) - Right enc -> case enc of - Nonchunked -> handleNonchunkedBody ctx resp post headers - Chunked -> handleChunkedBody ctx resp post - go mempty +receiveResponsePreserveBody !ctx = receiveHeaders ctx >>= \case + Left err -> pure (Left err) + Right (resp@Response.Response{headers}, post) -> case lookupTransferEncoding headers of + Left err -> pure (Left (Http err)) + Right enc -> case enc of + Nonchunked -> handleNonchunkedBody ctx resp post headers + Chunked -> handleChunkedBody ctx resp post + +receiveResponseDiscardBody :: + Resource + -> M (Either Exception Response) +receiveResponseDiscardBody !ctx = receiveHeaders ctx >>= \case + Left err -> pure (Left err) + Right (resp@Response.Response{headers}, post) -> case lookupTransferEncoding headers of + Left err -> pure (Left (Http err)) + Right enc -> case enc of + Nonchunked -> discardNonchunkedBody ctx resp post headers + Chunked -> discardChunkedBody ctx resp post handleChunkedBody :: Resource @@ -210,6 +247,10 @@ parserChunkedChunkLengthPostCr !n !chunks0 = Latin.opt >>= \case _ -> Parser.fail E.ExpectedCrlfAfterChunkLength Nothing -> pure (Continuation (PostCr n) chunks0) +-- Note: We could do much better. Upfront, we could allocate a +-- mutable byte array that is big enough to hold the entire body. +-- This would require changing the signature to make a primitive +-- offering reception into mutable byte arrays available. handleNonchunkedBody :: Resource -> Response -> Bytes -> Headers -> M (Either Exception (Bodied Response)) handleNonchunkedBody ctx resp !post !headers = case lookupContentLength headers of Left err -> pure (Left (Http err)) @@ -227,6 +268,23 @@ handleNonchunkedBody ctx resp !post !headers = case lookupContentLength headers Left err -> pure (Left (Receive err)) finish (ChunksCons post ChunksNil) (len - Bytes.length post) +-- This is not great. It relies on the GC to clean up the received +-- bytes for us. It would be better to reuse a mutable byte array +-- and receive into it repeatedly. +discardNonchunkedBody :: Resource -> Response -> Bytes -> Headers -> M (Either Exception Response) +discardNonchunkedBody ctx resp !post !headers = case lookupContentLength headers of + Left err -> pure (Left (Http err)) + Right len -> do + let finish n = case compare n 0 of + LT -> pure (Left (Http E.PipelinedResponses)) + EQ -> pure $ Right $ resp + GT -> receive ctx >>= \case + Right chunk -> case Bytes.length chunk of + 0 -> pure (Left (Http E.NonchunkedBodyEndOfInput)) + _ -> finish (n - Bytes.length chunk) + Left err -> pure (Left (Receive err)) + finish (len - Bytes.length post) + splitEndOfHeaders :: Bytes -> Maybe (Bytes, Bytes) splitEndOfHeaders !b = case Bytes.findTetragramIndex 0x0D 0x0A 0x0D 0x0A b of Nothing -> Nothing @@ -251,3 +309,34 @@ lookupContentLength !hdrs = Just i -> do when (i > 8_000_000_000) (Left E.ContentLengthTooLarge) Right i + +discardChunkedBody :: + Resource + -> Response + -> Bytes + -> M (Either Exception Response) +discardChunkedBody !ctx resp !input0 = do + let go :: Instruction -> Bytes -> M (Either Exception Response) + go instrA !inputA = case Parser.parseBytes (parserChunked (upgradeInstruction instrA)) inputA of + Parser.Failure e -> pure (Left (Http e)) + Parser.Success (Parser.Slice _ leftoverLen contB) -> + let instrB = downgradeContinuation contB in + case leftoverLen of + -- We expect that parserChunked consumes all input, so we check + -- here to be certain that it actually does. + 0 -> case instrB of + Done -> pure $ Right $ resp + _ -> receive ctx >>= \case + Right inputB -> case Bytes.length inputB of + 0 -> pure (Left (Http E.ChunkedBodyEndOfInput)) + _ -> go instrB inputB + Left err -> pure (Left (Receive err)) + _ -> pure (Left (Http E.ImplementationMistake)) + let instr0 = ChunkLength 0 + go instr0 input0 + +upgradeInstruction :: Instruction -> Continuation +upgradeInstruction i = Continuation i ChunksNil + +downgradeContinuation :: Continuation -> Instruction +downgradeContinuation (Continuation i _) = i