From a2de346e41c0828560b850c1ea1ada849f990025 Mon Sep 17 00:00:00 2001
From: Felix Paulusma <felix.paulusma@gmail.com>
Date: Thu, 12 Oct 2023 20:12:39 +0200
Subject: [PATCH] warp: if we're building the outgoing buffer, why not count
 the bytes?

---
 warp/Network/Wai/Handler/Warp/IO.hs       | 15 ++++++++-------
 warp/Network/Wai/Handler/Warp/Response.hs |  4 ++--
 2 files changed, 10 insertions(+), 9 deletions(-)

diff --git a/warp/Network/Wai/Handler/Warp/IO.hs b/warp/Network/Wai/Handler/Warp/IO.hs
index 8789722a5..eb76abedf 100644
--- a/warp/Network/Wai/Handler/Warp/IO.hs
+++ b/warp/Network/Wai/Handler/Warp/IO.hs
@@ -10,19 +10,20 @@ import Network.Wai.Handler.Warp.Buffer
 import Network.Wai.Handler.Warp.Imports
 import Network.Wai.Handler.Warp.Types
 
-toBufIOWith :: Int -> IORef WriteBuffer -> (ByteString -> IO ()) -> Builder -> IO ()
+toBufIOWith :: Int -> IORef WriteBuffer -> (ByteString -> IO ()) -> Builder -> IO Int
 toBufIOWith maxRspBufSize writeBufferRef io builder = do
   writeBuffer <- readIORef writeBufferRef
-  loop writeBuffer firstWriter
+  loop writeBuffer firstWriter 0
   where
     firstWriter = runBuilder builder
-    loop writeBuffer writer = do
+    loop writeBuffer writer bytesSent = do
       let buf = bufBuffer writeBuffer
           size = bufSize writeBuffer
       (len, signal) <- writer buf size
       bufferIO buf len io
+      let totalBytesSent = len + bytesSent
       case signal of
-        Done -> return ()
+        Done -> return totalBytesSent
         More minSize next
           | size < minSize -> do
               when (minSize > maxRspBufSize) $
@@ -40,8 +41,8 @@ toBufIOWith maxRspBufSize writeBufferRef io builder = do
                 biggerWriteBuffer <- createWriteBuffer minSize
                 writeIORef writeBufferRef biggerWriteBuffer
                 return biggerWriteBuffer
-              loop biggerWriteBuffer next
-          | otherwise -> loop writeBuffer next
+              loop biggerWriteBuffer next totalBytesSent
+          | otherwise -> loop writeBuffer next totalBytesSent
         Chunk bs next -> do
           io bs
-          loop writeBuffer next
+          loop writeBuffer next totalBytesSent
diff --git a/warp/Network/Wai/Handler/Warp/Response.hs b/warp/Network/Wai/Handler/Warp/Response.hs
index ebc3790e5..9a71dfdbe 100644
--- a/warp/Network/Wai/Handler/Warp/Response.hs
+++ b/warp/Network/Wai/Handler/Warp/Response.hs
@@ -224,8 +224,8 @@ sendRsp conn _ th ver s hs _ maxRspBufSize _ (RspBuilder body needsChunked) = do
                                  <> chunkedTransferTerminator
          | otherwise    = header <> body
         writeBufferRef = connWriteBuffer conn
-    toBufIOWith maxRspBufSize writeBufferRef (\bs -> connSendAll conn bs >> T.tickle th) hdrBdy
-    return (Just s, Nothing) -- fixme: can we tell the actual sent bytes?
+    len <- toBufIOWith maxRspBufSize writeBufferRef (\bs -> connSendAll conn bs >> T.tickle th) hdrBdy
+    return (Just s, Just $ fromIntegral len)
 
 ----------------------------------------------------------------