Reass.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE RecordWildCards #-}

module Network.QUIC.Stream.Reass (
    takeRecvStreamQwithSize,
    putRxStreamData,
    FlowCntl (..),
    tryReassemble,
) where

import qualified Data.ByteString as BS
import Data.Sequence (Seq)
import qualified Data.Sequence as Seq

import Network.QUIC.Imports
import Network.QUIC.Stream.Frag
import Network.QUIC.Stream.Misc
import Network.QUIC.Stream.Queue
import qualified Network.QUIC.Stream.Skew as Skew
import Network.QUIC.Stream.Types
import Network.QUIC.Types

----------------------------------------------------------------

getEndOfStream :: Stream -> IO Bool
getEndOfStream Stream{..} = readIORef $ endOfStream streamRecvQ

setEndOfStream :: Stream -> IO ()
setEndOfStream Stream{..} = writeIORef (endOfStream streamRecvQ) True

readPendingData :: Stream -> IO (Maybe ByteString)
readPendingData Stream{..} = readIORef $ pendingData streamRecvQ

writePendingData :: Stream -> ByteString -> IO ()
writePendingData Stream{..} bs = writeIORef (pendingData streamRecvQ) $ Just bs

clearPendingData :: Stream -> IO ()
clearPendingData Stream{..} = writeIORef (pendingData streamRecvQ) Nothing

----------------------------------------------------------------

takeRecvStreamQwithSize
    :: Stream
    -> Int
    -- ^ Number of bytes to receive.
    -> IO ByteString
takeRecvStreamQwithSize strm siz0 = do
    eos <- getEndOfStream strm
    if eos
        then return ""
        else do
            mb <- readPendingData strm
            case mb of
                Nothing -> do
                    b0 <- takeRecvStreamQ strm
                    if b0 == ""
                        then do
                            setEndOfStream strm
                            return ""
                        else handleBytes b0
                Just b0 -> do
                    clearPendingData strm
                    handleBytes b0
  where
    handleBytes b0 =
        let len = BS.length b0
         in case len `compare` siz0 of
                LT -> tryRead (siz0 - len) (b0 :)
                EQ -> return b0
                GT -> do
                    let (b1, b2) = BS.splitAt siz0 b0
                    writePendingData strm b2
                    return b1
    tryRead siz build = do
        mb <- tryTakeRecvStreamQ strm
        case mb of
            Nothing -> return $ BS.concat $ build []
            Just b -> do
                if b == ""
                    then do
                        setEndOfStream strm
                        return $ BS.concat $ build []
                    else do
                        let len = BS.length b
                        case len `compare` siz of
                            LT -> tryRead (siz - len) (build . (b :))
                            EQ -> return $ BS.concat $ build [b]
                            GT -> do
                                let (b1, b2) = BS.splitAt siz b
                                writePendingData strm b2
                                return $ BS.concat $ build [b1]

----------------------------------------------------------------
----------------------------------------------------------------

data FlowCntl = OverLimit | Duplicated | Reassembled

putRxStreamData :: Stream -> RxStreamData -> IO FlowCntl
putRxStreamData s rx@(RxStreamData _ off len _) = do
    lim <- getRxMaxStreamData s
    if len + off > lim
        then return OverLimit
        else do
            dup <- tryReassemble s rx put putFin
            if dup
                then return Duplicated
                else return Reassembled
  where
    put "" = return ()
    put d = do
        addRxStreamData s $ BS.length d
        putRecvStreamQ s d
    putFin = putRecvStreamQ s ""

-- fin of StreamState off fin means see-fin-already.
-- return value indicates duplication
tryReassemble
    :: Stream -> RxStreamData -> (StreamData -> IO ()) -> IO () -> IO Bool
tryReassemble Stream{} (RxStreamData "" _ _ False) _ _ = return True
tryReassemble Stream{..} x@(RxStreamData "" off _ True) _ putFin = do
    si0@(StreamState off0 fin0) <- readIORef streamStateRx
    let si1 = si0{streamFin = True}
    if fin0
        then do
            -- stdoutLogger "Illegal Fin" -- fixme
            return True
        else case off `compare` off0 of
            LT -> return True
            EQ -> do
                writeIORef streamStateRx si1
                putFin
                return False
            GT -> do
                writeIORef streamStateRx si1
                atomicModifyIORef'' streamReass (Skew.insert x)
                return False
tryReassemble Stream{..} x@(RxStreamData dat off len False) put putFin = do
    si0@(StreamState off0 _) <- readIORef streamStateRx
    case off `compare` off0 of
        LT -> return True
        EQ -> do
            put dat
            loop si0 (off0 + len)
            return False
        GT -> do
            atomicModifyIORef'' streamReass (Skew.insert x)
            return False
  where
    loop si0 xff = do
        mrxs <- atomicModifyIORef' streamReass (Skew.deleteMinIf xff)
        case mrxs of
            Nothing -> writeIORef streamStateRx si0{streamOffset = xff}
            Just rxs -> do
                mapM_ (put . rxstrmData) rxs
                let xff1 = nextOff rxs
                if hasFin rxs
                    then do
                        putFin
                    else do
                        loop si0 xff1
tryReassemble Stream{..} x@(RxStreamData dat off len True) put putFin = do
    si0@(StreamState off0 fin0) <- readIORef streamStateRx
    let si1 = si0{streamFin = True}
    if fin0
        then return True
        else case off `compare` off0 of
            LT -> return True
            EQ -> do
                let off1 = off0 + len
                writeIORef streamStateRx si1{streamOffset = off1}
                put dat
                putFin
                return False
            GT -> do
                writeIORef streamStateRx si1
                atomicModifyIORef'' streamReass (Skew.insert x)
                return False

hasFin :: Seq RxStreamData -> Bool
hasFin s = case Seq.viewr s of
    Seq.EmptyR -> False
    _ Seq.:> x -> rxstrmFin x