module LibP2P.MultistreamSelect.Negotiation
( NegotiationResult (..)
, ProtocolId
, StreamIO (..)
, negotiateInitiator
, negotiateResponder
, mkByteStreamIO
, mkMemoryStreamPair
, readExactBounded
, closeQuietly
) where
import Control.Concurrent.STM
import Control.Exception (IOException, SomeException, catch)
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import Data.Text (Text)
import Data.Word (Word64, Word8)
import LibP2P.Core.Varint (decodeUvarint, maxVarintBytes)
import LibP2P.MultistreamSelect.Wire
type ProtocolId = Text
maxMessageLength :: Word64
maxMessageLength :: Word64
maxMessageLength = Word64
1024
data NegotiationResult
= Accepted !ProtocolId
| NoProtocol
deriving (Int -> NegotiationResult -> ShowS
[NegotiationResult] -> ShowS
NegotiationResult -> String
(Int -> NegotiationResult -> ShowS)
-> (NegotiationResult -> String)
-> ([NegotiationResult] -> ShowS)
-> Show NegotiationResult
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> NegotiationResult -> ShowS
showsPrec :: Int -> NegotiationResult -> ShowS
$cshow :: NegotiationResult -> String
show :: NegotiationResult -> String
$cshowList :: [NegotiationResult] -> ShowS
showList :: [NegotiationResult] -> ShowS
Show, NegotiationResult -> NegotiationResult -> Bool
(NegotiationResult -> NegotiationResult -> Bool)
-> (NegotiationResult -> NegotiationResult -> Bool)
-> Eq NegotiationResult
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: NegotiationResult -> NegotiationResult -> Bool
== :: NegotiationResult -> NegotiationResult -> Bool
$c/= :: NegotiationResult -> NegotiationResult -> Bool
/= :: NegotiationResult -> NegotiationResult -> Bool
Eq)
data StreamIO = StreamIO
{ StreamIO -> ByteString -> IO ()
streamWrite :: ByteString -> IO ()
, StreamIO -> IO Word8
streamReadByte :: IO Word8
, StreamIO -> Int -> IO ByteString
streamReadChunk :: Int -> IO ByteString
, StreamIO -> IO ()
streamClose :: IO ()
}
mkByteStreamIO :: (ByteString -> IO ()) -> IO Word8 -> IO () -> StreamIO
mkByteStreamIO :: (ByteString -> IO ()) -> IO Word8 -> IO () -> StreamIO
mkByteStreamIO ByteString -> IO ()
write IO Word8
readByte IO ()
close = StreamIO
{ streamWrite :: ByteString -> IO ()
streamWrite = ByteString -> IO ()
write
, streamReadByte :: IO Word8
streamReadByte = IO Word8
readByte
, streamReadChunk :: Int -> IO ByteString
streamReadChunk = \Int
_ -> Word8 -> ByteString
BS.singleton (Word8 -> ByteString) -> IO Word8 -> IO ByteString
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO Word8
readByte
, streamClose :: IO ()
streamClose = IO ()
close
}
mkMemoryStreamPair :: IO (StreamIO, StreamIO)
mkMemoryStreamPair :: IO (StreamIO, StreamIO)
mkMemoryStreamPair = do
queueAtoB <- IO (TQueue Word8)
forall a. IO (TQueue a)
newTQueueIO :: IO (TQueue Word8)
queueBtoA <- newTQueueIO :: IO (TQueue Word8)
let writeToQueue TQueue Word8
q ByteString
bs = (Word8 -> IO ()) -> [Word8] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ (STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> (Word8 -> STM ()) -> Word8 -> IO ()
forall b c a. (b -> c) -> (a -> b) -> a -> c
. TQueue Word8 -> Word8 -> STM ()
forall a. TQueue a -> a -> STM ()
writeTQueue TQueue Word8
q) (ByteString -> [Word8]
BS.unpack ByteString
bs)
readFromQueue TQueue a
q = STM a -> IO a
forall a. STM a -> IO a
atomically (TQueue a -> STM a
forall a. TQueue a -> STM a
readTQueue TQueue a
q)
drainUpTo TQueue a
q Int
k
| Int
k Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= (Int
0 :: Int) = [a] -> STM [a]
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
| Bool
otherwise = do
mb <- TQueue a -> STM (Maybe a)
forall a. TQueue a -> STM (Maybe a)
tryReadTQueue TQueue a
q
case mb of
Maybe a
Nothing -> [a] -> STM [a]
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
Just a
b -> (a
b a -> [a] -> [a]
forall a. a -> [a] -> [a]
:) ([a] -> [a]) -> STM [a] -> STM [a]
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TQueue a -> Int -> STM [a]
drainUpTo TQueue a
q (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1)
readChunkFromQueue TQueue Word8
q Int
n = STM ByteString -> IO ByteString
forall a. STM a -> IO a
atomically (STM ByteString -> IO ByteString)
-> STM ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ do
b <- TQueue Word8 -> STM Word8
forall a. TQueue a -> STM a
readTQueue TQueue Word8
q
rest <- drainUpTo q (n - 1)
pure (BS.pack (b : rest))
pure
( StreamIO (writeToQueue queueAtoB) (readFromQueue queueBtoA) (readChunkFromQueue queueBtoA) (pure ())
, StreamIO (writeToQueue queueBtoA) (readFromQueue queueAtoB) (readChunkFromQueue queueAtoB) (pure ())
)
readChunkSize :: Int
readChunkSize :: Int
readChunkSize = Int
32768
readExactBounded
:: StreamIO
-> Int
-> Int
-> IO (Either String ByteString)
readExactBounded :: StreamIO -> Int -> Int -> IO (Either String ByteString)
readExactBounded StreamIO
stream Int
maxLen Int
n
| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 =
Either String ByteString -> IO (Either String ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String ByteString
forall a b. a -> Either a b
Left (String
"readExactBounded: negative length: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
n))
| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
> Int
maxLen =
Either String ByteString -> IO (Either String ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String ByteString
forall a b. a -> Either a b
Left (String
"readExactBounded: requested " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
n
String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
" bytes exceeds maximum " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> Int -> String
forall a. Show a => a -> String
show Int
maxLen))
| Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 = Either String ByteString -> IO (Either String ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> Either String ByteString
forall a b. b -> Either a b
Right ByteString
BS.empty)
| Bool
otherwise =
(ByteString -> Either String ByteString
forall a b. b -> Either a b
Right (ByteString -> Either String ByteString)
-> ([ByteString] -> ByteString)
-> [ByteString]
-> Either String ByteString
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [ByteString] -> ByteString
BS.concat ([ByteString] -> Either String ByteString)
-> IO [ByteString] -> IO (Either String ByteString)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> IO [ByteString]
go Int
n) IO (Either String ByteString)
-> (IOException -> IO (Either String ByteString))
-> IO (Either String ByteString)
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \(IOException
e :: IOException) ->
Either String ByteString -> IO (Either String ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String ByteString
forall a b. a -> Either a b
Left (String
"readExactBounded: read failed: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> IOException -> String
forall a. Show a => a -> String
show IOException
e))
where
go :: Int -> IO [ByteString]
go :: Int -> IO [ByteString]
go Int
remaining
| Int
remaining Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
0 = [ByteString] -> IO [ByteString]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []
| Bool
otherwise = do
chunk <- StreamIO -> Int -> IO ByteString
streamReadChunk StreamIO
stream (Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
readChunkSize Int
remaining)
if BS.null chunk
then fail "readExactBounded: streamReadChunk returned no bytes"
else (chunk :) <$> go (remaining - BS.length chunk)
closeQuietly :: StreamIO -> IO ()
closeQuietly :: StreamIO -> IO ()
closeQuietly StreamIO
stream = StreamIO -> IO ()
streamClose StreamIO
stream IO () -> (SomeException -> IO ()) -> IO ()
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \(SomeException
_ :: SomeException) -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
readMessage :: StreamIO -> IO (Either String Text)
readMessage :: StreamIO -> IO (Either String Text)
readMessage StreamIO
stream = do
varintResult <- StreamIO -> IO (Either String ByteString)
readVarint StreamIO
stream
case varintResult of
Left String
err -> Either String Text -> IO (Either String Text)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String Text
forall a b. a -> Either a b
Left String
err)
Right ByteString
varintBytes ->
case ByteString -> Either String (Word64, ByteString)
decodeUvarint ByteString
varintBytes of
Left String
err -> Either String Text -> IO (Either String Text)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String Text
forall a b. a -> Either a b
Left String
err)
Right (Word64
len, ByteString
_)
| Word64
len Word64 -> Word64 -> Bool
forall a. Ord a => a -> a -> Bool
> Word64
maxMessageLength ->
Either String Text -> IO (Either String Text)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String Text
forall a b. a -> Either a b
Left String
"readMessage: incoming message too large (max 1024 bytes)")
| Bool
otherwise -> do
payloadOrErr <-
StreamIO -> Int -> Int -> IO (Either String ByteString)
readExactBounded StreamIO
stream (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
maxMessageLength) (Word64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Word64
len)
case payloadOrErr of
Left String
err -> Either String Text -> IO (Either String Text)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String Text
forall a b. a -> Either a b
Left String
err)
Right ByteString
payload ->
case ByteString -> Either String (Text, ByteString)
decodeMessage (ByteString
varintBytes ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
payload) of
Left String
err -> Either String Text -> IO (Either String Text)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String Text
forall a b. a -> Either a b
Left String
err)
Right (Text
msg, ByteString
_) -> Either String Text -> IO (Either String Text)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> Either String Text
forall a b. b -> Either a b
Right Text
msg)
readVarint :: StreamIO -> IO (Either String ByteString)
readVarint :: StreamIO -> IO (Either String ByteString)
readVarint StreamIO
stream =
Int -> [Word8] -> IO (Either String ByteString)
go Int
0 [] IO (Either String ByteString)
-> (IOException -> IO (Either String ByteString))
-> IO (Either String ByteString)
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \(IOException
e :: IOException) ->
Either String ByteString -> IO (Either String ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String ByteString
forall a b. a -> Either a b
Left (String
"readVarint: read failed: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> IOException -> String
forall a. Show a => a -> String
show IOException
e))
where
go :: Int -> [Word8] -> IO (Either String ByteString)
go :: Int -> [Word8] -> IO (Either String ByteString)
go Int
n [Word8]
acc
| Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
maxVarintBytes =
Either String ByteString -> IO (Either String ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> Either String ByteString
forall a b. a -> Either a b
Left String
"readVarint: varint too long (exceeds 9 bytes)")
| Bool
otherwise = do
b <- StreamIO -> IO Word8
streamReadByte StreamIO
stream
if b < 0x80
then pure (Right (BS.pack (reverse (b : acc))))
else go (n + 1) (b : acc)
writeMessage :: StreamIO -> Text -> IO ()
writeMessage :: StreamIO -> Text -> IO ()
writeMessage StreamIO
stream Text
msg = StreamIO -> ByteString -> IO ()
streamWrite StreamIO
stream (Text -> ByteString
encodeMessage Text
msg)
negotiateInitiator :: StreamIO -> [ProtocolId] -> IO NegotiationResult
negotiateInitiator :: StreamIO -> [Text] -> IO NegotiationResult
negotiateInitiator StreamIO
stream [] = do
StreamIO -> Text -> IO ()
writeMessage StreamIO
stream Text
multistreamHeader
result <- StreamIO -> IO (Either String Text)
readMessage StreamIO
stream
case result of
Left String
_ -> NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NegotiationResult
NoProtocol
Right Text
_ -> NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NegotiationResult
NoProtocol
negotiateInitiator StreamIO
stream (Text
firstProto : [Text]
rest) = do
StreamIO -> ByteString -> IO ()
streamWrite StreamIO
stream (Text -> ByteString
encodeMessage Text
multistreamHeader ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Text -> ByteString
encodeMessage Text
firstProto)
headerReply <- StreamIO -> IO (Either String Text)
readMessage StreamIO
stream
case headerReply of
Left String
_ -> NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NegotiationResult
NoProtocol
Right Text
header
| Text
header Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
/= Text
multistreamHeader -> NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NegotiationResult
NoProtocol
| Bool
otherwise -> Text -> IO NegotiationResult -> IO NegotiationResult
awaitReply Text
firstProto ([Text] -> IO NegotiationResult
tryProtocols [Text]
rest)
where
tryProtocols :: [Text] -> IO NegotiationResult
tryProtocols [] = NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NegotiationResult
NoProtocol
tryProtocols (Text
proto : [Text]
remaining) = do
StreamIO -> Text -> IO ()
writeMessage StreamIO
stream Text
proto
Text -> IO NegotiationResult -> IO NegotiationResult
awaitReply Text
proto ([Text] -> IO NegotiationResult
tryProtocols [Text]
remaining)
awaitReply :: Text -> IO NegotiationResult -> IO NegotiationResult
awaitReply Text
proto IO NegotiationResult
onNa = do
result <- StreamIO -> IO (Either String Text)
readMessage StreamIO
stream
case result of
Left String
_ -> NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NegotiationResult
NoProtocol
Right Text
response
| Text
response Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
proto -> NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> NegotiationResult
Accepted Text
proto)
| Text
response Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
== Text
naMessage -> IO NegotiationResult
onNa
| Bool
otherwise -> NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NegotiationResult
NoProtocol
negotiateResponder :: StreamIO -> [ProtocolId] -> IO NegotiationResult
negotiateResponder :: StreamIO -> [Text] -> IO NegotiationResult
negotiateResponder StreamIO
stream [Text]
supported = do
result <- StreamIO -> IO (Either String Text)
readMessage StreamIO
stream
case result of
Left String
_ -> NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NegotiationResult
NoProtocol
Right Text
header
| Text
header Text -> Text -> Bool
forall a. Eq a => a -> a -> Bool
/= Text
multistreamHeader -> NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NegotiationResult
NoProtocol
| Bool
otherwise -> do
StreamIO -> Text -> IO ()
writeMessage StreamIO
stream Text
multistreamHeader
IO NegotiationResult
handleProposals
where
handleProposals :: IO NegotiationResult
handleProposals = do
result <- StreamIO -> IO (Either String Text)
readMessage StreamIO
stream
case result of
Left String
_ -> NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure NegotiationResult
NoProtocol
Right Text
proposal
| Text
proposal Text -> [Text] -> Bool
forall a. Eq a => a -> [a] -> Bool
forall (t :: * -> *) a. (Foldable t, Eq a) => a -> t a -> Bool
`elem` [Text]
supported -> do
StreamIO -> Text -> IO ()
writeMessage StreamIO
stream Text
proposal
NegotiationResult -> IO NegotiationResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Text -> NegotiationResult
Accepted Text
proposal)
| Bool
otherwise -> do
StreamIO -> Text -> IO ()
writeMessage StreamIO
stream Text
naMessage
IO NegotiationResult
handleProposals