{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}

module Simplex.Messaging.Transport.Buffer
  ( TBuffer (..),
    newTBuffer,
    peekBuffered,
    getBuffered,
    withTimedErr,
    getLnBuffered,
    trimCR,
  ) where

import Control.Concurrent.STM
import qualified Control.Exception as E
import Control.Monad (forM_)
import Data.ByteString.Char8 (ByteString)
import qualified Data.ByteString.Char8 as B
import GHC.IO.Exception (IOErrorType (..), IOException (..), ioException)
import System.Timeout (timeout)

data TBuffer = TBuffer
  { TBuffer -> TVar ByteString
buffer :: TVar ByteString,
    TBuffer -> TMVar ()
getLock :: TMVar ()
  }

newTBuffer :: IO TBuffer
newTBuffer :: IO TBuffer
newTBuffer = do
  TVar ByteString
buffer <- ByteString -> IO (TVar ByteString)
forall a. a -> IO (TVar a)
newTVarIO ByteString
""
  TMVar ()
getLock <- () -> IO (TMVar ())
forall a. a -> IO (TMVar a)
newTMVarIO ()
  TBuffer -> IO TBuffer
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure TBuffer {TVar ByteString
buffer :: TVar ByteString
buffer :: TVar ByteString
buffer, TMVar ()
getLock :: TMVar ()
getLock :: TMVar ()
getLock}

withBufferLock :: TBuffer -> IO a -> IO a
withBufferLock :: forall a. TBuffer -> IO a -> IO a
withBufferLock TBuffer {TMVar ()
getLock :: TBuffer -> TMVar ()
getLock :: TMVar ()
getLock} =
  IO () -> IO () -> IO a -> IO a
forall a b c. IO a -> IO b -> IO c -> IO c
E.bracket_
    (STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TMVar () -> STM ()
forall a. TMVar a -> STM a
takeTMVar TMVar ()
getLock)
    (STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TMVar () -> () -> STM ()
forall a. TMVar a -> a -> STM ()
putTMVar TMVar ()
getLock ())

-- | Attempt to read some bytes, appending it to the existing buffer
peekBuffered :: TBuffer -> Int -> IO ByteString -> IO (ByteString, Maybe ByteString)
peekBuffered :: TBuffer
-> Int -> IO ByteString -> IO (ByteString, Maybe ByteString)
peekBuffered tb :: TBuffer
tb@TBuffer {TVar ByteString
buffer :: TBuffer -> TVar ByteString
buffer :: TVar ByteString
buffer} Int
t IO ByteString
getChunk = TBuffer
-> IO (ByteString, Maybe ByteString)
-> IO (ByteString, Maybe ByteString)
forall a. TBuffer -> IO a -> IO a
withBufferLock TBuffer
tb (IO (ByteString, Maybe ByteString)
 -> IO (ByteString, Maybe ByteString))
-> IO (ByteString, Maybe ByteString)
-> IO (ByteString, Maybe ByteString)
forall a b. (a -> b) -> a -> b
$ do
  ByteString
old <- TVar ByteString -> IO ByteString
forall a. TVar a -> IO a
readTVarIO TVar ByteString
buffer
  Maybe ByteString
next_ <- Int -> IO ByteString -> IO (Maybe ByteString)
forall a. Int -> IO a -> IO (Maybe a)
timeout Int
t IO ByteString
getChunk
  Maybe ByteString -> (ByteString -> IO ()) -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ Maybe ByteString
next_ ((ByteString -> IO ()) -> IO ()) -> (ByteString -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \ByteString
next -> STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar ByteString -> ByteString -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar ByteString
buffer (ByteString -> STM ()) -> ByteString -> STM ()
forall a b. (a -> b) -> a -> b
$! ByteString
old ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
next
  (ByteString, Maybe ByteString) -> IO (ByteString, Maybe ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString
old, Maybe ByteString
next_)

getBuffered :: TBuffer -> Int -> Maybe Int -> IO ByteString -> IO ByteString
getBuffered :: TBuffer -> Int -> Maybe Int -> IO ByteString -> IO ByteString
getBuffered tb :: TBuffer
tb@TBuffer {TVar ByteString
buffer :: TBuffer -> TVar ByteString
buffer :: TVar ByteString
buffer} Int
n Maybe Int
t_ IO ByteString
getChunk = TBuffer -> IO ByteString -> IO ByteString
forall a. TBuffer -> IO a -> IO a
withBufferLock TBuffer
tb (IO ByteString -> IO ByteString) -> IO ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ do
  ByteString
b <- Bool -> ByteString -> IO ByteString
readChunks Bool
True (ByteString -> IO ByteString) -> IO ByteString -> IO ByteString
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TVar ByteString -> IO ByteString
forall a. TVar a -> IO a
readTVarIO TVar ByteString
buffer
  let (ByteString
s, ByteString
b') = Int -> ByteString -> (ByteString, ByteString)
B.splitAt Int
n ByteString
b
  STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar ByteString -> ByteString -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar ByteString
buffer (ByteString -> STM ()) -> ByteString -> STM ()
forall a b. (a -> b) -> a -> b
$! ByteString
b'
  -- This would prevent the need to pad auth tag in HTTP2
  -- threadDelay 150
  ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
s
  where
    readChunks :: Bool -> ByteString -> IO ByteString
    readChunks :: Bool -> ByteString -> IO ByteString
readChunks Bool
firstChunk ByteString
b
      | ByteString -> Int
B.length ByteString
b Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
b
      | Bool
otherwise =
          IO ByteString
get IO ByteString -> (ByteString -> IO ByteString) -> IO ByteString
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
            ByteString
"" -> ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
b
            ByteString
s -> Bool -> ByteString -> IO ByteString
readChunks Bool
False (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ ByteString
b ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
s
      where
        get :: IO ByteString
get
          | Bool
firstChunk = IO ByteString
getChunk
          | Bool
otherwise = Maybe Int -> IO ByteString -> IO ByteString
forall a. Maybe Int -> IO a -> IO a
withTimedErr Maybe Int
t_ IO ByteString
getChunk

withTimedErr :: Maybe Int -> IO a -> IO a
withTimedErr :: forall a. Maybe Int -> IO a -> IO a
withTimedErr Maybe Int
t_ IO a
a = case Maybe Int
t_ of
  Just Int
t -> Int -> IO a -> IO (Maybe a)
forall a. Int -> IO a -> IO (Maybe a)
timeout Int
t IO a
a IO (Maybe a) -> (Maybe a -> IO a) -> IO a
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= IO a -> (a -> IO a) -> Maybe a -> IO a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe IO a
forall {a}. IO a
err a -> IO a
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure
  Maybe Int
Nothing -> IO a
a
  where
    err :: IO a
err = IOException -> IO a
forall a. IOException -> IO a
ioException (Maybe Handle
-> IOErrorType
-> String
-> String
-> Maybe CInt
-> Maybe String
-> IOException
IOError Maybe Handle
forall a. Maybe a
Nothing IOErrorType
TimeExpired String
"" String
"get timeout" Maybe CInt
forall a. Maybe a
Nothing Maybe String
forall a. Maybe a
Nothing)

-- This function is only used in test and needs to be improved before it can be used in production,
-- it will never complete if TLS connection is closed before there is newline.
getLnBuffered :: TBuffer -> IO ByteString -> IO ByteString
getLnBuffered :: TBuffer -> IO ByteString -> IO ByteString
getLnBuffered tb :: TBuffer
tb@TBuffer {TVar ByteString
buffer :: TBuffer -> TVar ByteString
buffer :: TVar ByteString
buffer} IO ByteString
getChunk = TBuffer -> IO ByteString -> IO ByteString
forall a. TBuffer -> IO a -> IO a
withBufferLock TBuffer
tb (IO ByteString -> IO ByteString) -> IO ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ do
  ByteString
b <- ByteString -> IO ByteString
readChunks (ByteString -> IO ByteString) -> IO ByteString -> IO ByteString
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< TVar ByteString -> IO ByteString
forall a. TVar a -> IO a
readTVarIO TVar ByteString
buffer
  let (ByteString
s, ByteString
b') = (Char -> Bool) -> ByteString -> (ByteString, ByteString)
B.break (Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'\n') ByteString
b
  STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar ByteString -> ByteString -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar ByteString
buffer (ByteString -> STM ()) -> ByteString -> STM ()
forall a b. (a -> b) -> a -> b
$! Int -> ByteString -> ByteString
B.drop Int
1 ByteString
b' -- drop '\n' we made a break at
  ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString
trimCR ByteString
s
  where
    readChunks :: ByteString -> IO ByteString
    readChunks :: ByteString -> IO ByteString
readChunks ByteString
b
      | Char -> ByteString -> Bool
B.elem Char
'\n' ByteString
b = ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
b
      | Bool
otherwise = ByteString -> IO ByteString
readChunks (ByteString -> IO ByteString)
-> (ByteString -> ByteString) -> ByteString -> IO ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ByteString
b ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<>) (ByteString -> IO ByteString) -> IO ByteString -> IO ByteString
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< IO ByteString
getChunk

-- | Trim trailing CR from ByteString.
trimCR :: ByteString -> ByteString
trimCR :: ByteString -> ByteString
trimCR ByteString
"" = ByteString
""
trimCR ByteString
s = if ByteString -> Char
B.last ByteString
s Char -> Char -> Bool
forall a. Eq a => a -> a -> Bool
== Char
'\r' then HasCallStack => ByteString -> ByteString
ByteString -> ByteString
B.init ByteString
s else ByteString
s