mirror of
https://github.com/simplex-chat/simplexmq.git
synced 2026-09-01 20:18:26 +00:00
Evgeny Poberezkin
94540a2c71
* servers: add TCP timeouts to avoid memory leaks * fix tests * only use RecvTimeOut * servers: simple timeout for TCP transport * revert dependency change * simplify * simplify * simplify 2
77 lines
2.4 KiB
Haskell
77 lines
2.4 KiB
Haskell
{-# LANGUAGE LambdaCase #-}
|
|
{-# LANGUAGE NamedFieldPuns #-}
|
|
{-# LANGUAGE OverloadedStrings #-}
|
|
|
|
module Simplex.Messaging.Transport.Buffer where
|
|
|
|
import Control.Concurrent.STM
|
|
import qualified Control.Exception as E
|
|
import Data.ByteString.Char8 (ByteString)
|
|
import qualified Data.ByteString.Char8 as B
|
|
import System.Timeout (timeout)
|
|
import GHC.IO.Exception (ioException, IOException (..), IOErrorType (..))
|
|
|
|
data TBuffer = TBuffer
|
|
{ buffer :: TVar ByteString,
|
|
getLock :: TMVar ()
|
|
}
|
|
|
|
newTBuffer :: STM TBuffer
|
|
newTBuffer = do
|
|
buffer <- newTVar ""
|
|
getLock <- newTMVar ()
|
|
pure TBuffer {buffer, getLock}
|
|
|
|
withBufferLock :: TBuffer -> IO a -> IO a
|
|
withBufferLock TBuffer {getLock} =
|
|
E.bracket_
|
|
(atomically $ takeTMVar getLock)
|
|
(atomically $ putTMVar getLock ())
|
|
|
|
getBuffered :: TBuffer -> Int -> Maybe Int -> IO ByteString -> IO ByteString
|
|
getBuffered tb@TBuffer {buffer} n t_ getChunk = withBufferLock tb $ do
|
|
b <- readChunks True =<< readTVarIO buffer
|
|
let (s, b') = B.splitAt n b
|
|
atomically $ writeTVar buffer $! b'
|
|
-- This would prevent the need to pad auth tag in HTTP2
|
|
-- threadDelay 150
|
|
pure s
|
|
where
|
|
readChunks :: Bool -> ByteString -> IO ByteString
|
|
readChunks firstChunk b
|
|
| B.length b >= n = pure b
|
|
| otherwise =
|
|
get >>= \case
|
|
"" -> pure b
|
|
s -> readChunks False $ b <> s
|
|
where
|
|
get
|
|
| firstChunk = getChunk
|
|
| otherwise = withTimedErr t_ getChunk
|
|
|
|
withTimedErr :: Maybe Int -> IO a -> IO a
|
|
withTimedErr t_ a = case t_ of
|
|
Just t -> timeout t a >>= maybe err pure
|
|
Nothing -> a
|
|
where
|
|
err = ioException (IOError Nothing TimeExpired "" "get timeout" Nothing 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 tb@TBuffer {buffer} getChunk = withBufferLock tb $ do
|
|
b <- readChunks =<< readTVarIO buffer
|
|
let (s, b') = B.break (== '\n') b
|
|
atomically $ writeTVar buffer $! B.drop 1 b' -- drop '\n' we made a break at
|
|
pure $ trimCR s
|
|
where
|
|
readChunks :: ByteString -> IO ByteString
|
|
readChunks b
|
|
| B.elem '\n' b = pure b
|
|
| otherwise = readChunks . (b <>) =<< getChunk
|
|
|
|
-- | Trim trailing CR from ByteString.
|
|
trimCR :: ByteString -> ByteString
|
|
trimCR "" = ""
|
|
trimCR s = if B.last s == '\r' then B.init s else s
|