tricu

ByteStream.hs (4710B)

---

 {-# LANGUAGE BangPatterns #-}
 -- Native bulk bytes; source owns paths, digest selection and response policy.
 module IOV2.ByteStream (snapshot, snapshotWith, transfer) where
 
 import qualified IOV2.Resources as R
 import Control.Exception
 import Control.Monad (unless, when)
 import Crypto.Hash (Context, SHA256, hashInit, hashUpdate, hashFinalize)
 import qualified Data.ByteString as BS
 import Data.IORef
 import System.Directory (removeFile)
 import System.IO
 
 -- A whole-stream SHA-256 snapshot, including an EOF probe at the byte budget.
 -- Register before consuming input. The exposed file has no write/resize method;
 -- only the private handle can prepare it, before its token is returned.
 snapshot :: R.Runtime -> R.Scope -> R.Token -> FilePath -> Integer -> String
   -> IO (Either R.Problem (R.Token, Integer))
 snapshot = snapshotWith hClose
 
 snapshotWith :: (Handle -> IO ()) -> R.Runtime -> R.Scope -> R.Token -> FilePath
   -> Integer -> String -> IO (Either R.Problem (R.Token, Integer))
 snapshotWith release runtime scope reader directory limit expected
   | limit < 0 || limit > R.positionLimit = pure (Left (fault "invalid-bound" "invalid byte limit" "none"))
   | length expected /= 64 || any (`notElem` ("0123456789abcdef" :: String)) expected =
       pure (Left (fault "invalid-argument" "expected lowercase SHA-256" "none"))
   | null directory = pure (Left (fault "invalid-argument" "empty snapshot directory" "none"))
   | otherwise = R.withReaderLease runtime scope reader (\receive -> mask $ \restore -> do
       slot <- newIORef Nothing
       acquired <- R.acquire runtime scope $ do
         (path,h) <- openBinaryTempFile directory ".tricu-bytes-"
         (do
           removeFile path
           hSetBuffering h NoBuffering
           writeIORef slot (Just h)
           pure (R.File (R.FileOps (Just (BS.hGetSome h)) Nothing
             (Just (hSeek h AbsoluteSeek)) (Just (hTell h)) (hFileSize h)
             Nothing Nothing (release h)))) `onException` hClose h
       token <- either R.resourceFault pure acquired
       Just h <- readIORef slot
       -- A lease protects the registered spool throughout preparation. Exceptions
       -- leave it to scope settlement, retaining cancellation/cleanup precedence.
       result <- restore $ R.withReaderLease runtime scope token $ \_ -> do
         let copy !count !context = do
               allowInterrupt
               bytes <- receive (fromInteger (min (toInteger R.chunkLimit) (limit-count+1)))
               if BS.null bytes then pure (count,context) else do
                 let nextCount = count + toInteger (BS.length bytes)
                 when (nextCount > limit) (R.resourceFault (fault "invalid-bound" "stream exceeds byte limit" "unknown"))
                 BS.hPut h bytes
                 next <- evaluate (hashUpdate context bytes)
                 copy nextCount next
         (count,context) <- copy 0 (hashInit :: Context SHA256)
         unless (show (hashFinalize context) == expected)
           (R.resourceFault (fault "integrity-failed" "SHA-256 mismatch" "unknown"))
         hSeek h AbsoluteSeek 0
         pure count
       case result of
         Right count -> pure (token,count)
         Left primary -> do
           closed <- R.closeResource runtime scope token
           R.resourceFault $ case closed of
             Right () -> primary
             Left cleanup -> primary { R.detail = take 240 (R.detail primary) ++ "; cleanup: " ++ take 240 (show cleanup) })
   where fault category detail progress = R.Problem category "snapshot" detail progress
 
 -- Exact byte count, starting at the current cursor; no EOF probe. Both leases
 -- cover the transfer. A failure may follow partial writes: callers close rather
 -- than retrying or emitting a second response. No encoded byte trees are built.
 transfer :: R.Runtime -> R.Scope -> R.Token -> R.Token -> Integer -> IO (Either R.Problem ())
 transfer runtime scope reader writer count
   | count < 0 || count > R.positionLimit = pure (Left (R.Problem "invalid-bound" "transfer" "invalid byte count" "none"))
   | otherwise = R.withReaderLease runtime scope reader (\receive ->
       R.withWriterLease runtime scope writer (\send -> do
         let sendAll bytes = unless (BS.null bytes) $ do
               n <- send bytes
               sendAll (BS.drop n bytes)
             copy !remaining = when (remaining > 0) $ do
               allowInterrupt
               bytes <- receive (fromInteger (min remaining (toInteger R.chunkLimit)))
               when (BS.null bytes) (R.resourceFault (R.Problem "invalid-data" "transfer" "unexpected EOF" "unknown"))
               sendAll bytes
               copy (remaining - toInteger (BS.length bytes))
         copy count) >>= either R.resourceFault pure)