From fadd7434be6bfc11483dc188aa6e4267e01d5c40 Mon Sep 17 00:00:00 2001
From: Benjamin Barenblat <bbaren@google.com>
Date: Sat, 1 Sep 2018 15:14:49 -0400
Subject: Consolidate buffer-handling functions

---
 src/BTLS/Buffer.hs | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

(limited to 'src/BTLS/Buffer.hs')

diff --git a/src/BTLS/Buffer.hs b/src/BTLS/Buffer.hs
index 1a334e9..7168a10 100644
--- a/src/BTLS/Buffer.hs
+++ b/src/BTLS/Buffer.hs
@@ -15,12 +15,13 @@
 module BTLS.Buffer
   ( unsafeUseAsCBuffer
   , packCUStringLen
+  , onBufferOfMaxSize
   ) where
 
 import Data.ByteString (ByteString)
 import qualified Data.ByteString as ByteString
 import qualified Data.ByteString.Unsafe as ByteString
-import Foreign (Ptr, castPtr)
+import Foreign (Storable(peek), Ptr, alloca, allocaArray, castPtr)
 import Foreign.C.Types
 
 unsafeUseAsCBuffer :: ByteString -> ((Ptr a, CULong) -> IO b) -> IO b
@@ -31,3 +32,21 @@ unsafeUseAsCBuffer bs f =
 packCUStringLen :: Integral n => (Ptr CUChar, n) -> IO ByteString
 packCUStringLen (pStr, len) =
   ByteString.packCStringLen (castPtr pStr, fromIntegral len)
+
+-- | Allocates a buffer, runs a function 'f' to partially fill it, and packs the
+-- filled data into a 'ByteString'. 'f' must write the size of the filled data,
+-- in bytes and not including any trailing null, into its second argument.
+--
+-- If 'f' is safe to use under 'unsafeLocalState', this whole function is safe
+-- to use under 'unsafeLocalState'.
+onBufferOfMaxSize ::
+     (Integral size, Storable size)
+  => Int
+  -> (Ptr CUChar -> Ptr size -> IO ())
+  -> IO ByteString
+onBufferOfMaxSize maxSize f =
+  allocaArray maxSize $ \pOut ->
+    alloca $ \pOutLen -> do
+      f pOut pOutLen
+      outLen <- peek pOutLen
+      packCUStringLen (pOut, outLen)
-- 
cgit v1.2.3