Transport-Level Networking for Haskell
+Zero-copy Storable serialization. Reliable UDP transport. Effect-abstracted design for pure testing.
+Quick Start · Networking · Serialization · Testing · Architecture
++ +[](https://github.com/gondola-bros-entertainment/gbnet-hs/actions/workflows/ci.yml) + + + +
++ MIT License · Gondola Bros Entertainment +
diff --git a/bench/Main.hs b/bench/Main.hs new file mode 100644 --- /dev/null +++ b/bench/Main.hs @@ -0,0 +1,377 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} + +module Main where + +import Control.DeepSeq (NFData (..)) +import Criterion.Main +import qualified Data.ByteString as BS +import Data.List (foldl') +import Data.Word (Word16, Word32, Word64) +import Foreign.Storable (Storable (..)) +import GBNet.Channel + ( Channel, + channelSend, + defaultChannelConfig, + getRetransmitMessages, + newChannel, + onMessageReceived, + unreliableConfig, + ) +import GBNet.Class (MonoTime (..)) +import GBNet.Congestion + ( batchMessages, + ccRefillBudget, + ccUpdate, + cwOnAck, + cwOnLoss, + newCongestionController, + newCongestionWindow, + unbatchMessages, + ) +import GBNet.Config (defaultNetworkConfig) +import GBNet.Connection + ( Connection, + drainSendQueue, + newConnection, + sendMessage, + ) +import GBNet.Fragment + ( FragmentHeader (..), + deserializeFragmentHeader, + fragmentMessage, + newFragmentAssembler, + processFragment, + serializeFragmentHeader, + ) +import GBNet.Packet + ( PacketHeader (..), + PacketType (Payload), + deserializeHeader, + serializeHeader, + ) +import GBNet.Reliability + ( ReliableEndpoint, + SBEntry (..), + SentPacketRecord (..), + SequenceBuffer, + newReliableEndpoint, + newSequenceBuffer, + onPacketSent, + onPacketsReceived, + processAcks, + sbExists, + sbInsert, + sbInsertMany, + updateRtt, + ) +import GBNet.Replication.Delta (DeltaTracker, NetworkDelta (..), deltaEncode, newDeltaTracker) +import GBNet.Replication.Interest (InterestManager (..), newRadiusInterest) +import GBNet.Security (appendCrc32, validateAndStripCrc32) +import GBNet.Serialize (castPtr, deserialize, plusPtr, serialize) +import GBNet.Serialize.TH (deriveStorable) +import GBNet.Types (ChannelId (..), MessageId (..), SequenceNum (..)) + +-------------------------------------------------------------------------------- +-- TH-derived benchmark types +-------------------------------------------------------------------------------- + +-- Storable-based serialization +data Vec3S = Vec3S !Float !Float !Float + deriving (Eq, Show) + +deriveStorable ''Vec3S + +-- Nested type - demonstrates composition +data Transform = Transform !Vec3S !Float -- position + rotation angle + deriving (Eq, Show) + +deriveStorable ''Transform + +-------------------------------------------------------------------------------- +-- NFData instances for benchmark-only types +-------------------------------------------------------------------------------- + +instance NFData Vec3S where rnf (Vec3S x y z) = rnf x `seq` rnf y `seq` rnf z + +instance NFData Transform where rnf (Transform v r) = rnf v `seq` rnf r + +-- | Trivial delta instance for benchmarking tracker mechanics. +-- Uses full state as delta (no field-level compression). +instance NetworkDelta Vec3S where + type Delta Vec3S = Vec3S + diff current _baseline = current + apply _state delta = delta + +-------------------------------------------------------------------------------- +-- Setup helpers +-------------------------------------------------------------------------------- + +-- | A sample packet header for benchmarking. +sampleHeader :: PacketHeader +sampleHeader = + PacketHeader + { packetType = Payload, + sequenceNum = SequenceNum 42, + ack = SequenceNum 40, + ackBitfield = 0xDEADBEEF + } + +-- | Pre-serialized header bytes. +headerBytes :: BS.ByteString +headerBytes = serializeHeader sampleHeader + +-- | A sample fragment header for benchmarking. +sampleFragmentHeader :: FragmentHeader +sampleFragmentHeader = + FragmentHeader + { fhMessageId = MessageId 0xDEADBEEF, + fhFragmentIndex = 3, + fhFragmentCount = 10 + } + +-- | Pre-serialized fragment header bytes. +fragmentHeaderBytes :: BS.ByteString +fragmentHeaderBytes = serializeFragmentHeader sampleFragmentHeader + +-- | 64-byte payload for channel benchmarks. +payload64 :: BS.ByteString +payload64 = BS.replicate 64 0xAB + +-- | 1KB payload for fragment benchmarks. +payload1k :: BS.ByteString +payload1k = BS.replicate 1024 0xCD + +-- | 64-byte payload with CRC appended. +payload64WithCrc :: BS.ByteString +payload64WithCrc = appendCrc32 payload64 + +-- | Build a fresh Connection for benchmarking. +buildConnection :: Connection +buildConnection = newConnection defaultNetworkConfig 0x12345678 (MonoTime 0) + +-- | MTU for fragmentation benchmarks. +benchMtu :: Int +benchMtu = 1200 + +-- | Build a ReliableEndpoint with N in-flight sent packets. +buildEndpointWithInFlight :: Int -> ReliableEndpoint +buildEndpointWithInFlight n = + foldl' sendOne (newReliableEndpoint 256) [0 .. fromIntegral (n - 1)] + where + sendOne ep i = + let seqNum = SequenceNum i + sendTime = MonoTime (fromIntegral i * 1000000) -- 1ms apart + in onPacketSent seqNum sendTime (ChannelId 0) seqNum 64 ep + +-- | Build a SequenceBuffer with N sequential entries. +buildSequenceBuffer :: Int -> SequenceBuffer () +buildSequenceBuffer n = + foldl' (\buf i -> sbInsert (SequenceNum (fromIntegral i)) () buf) (newSequenceBuffer 256) [0 .. n - 1] + +-- | Build a Channel pre-loaded with N pending reliable messages. +buildChannelWithPending :: Int -> Channel +buildChannelWithPending n = + let ch = newChannel (ChannelId 0) defaultChannelConfig + in foldl' sendMsg ch [0 .. n - 1] + where + sendMsg ch i = + case channelSend payload64 (MonoTime (fromIntegral i * 1000000)) ch of + Right (_, ch') -> ch' + Left _ -> ch + +-------------------------------------------------------------------------------- +-- Main +-------------------------------------------------------------------------------- + +main :: IO () +main = + defaultMain + [ -- Group 1: Reliability (ACK processing) + bgroup + "reliability" + [ env (pure (newReliableEndpoint 256, map SequenceNum [0 .. 99])) $ \ ~(ep, seqs) -> + bench "onPacketsReceived/100-sequential" $ + nf (onPacketsReceived seqs) ep, + env (pure (newReliableEndpoint 256, map (\i -> SequenceNum (i * 3)) [0 .. 99])) $ \ ~(ep, seqs) -> + bench "onPacketsReceived/100-gaps" $ + nf (onPacketsReceived seqs) ep, + env (pure (newReliableEndpoint 256)) $ \ep -> + bench "onPacketsReceived/single" $ + nf (onPacketsReceived [SequenceNum 42]) ep, + env (pure (buildEndpointWithInFlight 10)) $ \ep -> + bench "processAcks/10-inflight" $ + nf + (processAcks (SequenceNum 9) 0x1FF (MonoTime 50000000)) + ep, + env (pure (buildEndpointWithInFlight 100)) $ \ep -> + bench "processAcks/100-inflight" $ + nf + (processAcks (SequenceNum 99) maxBound (MonoTime 500000000)) + ep, + env (pure (newReliableEndpoint 256)) $ \ep -> + bench "updateRtt" $ + nf (updateRtt 25.0 . updateRtt 30.0) ep + ], + -- Group 2: SequenceBuffer + bgroup + "sequencebuffer" + [ env (pure (newSequenceBuffer 256 :: SequenceBuffer ())) $ \buf -> + bench "sbInsert/sequential" $ + nf + ( \b -> + foldl' + (\acc i -> sbInsert (SequenceNum i) () acc) + b + [0 .. 99] + ) + buf, + env (pure (newSequenceBuffer 256 :: SequenceBuffer ())) $ \buf -> + bench "sbInsert/wraparound" $ + nf + ( \b -> + foldl' + (\acc i -> sbInsert (SequenceNum (maxBound - 50 + i)) () acc) + b + [0 .. 99] + ) + buf, + env (pure (newSequenceBuffer 256 :: SequenceBuffer (), [(SequenceNum i, ()) | i <- [0 .. 99]])) $ \ ~(buf, items) -> + bench "sbInsertMany/100" $ + nf (sbInsertMany items) buf, + env (pure (buildSequenceBuffer 100)) $ \buf -> + bench "sbExists/hit" $ + whnf (sbExists (SequenceNum 50)) buf, + env (pure (buildSequenceBuffer 100)) $ \buf -> + bench "sbExists/miss" $ + whnf (sbExists (SequenceNum 200)) buf + ], + -- Group 3: Channel (message pipeline) + bgroup + "channel" + [ env (pure (newChannel (ChannelId 0) defaultChannelConfig)) $ \ch -> + bench "channelSend/reliable" $ + nf (channelSend payload64 (MonoTime 1000000)) ch, + env (pure (newChannel (ChannelId 0) unreliableConfig)) $ \ch -> + bench "channelSend/unreliable" $ + nf (channelSend payload64 (MonoTime 1000000)) ch, + env (pure (newChannel (ChannelId 0) defaultChannelConfig)) $ \ch -> + bench "onMessageReceived/ordered" $ + nf + ( \c -> + foldl' + (\acc i -> onMessageReceived (SequenceNum i) payload64 (MonoTime 1000000) acc) + c + [0 .. 49] + ) + ch, + env (pure (buildChannelWithPending 50, MonoTime 10000000000)) $ \ ~(ch, now) -> + bench "getRetransmitMessages/50-pending" $ + nf (\(c, t) -> getRetransmitMessages t 100.0 c) (ch, now) + ], + -- Group 4: Packet Header + bgroup + "packetheader" + [ env (pure sampleHeader) $ \hdr -> + bench "serialize" $ nf serializeHeader hdr, + env (pure headerBytes) $ \bs -> + bench "deserialize" $ nf deserializeHeader bs + ], + -- Group 5: Fragment Header + bgroup + "fragmentheader" + [ env (pure sampleFragmentHeader) $ \hdr -> + bench "serialize" $ nf serializeFragmentHeader hdr, + env (pure fragmentHeaderBytes) $ \bs -> + bench "deserialize" $ nf deserializeFragmentHeader bs + ], + -- Group 6: Storable serialization (flat and nested) + bgroup + "storable" + [ env (pure (Vec3S 1.0 (-2.5) 100.0)) $ \v -> + bench "vec3/serialize" $ nf serialize v, + env (pure (serialize (Vec3S 1.0 (-2.5) 100.0))) $ \bs -> + bench "vec3/deserialize" $ nf (deserialize :: BS.ByteString -> Either String Vec3S) bs, + env (pure (Transform (Vec3S 1.0 2.0 3.0) 45.0)) $ \t -> + bench "transform/serialize" $ nf serialize t, + env (pure (serialize (Transform (Vec3S 1.0 2.0 3.0) 45.0))) $ \bs -> + bench "transform/deserialize" $ nf (deserialize :: BS.ByteString -> Either String Transform) bs + ], + -- Group 7: Connection operations + bgroup + "connection" + [ env (pure buildConnection) $ \conn -> + bench "sendMessage/64B" $ + nf (sendMessage (ChannelId 0) payload64 (MonoTime 1000000)) conn, + env (pure buildConnection) $ \conn -> + bench "drainSendQueue" $ + nf drainSendQueue conn + ], + -- Group 8: Fragmentation + bgroup + "fragment" + [ env (pure payload1k) $ \payload -> + bench "fragmentMessage/1KB" $ + nf (fragmentMessage (MessageId 1) payload) benchMtu, + env (pure (newFragmentAssembler 5000.0 256)) $ \assembler -> + bench "processFragment/single" $ + nf (processFragment payload64 (MonoTime 1000000)) assembler + ], + -- Group 9: Security (CRC32C) + bgroup + "security" + [ env (pure payload64) $ \payload -> + bench "crc32c/append/64B" $ nf appendCrc32 payload, + env (pure payload1k) $ \payload -> + bench "crc32c/append/1KB" $ nf appendCrc32 payload, + env (pure payload64WithCrc) $ \payload -> + bench "crc32c/validate/64B" $ nf validateAndStripCrc32 payload + ], + -- Group 10: Congestion control + bgroup + "congestion" + [ env (pure (newCongestionController 60.0 0.1 250.0 10000.0)) $ \cc -> + bench "ccUpdate/good" $ + nf (ccUpdate 0.02 50.0 (MonoTime 1000000000)) cc, + env (pure (newCongestionController 60.0 0.1 250.0 10000.0)) $ \cc -> + bench "ccUpdate/bad" $ + nf (ccUpdate 0.15 300.0 (MonoTime 1000000000)) cc, + env (pure (newCongestionController 60.0 0.1 250.0 10000.0)) $ \cc -> + bench "ccRefillBudget" $ + nf (ccRefillBudget 1200) cc, + env (pure (newCongestionWindow 1200)) $ \cw -> + bench "cwOnAck/1200B" $ + nf (cwOnAck 1200) cw, + env (pure (newCongestionWindow 1200)) $ \cw -> + bench "cwOnLoss" $ + nf cwOnLoss cw + ], + -- Group 11: Delta replication + bgroup + "delta" + [ env (pure (newDeltaTracker 64 :: DeltaTracker Vec3S)) $ \dt -> + bench "deltaEncode" $ + nf (\t -> deltaEncode 0 (Vec3S 1.0 2.0 3.0) t) dt + ], + -- Group 12: Interest manager + bgroup + "interest" + [ env (pure (newRadiusInterest 100.0)) $ \interest -> + bench "relevant/inRange" $ + nf (\i -> relevant i (50.0, 50.0, 0.0) (100.0, 100.0, 0.0)) interest, + env (pure (newRadiusInterest 100.0)) $ \interest -> + bench "relevant/outOfRange" $ + nf (\i -> relevant i (0.0, 0.0, 0.0) (500.0, 500.0, 0.0)) interest + ], + -- Group 13: Message batching + bgroup + "batching" + [ env (pure (replicate 10 payload64)) $ \msgs -> + bench "batchMessages/10x64B" $ + nf (\m -> batchMessages m 1200) msgs, + env (pure (batchMessages (replicate 10 payload64) 1200)) $ \ ~batches -> + bench "unbatchMessages/10x64B" $ + nf (map unbatchMessages) batches + ] + ] diff --git a/gbnet-hs.cabal b/gbnet-hs.cabal new file mode 100644 --- /dev/null +++ b/gbnet-hs.cabal @@ -0,0 +1,117 @@ +cabal-version: 3.0 +name: gbnet-hs +version: 0.1.0.0 +synopsis: Transport-level networking library with zero-copy Storable serialization +description: + A transport-level networking library providing reliable UDP with + zero-copy Storable serialization, dual-layer congestion control, and + effect-abstracted design for pure deterministic testing. + +license: MIT +license-file: LICENSE +author: Devon Tomlin +maintainer: devon.tomlin@novavero.ai +homepage: https://github.com/Gondola-Bros-Entertainment/gbnet-hs +bug-reports: https://github.com/Gondola-Bros-Entertainment/gbnet-hs/issues +category: Network, Game +stability: experimental +build-type: Simple +tested-with: GHC == 9.6.7 +extra-doc-files: + README.md + CHANGELOG.md + +library + exposed-modules: + GBNet + GBNet.Serialize + GBNet.Serialize.TH + GBNet.Packet + GBNet.Util + GBNet.Reliability + GBNet.Channel + GBNet.Config + GBNet.Congestion + GBNet.Stats + GBNet.Connection + GBNet.Fragment + GBNet.Security + GBNet.Simulator + GBNet.Socket + GBNet.Peer + GBNet.Replication.Delta + GBNet.Replication.Interest + GBNet.Replication.Priority + GBNet.Replication.Interpolation + GBNet.Class + GBNet.Types + GBNet.TestNet + GBNet.Net + GBNet.Net.IO + + other-modules: + GBNet.ZeroCopy + GBNet.Peer.Internal + GBNet.Peer.Protocol + GBNet.Peer.Handshake + GBNet.Peer.Migration + + build-depends: + base >= 4.16 && < 5 + , bytestring >= 0.11 && < 0.13 + , containers >= 0.6 && < 0.8 + , crc32c >= 0.2 && < 0.3 + , deepseq >= 1.4 && < 1.6 + , mtl >= 2.2 && < 2.4 + , network >= 3.1 && < 3.3 + , stm >= 2.5 && < 2.6 + , template-haskell >= 2.18 && < 2.23 + , transformers >= 0.5 && < 0.7 + , optics >= 0.4.2 && < 0.5 + , optics-th >= 0.4.1 && < 0.5 + , vector >= 0.12 && < 0.14 + + hs-source-dirs: src + default-language: Haskell2010 + -- Aggressive specialisation monomorphises MonadNetwork/MonadTime at call + -- sites, eliminating dictionary overhead. Users wanting -O2 can add it + -- in their own cabal.project (package gbnet-hs ghc-options: -O2). + ghc-options: + -Wall + -Wcompat + -Wincomplete-record-updates + -Wincomplete-uni-patterns + -fspecialise-aggressively + -fexpose-all-unfoldings + +test-suite gbnet-test + type: exitcode-stdio-1.0 + main-is: Main.hs + hs-source-dirs: test + build-depends: + base >= 4.16 && < 5 + , gbnet-hs + , bytestring >= 0.11 && < 0.13 + , network >= 3.1 && < 3.3 + , QuickCheck >= 2.14 && < 2.16 + , containers >= 0.6 && < 0.8 + default-language: Haskell2010 + ghc-options: -Wall -Wcompat + +benchmark gbnet-bench + type: exitcode-stdio-1.0 + main-is: Main.hs + hs-source-dirs: bench + default-language: Haskell2010 + ghc-options: -rtsopts -Wall -Wcompat + build-depends: + base >= 4.16 && < 5 + , bytestring >= 0.11 && < 0.13 + , criterion >= 1.5 && < 1.7 + , deepseq >= 1.4 && < 1.6 + , gbnet-hs + +source-repository head + type: git + location: https://github.com/Gondola-Bros-Entertainment/gbnet-hs + branch: main diff --git a/src/GBNet.hs b/src/GBNet.hs new file mode 100644 --- /dev/null +++ b/src/GBNet.hs @@ -0,0 +1,250 @@ +-- | +-- Module : GBNet +-- Description : Game networking library with zero-copy Storable serialization +-- +-- This is the main entry point for gbnet-hs. Import this module for +-- convenient access to the full API. +-- +-- @ +-- import GBNet +-- +-- main = do +-- let addr = localhost 7777 +-- Right (peer, sock) <- newPeer addr defaultNetworkConfig =<< getMonoTimeIO +-- netState <- newNetState sock addr +-- evalNetT (gameLoop peer) netState +-- @ +-- +-- For explicit imports, use the individual modules: +-- +-- @ +-- import GBNet.Class (MonadNetwork, MonadTime) +-- import GBNet.Net (NetT, runNetT, NetState) +-- import GBNet.Peer (NetPeer, peerTick, PeerEvent(..)) +-- @ +module GBNet + ( -- * Effect Classes + MonadTime (..), + MonadNetwork (..), + MonoTime (..), + NetError (..), + + -- * Domain Types + ChannelId (..), + channelIdToInt, + SequenceNum (..), + MessageId (..), + + -- * Network Monad + NetT, + runNetT, + evalNetT, + execNetT, + NetState, + newNetState, + + -- * IO Backend + initNetState, + getMonoTimeIO, + + -- * Socket Address Utilities + SockAddr (..), + localhost, + ipv4, + anyAddr, + + -- * Socket + SocketError (..), + UdpSocket, + + -- * Peer Networking + NetPeer, + PeerId (..), + peerIdFromAddr, + PeerEvent (..), + ConnectionDirection (..), + DisconnectReason (..), + + -- ** Peer Creation + newPeer, + newPeerState, + + -- ** Peer Operations (Polymorphic) + peerTick, + peerRecvAllM, + peerSendAllM, + peerShutdownM, + + -- ** Peer Operations (Pure) + peerConnect, + peerDisconnect, + peerProcess, + peerSend, + peerBroadcast, + + -- ** Pure Processing Types + PeerResult (..), + IncomingPacket (..), + RawPacket (..), + + -- ** Peer Queries + peerCount, + peerIsConnected, + peerStats, + peerLocalAddr, + peerConnectedIds, + peerConfig, + + -- * Connection + ConnectionState (..), + ConnectionError (..), + ChannelError (..), + + -- * Configuration + NetworkConfig (..), + defaultNetworkConfig, + ConfigError (..), + validateConfig, + SimulationConfig (..), + defaultSimulationConfig, + ChannelConfig (..), + defaultChannelConfig, + unreliableConfig, + reliableOrderedConfig, + reliableSequencedConfig, + DeliveryMode (..), + + -- * Simulation + NetworkSimulator, + newNetworkSimulator, + simulatorProcessSend, + simulatorReceiveReady, + simulatorPendingCount, + simulatorConfig, + + -- * Testing + TestNet, + runTestNet, + TestNetState (..), + initialTestNetState, + TestWorld (..), + newTestWorld, + runPeerInWorld, + deliverPackets, + worldAdvanceTime, + + -- * Serialization (Storable-based, zero-copy) + serialize, + deserialize, + deriveStorable, + + -- * Statistics + NetworkStats (..), + CongestionLevel (..), + + -- * Replication: Delta Compression + NetworkDelta (..), + BaselineSeq, + DeltaTracker, + newDeltaTracker, + deltaEncode, + deltaOnAck, + BaselineManager, + newBaselineManager, + pushBaseline, + getBaseline, + deltaDecode, + + -- * Replication: Interest Management + InterestManager (..), + RadiusInterest, + newRadiusInterest, + GridInterest, + newGridInterest, + filterRelevant, + + -- * Replication: Priority Accumulator + PriorityAccumulator, + newPriorityAccumulator, + register, + unregister, + accumulate, + drainTop, + getPriority, + + -- * Replication: Snapshot Interpolation + Interpolatable (..), + SnapshotBuffer, + newSnapshotBuffer, + newSnapshotBufferWithConfig, + pushSnapshot, + sampleSnapshot, + snapshotReady, + ) +where + +import Data.Word (Word16, Word8) +import GBNet.Channel (ChannelConfig (..), ChannelError (..), DeliveryMode (..), defaultChannelConfig, reliableOrderedConfig, reliableSequencedConfig, unreliableConfig) +import GBNet.Class +import GBNet.Config (ConfigError (..), NetworkConfig (..), SimulationConfig (..), defaultNetworkConfig, defaultSimulationConfig, validateConfig) +import GBNet.Connection (ConnectionError (..), ConnectionState (..), DisconnectReason (..)) +import GBNet.Net +import GBNet.Net.IO (initNetState) +import GBNet.Peer +import GBNet.Replication.Delta (BaselineManager, BaselineSeq, DeltaTracker, NetworkDelta (..), deltaDecode, deltaEncode, deltaOnAck, getBaseline, newBaselineManager, newDeltaTracker, pushBaseline) +import GBNet.Replication.Interest (GridInterest, InterestManager (..), RadiusInterest, newGridInterest, newRadiusInterest) +import qualified GBNet.Replication.Interest as Interest +import GBNet.Replication.Interpolation (Interpolatable (..), SnapshotBuffer, newSnapshotBuffer, newSnapshotBufferWithConfig, pushSnapshot, sampleSnapshot, snapshotReady) +import GBNet.Replication.Priority (PriorityAccumulator, accumulate, drainTop, getPriority, newPriorityAccumulator, register, unregister) +import GBNet.Serialize (deserialize, serialize) +import GBNet.Serialize.TH (deriveStorable) +import GBNet.Simulator (NetworkSimulator, newNetworkSimulator, simulatorConfig, simulatorPendingCount, simulatorProcessSend, simulatorReceiveReady) +import GBNet.Socket (SocketError (..), UdpSocket) +import GBNet.Stats (CongestionLevel (..), NetworkStats (..)) +import GBNet.TestNet +import GBNet.Types +import Network.Socket (SockAddr (..)) +import qualified Network.Socket as NS + +-- | Create a localhost address on the given port. +-- +-- @ +-- let addr = localhost 7777 +-- @ +localhost :: Word16 -> SockAddr +localhost port = SockAddrInet (fromIntegral port) (NS.tupleToHostAddress (127, 0, 0, 1)) + +-- | Create an IPv4 address from tuple and port. +-- +-- @ +-- let addr = ipv4 (192, 168, 1, 100) 7777 +-- @ +ipv4 :: (Word8, Word8, Word8, Word8) -> Word16 -> SockAddr +ipv4 (a, b, c, d) port = + SockAddrInet (fromIntegral port) (NS.tupleToHostAddress (a, b, c, d)) + +-- | Bind to any interface on the given port. +-- +-- @ +-- let addr = anyAddr 7777 +-- @ +anyAddr :: Word16 -> SockAddr +anyAddr port = SockAddrInet (fromIntegral port) 0 + +-- | Filter entities to only those relevant to an observer. +-- +-- @ +-- let nearby = filterRelevant interest observerPos entities +-- @ +filterRelevant :: + (InterestManager a) => + a -> + Interest.Position -> + [(Interest.Position, b)] -> + [b] +filterRelevant mgr observerPos = + map snd . filter (\(entPos, _) -> Interest.relevant mgr entPos observerPos) + +-- | Get the network configuration from a peer. +peerConfig :: NetPeer -> NetworkConfig +peerConfig = npConfig diff --git a/src/GBNet/Channel.hs b/src/GBNet/Channel.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Channel.hs @@ -0,0 +1,491 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedLabels #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- Module : GBNet.Channel +-- Description : Channel-based message delivery with multiple reliability modes +-- +-- Implements LiteNetLib-style delivery modes: Unreliable, UnreliableSequenced, +-- ReliableUnordered, ReliableOrdered, and ReliableSequenced. Each channel +-- manages its own sequence numbers, buffering, and retransmission. +module GBNet.Channel + ( -- * Delivery modes + DeliveryMode (..), + isReliable, + isSequenced, + isOrdered, + + -- * Configuration + ChannelConfig (..), + defaultChannelConfig, + unreliableConfig, + reliableOrderedConfig, + reliableSequencedConfig, + + -- * Channel message + ChannelMessage (..), + + -- * Channel + Channel (..), + newChannel, + resetChannel, + channelIsReliable, + channelSend, + getOutgoingMessage, + getRetransmitMessages, + onMessageReceived, + acknowledgeMessage, + channelReceive, + takePendingAcks, + channelUpdate, + + -- * Errors + ChannelError (..), + ) +where + +import Control.DeepSeq (NFData (..), rwhnf) +import qualified Data.ByteString as BS +import Data.Foldable (toList) +import Data.Map.Strict (Map) +import qualified Data.Map.Strict as Map +import Data.Sequence (Seq, (|>)) +import qualified Data.Sequence as Seq +import Data.Word (Word64, Word8) +import GBNet.Reliability (MonoTime, elapsedMs) +import GBNet.Types (ChannelId (..), SequenceNum (..)) +import GBNet.Util (sequenceGreaterThan) +import Optics ((%~), (&), (.~)) +import Optics.TH (makeFieldLabelsNoPrefix) + +-- | Delivery mode for channel messages. +data DeliveryMode + = -- | Fire and forget, no guarantees + Unreliable + | -- | Unreliable but drops out-of-order + UnreliableSequenced + | -- | Guaranteed delivery, any order + ReliableUnordered + | -- | Guaranteed delivery, strict order + ReliableOrdered + | -- | Guaranteed delivery, drops out-of-order + ReliableSequenced + deriving (Eq, Show, Enum, Bounded) + +instance NFData DeliveryMode where rnf = rwhnf + +-- | Check if a delivery mode guarantees delivery. +isReliable :: DeliveryMode -> Bool +isReliable Unreliable = False +isReliable UnreliableSequenced = False +isReliable _ = True + +-- | Check if a delivery mode drops out-of-order messages. +isSequenced :: DeliveryMode -> Bool +isSequenced UnreliableSequenced = True +isSequenced ReliableSequenced = True +isSequenced _ = False + +-- | Check if a delivery mode delivers in strict order. +isOrdered :: DeliveryMode -> Bool +isOrdered ReliableOrdered = True +isOrdered _ = False + +-- | Channel configuration. +data ChannelConfig = ChannelConfig + { ccDeliveryMode :: !DeliveryMode, + ccMaxMessageSize :: !Int, + ccMessageBufferSize :: !Int, + ccBlockOnFull :: !Bool, + -- | Milliseconds + ccOrderedBufferTimeout :: !Double, + ccMaxOrderedBufferSize :: !Int, + ccMaxReliableRetries :: !Int, + ccPriority :: !Word8 + } + deriving (Eq, Show) + +instance NFData ChannelConfig where + rnf (ChannelConfig dm mms mbs bof obt mobs mrr p) = + rnf dm `seq` + rnf mms `seq` + rnf mbs `seq` + rnf bof `seq` + rnf obt `seq` + rnf mobs `seq` + rnf mrr `seq` + rnf p + +-- | Default channel configuration (ReliableOrdered). +defaultChannelConfig :: ChannelConfig +defaultChannelConfig = + ChannelConfig + { ccDeliveryMode = ReliableOrdered, + ccMaxMessageSize = 1024, + ccMessageBufferSize = 256, + ccBlockOnFull = False, + ccOrderedBufferTimeout = 5000.0, + ccMaxOrderedBufferSize = 64, + ccMaxReliableRetries = 10, + ccPriority = 0 + } + +-- | Unreliable channel configuration. +unreliableConfig :: ChannelConfig +unreliableConfig = + defaultChannelConfig + { ccDeliveryMode = Unreliable, + ccBlockOnFull = False + } + +-- | Reliable ordered channel configuration. +reliableOrderedConfig :: ChannelConfig +reliableOrderedConfig = + defaultChannelConfig + { ccDeliveryMode = ReliableOrdered + } + +-- | Reliable sequenced channel configuration. +reliableSequencedConfig :: ChannelConfig +reliableSequencedConfig = + defaultChannelConfig + { ccDeliveryMode = ReliableSequenced + } + +-- | A message in the channel system. +data ChannelMessage = ChannelMessage + { cmSequence :: !SequenceNum, + cmData :: !BS.ByteString, + cmSendTime :: !MonoTime, + cmAcked :: !Bool, + cmRetryCount :: !Int, + cmReliable :: !Bool + } + deriving (Show) + +instance NFData ChannelMessage where + rnf (ChannelMessage s d t a r rel) = + rnf s `seq` rnf d `seq` rnf t `seq` rnf a `seq` rnf r `seq` rnf rel + +-- | Channel state for message delivery. +data Channel = Channel + { chConfig :: !ChannelConfig, + chChannelId :: !ChannelId, + chLocalSequence :: !SequenceNum, + chRemoteSequence :: !SequenceNum, + chSendBuffer :: !(Map SequenceNum ChannelMessage), + chReceiveBuffer :: !(Seq BS.ByteString), -- O(1) append via |> + chPendingAck :: ![SequenceNum], + chOrderedReceiveBuffer :: !(Map SequenceNum (BS.ByteString, MonoTime)), + chOrderedExpected :: !SequenceNum, + chTotalSent :: !Word64, + chTotalReceived :: !Word64, + chTotalDropped :: !Word64, + chTotalRetransmits :: !Word64 + } + deriving (Show) + +instance NFData Channel where + rnf (Channel cfg ci ls rs sb rb pa orb oe ts tr td trt) = + rnf cfg `seq` + rnf ci `seq` + rnf ls `seq` + rnf rs `seq` + rnf sb `seq` + rnf rb `seq` + rnf pa `seq` + rnf orb `seq` + rnf oe `seq` + rnf ts `seq` + rnf tr `seq` + rnf td `seq` + rnf trt + +makeFieldLabelsNoPrefix ''ChannelConfig +makeFieldLabelsNoPrefix ''ChannelMessage +makeFieldLabelsNoPrefix ''Channel + +-- | Create a new channel with the given configuration. +newChannel :: ChannelId -> ChannelConfig -> Channel +newChannel channelId config = + Channel + { chConfig = config, + chChannelId = channelId, + chLocalSequence = 0, + chRemoteSequence = 0, + chSendBuffer = Map.empty, + chReceiveBuffer = Seq.empty, + chPendingAck = [], + chOrderedReceiveBuffer = Map.empty, + chOrderedExpected = 0, + chTotalSent = 0, + chTotalReceived = 0, + chTotalDropped = 0, + chTotalRetransmits = 0 + } + +-- | Queue a message for sending. +-- Returns 'Left ChannelBufferFull' if the buffer is full and blocking is enabled. +-- Returns 'Left ChannelMessageTooLarge' if the message exceeds the max size. +channelSend :: BS.ByteString -> MonoTime -> Channel -> Either ChannelError (SequenceNum, Channel) +channelSend payload now ch + | BS.length payload > ccMaxMessageSize (chConfig ch) = Left ChannelMessageTooLarge + | bufferFull && ccBlockOnFull (chConfig ch) = Left ChannelBufferFull + | otherwise = Right (seqNum, ch') + where + bufferFull = Map.size (chSendBuffer ch) >= ccMessageBufferSize (chConfig ch) + seqNum = chLocalSequence ch + reliable = isReliable (ccDeliveryMode (chConfig ch)) + msg = + ChannelMessage + { cmSequence = seqNum, + cmData = payload, + cmSendTime = now, + cmAcked = False, + cmRetryCount = 0, + cmReliable = reliable + } + sendBuf = + if bufferFull + then Map.deleteMin (chSendBuffer ch) + else chSendBuffer ch + ch' = + ch + & #chLocalSequence + .~ (seqNum + 1) + & #chSendBuffer + .~ Map.insert seqNum msg sendBuf + & #chTotalSent + %~ (+ 1) + +-- | Get the next outgoing message that hasn't been sent yet. +getOutgoingMessage :: Channel -> Maybe (ChannelMessage, Channel) +getOutgoingMessage ch = + case Map.lookupMin (chSendBuffer ch) of + Nothing -> Nothing + Just (seqNum, msg) + | cmAcked msg -> + -- Already acked, remove and try next + getOutgoingMessage (ch & #chSendBuffer %~ Map.delete seqNum) + | cmRetryCount msg == 0 -> + -- First send + if cmReliable msg + then + -- Reliable: keep in buffer for retransmit + let msg' = msg & #cmRetryCount .~ 1 + ch' = ch & #chSendBuffer %~ Map.insert seqNum msg' + in Just (msg, ch') + else + -- Unreliable: remove from buffer immediately (fire and forget) + let ch' = ch & #chSendBuffer %~ Map.delete seqNum + in Just (msg, ch') + | otherwise -> Nothing -- Already sent, waiting for ack or retransmit + +-- | Get messages that need retransmission based on RTO. +getRetransmitMessages :: MonoTime -> Double -> Channel -> ([ChannelMessage], Channel) +getRetransmitMessages now rtoMs ch + | not (isReliable (ccDeliveryMode (chConfig ch))) = ([], ch) + | otherwise = Map.foldrWithKey checkRetransmit ([], ch) (chSendBuffer ch) + where + maxRetries = ccMaxReliableRetries (chConfig ch) + + checkRetransmit seqNum msg (acc, c) + | cmAcked msg = (acc, c) + | cmRetryCount msg == 0 = (acc, c) -- Not sent yet + | cmRetryCount msg > maxRetries = + -- Give up on this message + ( acc, + c + & #chSendBuffer + %~ Map.delete seqNum + & #chTotalDropped + %~ (+ 1) + ) + | elapsedMs (cmSendTime msg) now >= rtoMs = + -- Needs retransmit + let msg' = msg & #cmSendTime .~ now & #cmRetryCount %~ (+ 1) + c' = + c + & #chSendBuffer + %~ Map.insert seqNum msg' + & #chTotalRetransmits + %~ (+ 1) + in (msg : acc, c') + | otherwise = (acc, c) + +-- | Process a received message. Returns updated channel. +onMessageReceived :: SequenceNum -> BS.ByteString -> MonoTime -> Channel -> Channel +onMessageReceived seqNum payload now ch = + case ccDeliveryMode (chConfig ch) of + Unreliable -> + ch + & #chReceiveBuffer + %~ (|> payload) + & #chTotalReceived + %~ (+ 1) + UnreliableSequenced -> + if sequenceGreaterThan seqNum (chRemoteSequence ch) + then + ch + & #chReceiveBuffer + %~ (|> payload) + & #chRemoteSequence + .~ seqNum + & #chTotalReceived + %~ (+ 1) + else ch & #chTotalDropped %~ (+ 1) + ReliableUnordered -> + ch + & #chReceiveBuffer + %~ (|> payload) + & #chPendingAck + %~ (seqNum :) + & #chTotalReceived + %~ (+ 1) + ReliableOrdered -> + let ch' = ch & #chPendingAck %~ (seqNum :) + in if seqNum == chOrderedExpected ch' + then deliverOrdered payload ch' + else bufferOrdered seqNum payload now ch' + ReliableSequenced -> + if sequenceGreaterThan seqNum (chRemoteSequence ch) + then + ch + & #chReceiveBuffer + %~ (|> payload) + & #chPendingAck + %~ (seqNum :) + & #chRemoteSequence + .~ seqNum + & #chTotalReceived + %~ (+ 1) + else + ch + & #chPendingAck + %~ (seqNum :) + & #chTotalDropped + %~ (+ 1) + +-- | Deliver message and flush any buffered consecutive messages. +deliverOrdered :: BS.ByteString -> Channel -> Channel +deliverOrdered payload ch = + let ch' = + ch + & #chReceiveBuffer + %~ (|> payload) + & #chOrderedExpected + %~ (+ 1) + & #chTotalReceived + %~ (+ 1) + in flushOrderedBuffer ch' + +-- | Buffer an out-of-order message for later delivery. +bufferOrdered :: SequenceNum -> BS.ByteString -> MonoTime -> Channel -> Channel +bufferOrdered seqNum payload now ch + | Map.size (chOrderedReceiveBuffer ch) >= ccMaxOrderedBufferSize (chConfig ch) = + ch & #chTotalDropped %~ (+ 1) + | otherwise = + ch & #chOrderedReceiveBuffer %~ Map.insert seqNum (payload, now) + +-- | Flush consecutive messages from the ordered buffer. +flushOrderedBuffer :: Channel -> Channel +flushOrderedBuffer ch = + case Map.lookup (chOrderedExpected ch) (chOrderedReceiveBuffer ch) of + Nothing -> ch + Just (payload, _) -> + let ch' = + ch + & #chOrderedReceiveBuffer + %~ Map.delete (chOrderedExpected ch) + & #chReceiveBuffer + %~ (|> payload) + & #chOrderedExpected + %~ (+ 1) + & #chTotalReceived + %~ (+ 1) + in flushOrderedBuffer ch' + +-- | Acknowledge a message by sequence number. +acknowledgeMessage :: SequenceNum -> Channel -> Channel +acknowledgeMessage seqNum ch = + ch & #chSendBuffer %~ Map.adjust (\msg -> msg & #cmAcked .~ True) seqNum + +-- | Take all received messages from the channel. +channelReceive :: Channel -> ([BS.ByteString], Channel) +channelReceive ch = (toList (chReceiveBuffer ch), ch & #chReceiveBuffer .~ Seq.empty) + +-- | Take pending ack sequence numbers. +takePendingAcks :: Channel -> ([SequenceNum], Channel) +takePendingAcks ch = (chPendingAck ch, ch & #chPendingAck .~ []) + +-- | Update channel state (flush old ordered messages, clean up acked). +channelUpdate :: MonoTime -> Channel -> Channel +channelUpdate now = flushTimedOutOrdered now . cleanupAcked + +-- | Remove acked messages from send buffer. +cleanupAcked :: Channel -> Channel +cleanupAcked ch = + ch & #chSendBuffer %~ Map.filter (not . cmAcked) + +-- | Flush ordered messages that have timed out waiting. +flushTimedOutOrdered :: MonoTime -> Channel -> Channel +flushTimedOutOrdered now ch + | not (isOrdered (ccDeliveryMode (chConfig ch))) = ch + | Map.null (chOrderedReceiveBuffer ch) = ch + | otherwise = + let timeout = ccOrderedBufferTimeout (chConfig ch) + (toFlush, toKeep) = Map.partition (timedOut timeout) (chOrderedReceiveBuffer ch) + in if Map.null toFlush + then ch + else + let payloads = map fst (Map.elems toFlush) + newExpected = case Map.lookupMax toFlush of + Nothing -> chOrderedExpected ch + Just (maxSeq, _) -> maxSeq + 1 + in ch + & #chOrderedReceiveBuffer + .~ toKeep + & #chReceiveBuffer + %~ (<> Seq.fromList payloads) + & #chOrderedExpected + .~ newExpected + & #chTotalReceived + %~ (+ fromIntegral (length payloads)) + where + timedOut timeout (_, arriveTime) = elapsedMs arriveTime now >= timeout + +-- | Reset channel to initial state. +resetChannel :: Channel -> Channel +resetChannel ch = + ch + { chLocalSequence = 0, + chRemoteSequence = 0, + chSendBuffer = Map.empty, + chReceiveBuffer = Seq.empty, + chPendingAck = [], + chOrderedReceiveBuffer = Map.empty, + chOrderedExpected = 0, + chTotalSent = 0, + chTotalReceived = 0, + chTotalDropped = 0, + chTotalRetransmits = 0 + } + +-- | Check if channel uses reliable delivery. +channelIsReliable :: Channel -> Bool +channelIsReliable ch = isReliable (ccDeliveryMode (chConfig ch)) +{-# INLINE channelIsReliable #-} + +-- | Channel errors. +data ChannelError + = ChannelBufferFull + | ChannelMessageTooLarge + deriving (Eq, Show) + +instance NFData ChannelError where rnf = rwhnf diff --git a/src/GBNet/Class.hs b/src/GBNet/Class.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Class.hs @@ -0,0 +1,84 @@ +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} + +-- | +-- Module : GBNet.Class +-- Description : Effect abstractions for networking +-- +-- Typeclasses that abstract over network IO, enabling: +-- - Pure testing with deterministic "networks" +-- - Swappable backends (real sockets, mock, simulation) +-- - Clean separation of pure logic from IO +module GBNet.Class + ( -- * Time + MonoTime (..), + MonadTime (..), + getMonoTimeIO, + + -- * Network IO + MonadNetwork (..), + NetError (..), + ) +where + +import Control.DeepSeq (NFData) +import Control.Monad.State.Strict (StateT (..)) +import Control.Monad.Trans.Class (lift) +import Data.ByteString (ByteString) +import Data.Word (Word64) +import GHC.Clock (getMonotonicTimeNSec) +import Network.Socket (SockAddr) + +-- | Monotonic time in nanoseconds. +-- Derives 'Num' because arithmetic on timestamps is pervasive. +newtype MonoTime = MonoTime {unMonoTime :: Word64} + deriving stock (Eq, Ord, Show) + deriving newtype (Bounded, Enum, NFData, Num) + +-- | Network errors. +data NetError + = NetSendFailed !String + | NetSocketClosed + | NetTimeout + deriving (Eq, Show) + +-- | Monad that can provide monotonic time. +class (Monad m) => MonadTime m where + -- | Get current monotonic time in nanoseconds. + getMonoTime :: m MonoTime + +-- | Monad that can perform network IO. +-- +-- This abstraction allows: +-- - Real IO with UDP sockets +-- - Pure testing with simulated networks +-- - Network condition simulation (latency, loss, reordering) +class (MonadTime m) => MonadNetwork m where + -- | Send raw bytes to an address. + -- Returns Left on failure, Right () on success. + netSend :: SockAddr -> ByteString -> m (Either NetError ()) + + -- | Receive bytes (non-blocking). + -- Returns Nothing if no data available, Just (data, sender) otherwise. + netRecv :: m (Maybe (ByteString, SockAddr)) + + -- | Close the network (cleanup). + netClose :: m () + +-- | Lift MonadTime through StateT. +instance (MonadTime m) => MonadTime (StateT s m) where + getMonoTime = lift getMonoTime + +-- | Lift MonadNetwork through StateT. +instance (MonadNetwork m) => MonadNetwork (StateT s m) where + netSend addr bs = lift (netSend addr bs) + netRecv = lift netRecv + netClose = lift netClose + +-- | Get current monotonic time in nanoseconds (IO helper). +getMonoTimeIO :: IO MonoTime +getMonoTimeIO = MonoTime <$> getMonotonicTimeNSec + +-- | MonadTime instance for IO. +instance MonadTime IO where + getMonoTime = getMonoTimeIO diff --git a/src/GBNet/Config.hs b/src/GBNet/Config.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Config.hs @@ -0,0 +1,369 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- Module : GBNet.Config +-- Description : Network configuration constants and structures +-- +-- 'NetworkConfig' controls all tunable parameters: timeouts, MTU, channels, +-- congestion, etc. 'ChannelConfig' is re-exported from "GBNet.Channel". +module GBNet.Config + ( -- * Constants + defaultProtocolId, + defaultMaxClients, + defaultConnectionTimeoutMs, + defaultKeepaliveIntervalMs, + defaultConnectionRequestTimeoutMs, + defaultConnectionRequestMaxRetries, + defaultMtu, + defaultFragmentThreshold, + defaultFragmentTimeoutMs, + defaultMaxFragments, + defaultMaxReassemblyBufferSize, + defaultPacketBufferSize, + defaultAckBufferSize, + defaultMaxSequenceDistance, + defaultReliableRetryTimeMs, + defaultMaxReliableRetries, + defaultMaxChannels, + defaultSendRateHz, + defaultMaxPacketRateHz, + defaultCongestionThreshold, + defaultCongestionGoodRttThresholdMs, + defaultCongestionBadLossThreshold, + defaultCongestionRecoveryTimeMs, + defaultDisconnectRetries, + defaultDisconnectRetryTimeoutMs, + defaultMaxInFlight, + defaultChannelPriority, + maxBackoffExponent, + minMtu, + maxMtu, + maxChannelCount, + defaultDeltaBaselineTimeoutMs, + defaultMaxBaselineSnapshots, + + -- * Configuration + NetworkConfig (..), + defaultNetworkConfig, + + -- * Configuration errors + ConfigError (..), + validateConfig, + + -- * Simulation + SimulationConfig (..), + defaultSimulationConfig, + ) +where + +import Data.Word (Word16, Word32, Word8) +import GBNet.Channel (ChannelConfig, defaultChannelConfig) +import Optics.TH (makeFieldLabelsNoPrefix) + +-- Constants + +-- | Default protocol identifier for packet validation. +defaultProtocolId :: Word32 +defaultProtocolId = 0x12345678 + +-- | Default maximum simultaneous client connections. +defaultMaxClients :: Int +defaultMaxClients = 64 + +-- | Default connection timeout in milliseconds. +defaultConnectionTimeoutMs :: Double +defaultConnectionTimeoutMs = 10000.0 + +-- | Default keepalive interval in milliseconds. +defaultKeepaliveIntervalMs :: Double +defaultKeepaliveIntervalMs = 1000.0 + +-- | Default connection request timeout in milliseconds. +defaultConnectionRequestTimeoutMs :: Double +defaultConnectionRequestTimeoutMs = 5000.0 + +-- | Default maximum connection request retries before giving up. +defaultConnectionRequestMaxRetries :: Int +defaultConnectionRequestMaxRetries = 5 + +-- | Default maximum transmission unit in bytes. +defaultMtu :: Int +defaultMtu = 1200 + +-- | Default fragment threshold in bytes; messages above this are fragmented. +defaultFragmentThreshold :: Int +defaultFragmentThreshold = 1024 + +-- | Default fragment reassembly timeout in milliseconds. +defaultFragmentTimeoutMs :: Double +defaultFragmentTimeoutMs = 5000.0 + +-- | Default maximum fragments per message. +defaultMaxFragments :: Int +defaultMaxFragments = 256 + +-- | Default maximum reassembly buffer size in bytes. +defaultMaxReassemblyBufferSize :: Int +defaultMaxReassemblyBufferSize = 1024 * 1024 + +-- | Default packet buffer size (number of entries). +defaultPacketBufferSize :: Int +defaultPacketBufferSize = 256 + +-- | Default ACK buffer size (number of entries). +defaultAckBufferSize :: Int +defaultAckBufferSize = 256 + +-- | Default maximum sequence distance before treating packets as stale. +defaultMaxSequenceDistance :: Word16 +defaultMaxSequenceDistance = 32768 + +-- | Default reliable retry time in milliseconds. +defaultReliableRetryTimeMs :: Double +defaultReliableRetryTimeMs = 100.0 + +-- | Default maximum reliable message retries before dropping. +defaultMaxReliableRetries :: Int +defaultMaxReliableRetries = 10 + +-- | Default maximum channel count. +defaultMaxChannels :: Int +defaultMaxChannels = 8 + +-- | Default send rate in packets per second. +defaultSendRateHz :: Double +defaultSendRateHz = 60.0 + +-- | Default maximum packet rate in packets per second. +defaultMaxPacketRateHz :: Double +defaultMaxPacketRateHz = 120.0 + +-- | Default congestion detection threshold (fraction of budget used). +defaultCongestionThreshold :: Double +defaultCongestionThreshold = 0.1 + +-- | Default RTT threshold in milliseconds for "good" congestion mode. +defaultCongestionGoodRttThresholdMs :: Double +defaultCongestionGoodRttThresholdMs = 250.0 + +-- | Default packet loss threshold for "bad" congestion mode. +defaultCongestionBadLossThreshold :: Double +defaultCongestionBadLossThreshold = 0.1 + +-- | Default congestion recovery time in milliseconds. +defaultCongestionRecoveryTimeMs :: Double +defaultCongestionRecoveryTimeMs = 10000.0 + +-- | Default number of disconnect packet retries. +defaultDisconnectRetries :: Int +defaultDisconnectRetries = 3 + +-- | Default disconnect retry timeout in milliseconds. +defaultDisconnectRetryTimeoutMs :: Double +defaultDisconnectRetryTimeoutMs = 500.0 + +-- | Default maximum packets in flight before back-pressure. +defaultMaxInFlight :: Int +defaultMaxInFlight = 256 + +-- | Default channel priority (0 = lowest, 255 = highest). +defaultChannelPriority :: Word8 +defaultChannelPriority = 128 + +-- | Maximum exponential backoff exponent (caps at 2^5 = 32x RTO). +maxBackoffExponent :: Int +maxBackoffExponent = 5 + +-- | Minimum allowed MTU (RFC 791 minimum IP datagram size). +minMtu :: Int +minMtu = 576 + +-- | Maximum allowed MTU. +maxMtu :: Int +maxMtu = 65535 + +-- | Maximum channel count, constrained by 3-bit wire format (payload header). +maxChannelCount :: Int +maxChannelCount = 8 + +-- | Default maximum pending connection requests. +defaultMaxPending :: Int +defaultMaxPending = 256 + +-- | Default rate limit for connection requests per second. +defaultRateLimitPerSecond :: Int +defaultRateLimitPerSecond = 10 + +-- | Default delta baseline timeout in milliseconds. +defaultDeltaBaselineTimeoutMs :: Double +defaultDeltaBaselineTimeoutMs = 2000.0 + +-- | Default maximum baseline snapshots for delta compression. +defaultMaxBaselineSnapshots :: Int +defaultMaxBaselineSnapshots = 32 + +-- | Configuration validation errors. +data ConfigError + = FragmentThresholdExceedsMtu + | InvalidChannelCount + | InvalidPacketBufferSize + | InvalidMtu + | TimeoutNotGreaterThanKeepalive + | InvalidMaxClients + | ChannelConfigsExceedMaxChannels + | InvalidSendRate + | InvalidMaxPacketRate + | InvalidMaxInFlight + | InvalidFragmentThreshold + | SendRateExceedsMaxPacketRate + | InvalidCongestionThreshold + deriving (Eq, Show) + +-- | Top-level network configuration. +data NetworkConfig = NetworkConfig + { ncProtocolId :: !Word32, + ncMaxClients :: !Int, + ncConnectionTimeoutMs :: !Double, + ncKeepaliveIntervalMs :: !Double, + ncConnectionRequestTimeoutMs :: !Double, + ncConnectionRequestMaxRetries :: !Int, + ncMtu :: !Int, + ncFragmentThreshold :: !Int, + ncFragmentTimeoutMs :: !Double, + ncMaxFragments :: !Int, + ncMaxReassemblyBufferSize :: !Int, + ncPacketBufferSize :: !Int, + ncAckBufferSize :: !Int, + ncMaxSequenceDistance :: !Word16, + ncReliableRetryTimeMs :: !Double, + ncMaxReliableRetries :: !Int, + ncMaxInFlight :: !Int, + ncMaxChannels :: !Int, + ncDefaultChannelConfig :: !ChannelConfig, + ncChannelConfigs :: ![ChannelConfig], + ncSendRate :: !Double, + ncMaxPacketRate :: !Double, + ncCongestionThreshold :: !Double, + ncCongestionGoodRttThreshold :: !Double, + ncCongestionBadLossThreshold :: !Double, + ncCongestionRecoveryTimeMs :: !Double, + ncDisconnectRetries :: !Int, + ncDisconnectRetryTimeoutMs :: !Double, + ncMaxPending :: !Int, + ncRateLimitPerSecond :: !Int, + ncUseCwndCongestion :: !Bool, + ncSimulation :: !(Maybe SimulationConfig), + ncEnableConnectionMigration :: !Bool, + ncDeltaBaselineTimeoutMs :: !Double, + ncMaxBaselineSnapshots :: !Int + } + deriving (Eq, Show) + +-- | Default network configuration. +defaultNetworkConfig :: NetworkConfig +defaultNetworkConfig = + NetworkConfig + { ncProtocolId = defaultProtocolId, + ncMaxClients = defaultMaxClients, + ncConnectionTimeoutMs = defaultConnectionTimeoutMs, + ncKeepaliveIntervalMs = defaultKeepaliveIntervalMs, + ncConnectionRequestTimeoutMs = defaultConnectionRequestTimeoutMs, + ncConnectionRequestMaxRetries = defaultConnectionRequestMaxRetries, + ncMtu = defaultMtu, + ncFragmentThreshold = defaultFragmentThreshold, + ncFragmentTimeoutMs = defaultFragmentTimeoutMs, + ncMaxFragments = defaultMaxFragments, + ncMaxReassemblyBufferSize = defaultMaxReassemblyBufferSize, + ncPacketBufferSize = defaultPacketBufferSize, + ncAckBufferSize = defaultAckBufferSize, + ncMaxSequenceDistance = defaultMaxSequenceDistance, + ncReliableRetryTimeMs = defaultReliableRetryTimeMs, + ncMaxReliableRetries = defaultMaxReliableRetries, + ncMaxInFlight = defaultMaxInFlight, + ncMaxChannels = defaultMaxChannels, + ncDefaultChannelConfig = defaultChannelConfig, + ncChannelConfigs = [], + ncSendRate = defaultSendRateHz, + ncMaxPacketRate = defaultMaxPacketRateHz, + ncCongestionThreshold = defaultCongestionThreshold, + ncCongestionGoodRttThreshold = defaultCongestionGoodRttThresholdMs, + ncCongestionBadLossThreshold = defaultCongestionBadLossThreshold, + ncCongestionRecoveryTimeMs = defaultCongestionRecoveryTimeMs, + ncDisconnectRetries = defaultDisconnectRetries, + ncDisconnectRetryTimeoutMs = defaultDisconnectRetryTimeoutMs, + ncMaxPending = defaultMaxPending, + ncRateLimitPerSecond = defaultRateLimitPerSecond, + ncUseCwndCongestion = False, + ncSimulation = Nothing, + ncEnableConnectionMigration = True, + ncDeltaBaselineTimeoutMs = defaultDeltaBaselineTimeoutMs, + ncMaxBaselineSnapshots = defaultMaxBaselineSnapshots + } + +-- | Validate configuration, returning an error if invalid. +validateConfig :: NetworkConfig -> Either ConfigError () +validateConfig cfg + | ncFragmentThreshold cfg > ncMtu cfg = + Left FragmentThresholdExceedsMtu + | ncMaxChannels cfg == 0 || ncMaxChannels cfg > maxChannelCount = + Left InvalidChannelCount + | ncPacketBufferSize cfg == 0 = + Left InvalidPacketBufferSize + | ncMtu cfg < minMtu || ncMtu cfg > maxMtu = + Left InvalidMtu + | ncConnectionTimeoutMs cfg <= ncKeepaliveIntervalMs cfg = + Left TimeoutNotGreaterThanKeepalive + | ncMaxClients cfg == 0 = + Left InvalidMaxClients + | length (ncChannelConfigs cfg) > ncMaxChannels cfg = + Left ChannelConfigsExceedMaxChannels + | not (isValidPositive (ncSendRate cfg)) = + Left InvalidSendRate + | not (isValidPositive (ncMaxPacketRate cfg)) = + Left InvalidMaxPacketRate + | ncMaxInFlight cfg == 0 = + Left InvalidMaxInFlight + | ncFragmentThreshold cfg == 0 = + Left InvalidFragmentThreshold + | ncSendRate cfg > ncMaxPacketRate cfg = + Left SendRateExceedsMaxPacketRate + | not (isFinite (ncCongestionGoodRttThreshold cfg)) + || not (isFinite (ncCongestionBadLossThreshold cfg)) + || not (isFinite (ncCongestionThreshold cfg)) = + Left InvalidCongestionThreshold + | otherwise = Right () + where + isValidPositive x = x > 0 && not (isNaN x) + isFinite x = not (isNaN x) && not (isInfinite x) + +-- | Network condition simulation configuration. +data SimulationConfig = SimulationConfig + { simPacketLoss :: !Double, + simLatencyMs :: !Int, + simJitterMs :: !Int, + simDuplicateChance :: !Double, + simOutOfOrderChance :: !Double, + simBandwidthLimitBytesPerSec :: !Int + } + deriving (Eq, Show) + +-- | Default simulation config (no simulation). +defaultSimulationConfig :: SimulationConfig +defaultSimulationConfig = + SimulationConfig + { simPacketLoss = 0.0, + simLatencyMs = 0, + simJitterMs = 0, + simDuplicateChance = 0.0, + simOutOfOrderChance = 0.0, + simBandwidthLimitBytesPerSec = 0 + } + +makeFieldLabelsNoPrefix ''NetworkConfig +makeFieldLabelsNoPrefix ''SimulationConfig diff --git a/src/GBNet/Congestion.hs b/src/GBNet/Congestion.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Congestion.hs @@ -0,0 +1,541 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedLabels #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- Module : GBNet.Congestion +-- Description : Binary congestion control and bandwidth tracking +-- +-- Gaffer-style Good/Bad mode congestion control, byte-budget gating, +-- adaptive recovery timer, message batching, and bandwidth tracking. +module GBNet.Congestion + ( -- * Constants + congestionRateReduction, + minSendRate, + batchHeaderSize, + batchLengthSize, + maxBatchMessages, + initialCwndPackets, + minCwndBytes, + minRecoverySecs, + maxRecoverySecs, + recoveryHalveIntervalSecs, + quickDropThresholdSecs, + initialSsthresh, + + -- * Congestion mode + CongestionMode (..), + CongestionPhase (..), + + -- * Binary congestion controller + CongestionController (..), + newCongestionController, + ccRefillBudget, + ccDeductBudget, + ccUpdate, + ccCanSend, + + -- * Congestion level query + ccCongestionLevel, + cwCongestionLevel, + + -- * Window-based congestion controller + CongestionWindow (..), + newCongestionWindow, + cwOnAck, + cwOnLoss, + cwOnSend, + cwCanSend, + cwUpdatePacing, + cwCanSendPaced, + cwSlowStartRestart, + + -- * Bandwidth tracking + BandwidthTracker (..), + newBandwidthTracker, + btRecord, + btBytesPerSecond, + + -- * Message batching + batchMessages, + unbatchMessages, + ) +where + +import Control.DeepSeq (NFData (..)) +import qualified Data.ByteString as BS +import qualified Data.ByteString.Builder as BSB +import qualified Data.ByteString.Lazy as BSL +import Data.Sequence (Seq) +import qualified Data.Sequence as Seq +import Data.Word (Word64, Word8) +import GBNet.Reliability (MonoTime, elapsedMs) +import GBNet.Stats (CongestionLevel (..)) +import Optics ((%~), (&), (.~), (?~)) +import Optics.TH (makeFieldLabelsNoPrefix) + +-- Constants + +-- | Multiplicative decrease factor on congestion (halve rate). +congestionRateReduction :: Double +congestionRateReduction = 0.5 + +-- | Minimum send rate floor in packets per second. +minSendRate :: Double +minSendRate = 1.0 + +-- | Batch header overhead: 1 byte for message count. +batchHeaderSize :: Int +batchHeaderSize = 1 + +-- | Per-message length prefix size in a batch: 2 bytes (big-endian u16). +batchLengthSize :: Int +batchLengthSize = 2 + +-- | Maximum messages per batch (limited by u8 count field). +maxBatchMessages :: Word8 +maxBatchMessages = 255 + +-- | Initial congestion window in packets (slow start). +initialCwndPackets :: Int +initialCwndPackets = 10 + +-- | Minimum congestion window in bytes (one MTU). +minCwndBytes :: Int +minCwndBytes = 1200 + +-- | Minimum adaptive recovery time in seconds. +minRecoverySecs :: Double +minRecoverySecs = 1.0 + +-- | Maximum adaptive recovery time in seconds. +maxRecoverySecs :: Double +maxRecoverySecs = 60.0 + +-- | Interval in seconds of sustained good conditions before halving recovery time. +recoveryHalveIntervalSecs :: Double +recoveryHalveIntervalSecs = 10.0 + +-- | If Good mode lasts less than this (seconds) before re-entering Bad, double recovery time. +quickDropThresholdSecs :: Double +quickDropThresholdSecs = 10.0 + +-- | Initial slow-start threshold (IEEE 754 positive infinity). +initialSsthresh :: Double +initialSsthresh = 1 / 0 + +-- | Additive increase per update tick when in Good mode (packets/sec). +sendRateIncrease :: Double +sendRateIncrease = 1.0 + +-- | Maximum send rate multiplier relative to base rate. +maxSendRateMultiplier :: Double +maxSendRateMultiplier = 4.0 + +-- | Binary congestion state. +data CongestionMode + = CongestionGood + | CongestionBad + deriving (Eq, Show) + +instance NFData CongestionMode where rnf x = x `seq` () + +-- | Phase for window-based congestion control. +data CongestionPhase + = SlowStart + | Avoidance + | Recovery + deriving (Eq, Show) + +instance NFData CongestionPhase where rnf x = x `seq` () + +-- | Binary congestion controller (Gaffer-style). +data CongestionController = CongestionController + { ccMode :: !CongestionMode, + ccGoodConditionsStart :: !(Maybe MonoTime), + ccLossThreshold :: !Double, + ccRttThresholdMs :: !Double, + ccBaseSendRate :: !Double, + ccCurrentSendRate :: !Double, + ccBudgetBytesRemaining :: !Int, + ccBytesPerTick :: !Int, + ccAdaptiveRecoverySecs :: !Double, + ccLastGoodEntry :: !(Maybe MonoTime), + ccLastBadEntry :: !(Maybe MonoTime) + } + deriving (Show) + +makeFieldLabelsNoPrefix ''CongestionController + +instance NFData CongestionController where + rnf (CongestionController m gs lt rt br cr bb bt ar lg lb) = + rnf m `seq` rnf gs `seq` rnf lt `seq` rnf rt `seq` rnf br `seq` + rnf cr `seq` rnf bb `seq` rnf bt `seq` rnf ar `seq` rnf lg `seq` rnf lb + +-- | Create a new congestion controller. +newCongestionController :: Double -> Double -> Double -> Double -> CongestionController +newCongestionController baseSendRate lossThreshold rttThresholdMs recoveryTimeMs = + CongestionController + { ccMode = CongestionGood, + ccGoodConditionsStart = Nothing, + ccLossThreshold = lossThreshold, + ccRttThresholdMs = rttThresholdMs, + ccBaseSendRate = baseSendRate, + ccCurrentSendRate = baseSendRate, + ccBudgetBytesRemaining = 0, + ccBytesPerTick = 0, + ccAdaptiveRecoverySecs = recoveryTimeMs / 1000.0, + ccLastGoodEntry = Nothing, + ccLastBadEntry = Nothing + } + +-- | Refill the byte budget at the start of each tick. +ccRefillBudget :: Int -> CongestionController -> CongestionController +ccRefillBudget mtu cc = + let bytesPerTick = floor (ccCurrentSendRate cc * fromIntegral mtu) + in cc + & #ccBytesPerTick + .~ bytesPerTick + & #ccBudgetBytesRemaining + .~ bytesPerTick + +-- | Deduct bytes from the send budget. +ccDeductBudget :: Int -> CongestionController -> CongestionController +ccDeductBudget bytes cc = + cc & #ccBudgetBytesRemaining %~ subtract bytes +{-# INLINE ccDeductBudget #-} + +-- | Update congestion state based on current network conditions. +ccUpdate :: Double -> Double -> MonoTime -> CongestionController -> CongestionController +ccUpdate packetLoss rttMs now cc = + let isBad = packetLoss > ccLossThreshold cc || rttMs > ccRttThresholdMs cc + in case ccMode cc of + CongestionGood + | isBad -> + -- Quick re-entry to Bad doubles recovery timer + let recoveryMult = case ccLastGoodEntry cc of + Just goodEntry -> + if elapsedMs goodEntry now < quickDropThresholdSecs * 1000.0 + then 2.0 + else 1.0 + Nothing -> 1.0 + newRecovery = min maxRecoverySecs (ccAdaptiveRecoverySecs cc * recoveryMult) + newRate = max minSendRate (ccCurrentSendRate cc * congestionRateReduction) + in cc + & #ccMode + .~ CongestionBad + & #ccLastBadEntry + ?~ now + & #ccCurrentSendRate + .~ newRate + & #ccGoodConditionsStart + .~ Nothing + & #ccAdaptiveRecoverySecs + .~ newRecovery + | otherwise -> + -- Additive increase: ramp rate up toward max capacity + let maxRate = ccBaseSendRate cc * maxSendRateMultiplier + newRate = min maxRate (ccCurrentSendRate cc + sendRateIncrease) + -- Halve recovery time after sustained good conditions + cc' = cc & #ccCurrentSendRate .~ newRate + in case ccLastGoodEntry cc' of + Just goodEntry -> + let elapsed = elapsedMs goodEntry now / 1000.0 + intervals = floor (elapsed / recoveryHalveIntervalSecs) :: Int + in if intervals > 0 + then + let newRecovery = + max + minRecoverySecs + (ccAdaptiveRecoverySecs cc' / (2.0 ^ intervals)) + in cc' + & #ccAdaptiveRecoverySecs + .~ newRecovery + & #ccLastGoodEntry + ?~ now + else cc' + Nothing -> cc' + CongestionBad + | not isBad -> + case ccGoodConditionsStart cc of + Nothing -> + cc & #ccGoodConditionsStart ?~ now + Just start -> + let requiredMs = ccAdaptiveRecoverySecs cc * 1000.0 + in if elapsedMs start now >= requiredMs + then + cc + & #ccMode + .~ CongestionGood + & #ccLastGoodEntry + ?~ now + & #ccCurrentSendRate + .~ ccBaseSendRate cc + & #ccGoodConditionsStart + .~ Nothing + else cc + | otherwise -> + cc & #ccGoodConditionsStart .~ Nothing + +-- | Check if a packet can be sent given packets sent and size. +ccCanSend :: Int -> Int -> CongestionController -> Bool +ccCanSend packetsSentThisCycle packetBytes cc = + fromIntegral packetsSentThisCycle < ccCurrentSendRate cc + && ccBudgetBytesRemaining cc >= packetBytes +{-# INLINE ccCanSend #-} + +-- | Query congestion level from the binary controller. +ccCongestionLevel :: CongestionController -> CongestionLevel +ccCongestionLevel cc = case ccMode cc of + CongestionBad -> CongestionCritical + CongestionGood + | budgetRatio < budgetElevatedThreshold -> CongestionElevated + | otherwise -> CongestionNone + where + budgetRatio + | ccBytesPerTick cc <= 0 = 1.0 + | otherwise = fromIntegral (ccBudgetBytesRemaining cc) / fromIntegral (ccBytesPerTick cc) :: Double + budgetElevatedThreshold = 0.25 + +-- | Query congestion level from the window-based controller. +cwCongestionLevel :: CongestionWindow -> CongestionLevel +cwCongestionLevel cw + | utilization > windowCriticalThreshold = CongestionCritical + | utilization > windowHighThreshold = CongestionHigh + | utilization > windowElevatedThreshold = CongestionElevated + | otherwise = CongestionNone + where + utilization + | cwCwnd cw <= 0 = 1.0 + | otherwise = fromIntegral (cwBytesInFlight cw) / cwCwnd cw + windowElevatedThreshold = 0.7 + windowHighThreshold = 0.85 + windowCriticalThreshold = 0.95 + +-- | Window-based congestion controller (TCP-like). +data CongestionWindow = CongestionWindow + { cwPhase :: !CongestionPhase, + cwCwnd :: !Double, + cwSsthresh :: !Double, + cwBytesInFlight :: !Word64, + cwMtu :: !Int, + cwLastSendTime :: !(Maybe MonoTime), + cwMinInterPacketDelay :: !Double -- milliseconds + } + deriving (Show) + +makeFieldLabelsNoPrefix ''CongestionWindow + +instance NFData CongestionWindow where + rnf (CongestionWindow p c s b m l d) = + rnf p `seq` rnf c `seq` rnf s `seq` rnf b `seq` rnf m `seq` rnf l `seq` rnf d + +-- | Create a new congestion window. +newCongestionWindow :: Int -> CongestionWindow +newCongestionWindow mtu = + CongestionWindow + { cwPhase = SlowStart, + cwCwnd = fromIntegral (initialCwndPackets * mtu), + cwSsthresh = initialSsthresh, + cwBytesInFlight = 0, + cwMtu = mtu, + cwLastSendTime = Nothing, + cwMinInterPacketDelay = 0.0 + } + +-- | Called when bytes are acknowledged. +cwOnAck :: Int -> CongestionWindow -> CongestionWindow +cwOnAck bytes cw = + let cw' = cw & #cwBytesInFlight %~ (\b -> b - fromIntegral (min bytes (fromIntegral b))) + in case cwPhase cw' of + SlowStart -> + let newCwnd = cwCwnd cw' + fromIntegral bytes + in if newCwnd >= cwSsthresh cw' + then cw' & #cwCwnd .~ newCwnd & #cwPhase .~ Avoidance + else cw' & #cwCwnd .~ newCwnd + Avoidance + | cwCwnd cw' > 0 -> + -- Additive increase: cwnd += mtu * bytes / cwnd + let increase = fromIntegral (cwMtu cw') * fromIntegral bytes / cwCwnd cw' + in cw' & #cwCwnd %~ (+ increase) + | otherwise -> cw' + Recovery -> + cw' -- Conservative in recovery + +-- | Called on packet loss detection. +cwOnLoss :: CongestionWindow -> CongestionWindow +cwOnLoss cw = + let newSsthresh = max (fromIntegral minCwndBytes) (cwCwnd cw / 2.0) + in cw + & #cwSsthresh + .~ newSsthresh + & #cwCwnd + .~ newSsthresh + & #cwPhase + .~ Recovery + +-- | Record bytes sent. +cwOnSend :: Int -> MonoTime -> CongestionWindow -> CongestionWindow +cwOnSend bytes now cw = + cw + & #cwBytesInFlight + %~ (+ fromIntegral bytes) + & #cwLastSendTime + ?~ now + +-- | Check if a packet can be sent. +cwCanSend :: Int -> CongestionWindow -> Bool +cwCanSend packetBytes cw = + cwBytesInFlight cw + fromIntegral packetBytes <= floor (cwCwnd cw) +{-# INLINE cwCanSend #-} + +-- | Update pacing delay from cwnd and RTT. +cwUpdatePacing :: Double -> CongestionWindow -> CongestionWindow +cwUpdatePacing rttMs cw + | cwCwnd cw > 0 && rttMs > 0 = + let packetsInWindow = cwCwnd cw / fromIntegral (cwMtu cw) + delay = + if packetsInWindow > 0 + then rttMs / packetsInWindow + else 0.0 + in cw & #cwMinInterPacketDelay .~ delay + | otherwise = cw + +-- | Check if enough time has elapsed for pacing. +cwCanSendPaced :: MonoTime -> CongestionWindow -> Bool +cwCanSendPaced now cw = + case cwLastSendTime cw of + Nothing -> True + Just lastSend -> elapsedMs lastSend now >= cwMinInterPacketDelay cw +{-# INLINE cwCanSendPaced #-} + +-- | Reset to slow start if idle too long (RFC 2861). +-- Prevents bursting a stale window after an idle period. +cwSlowStartRestart :: Double -> MonoTime -> CongestionWindow -> CongestionWindow +cwSlowStartRestart rtoMs now cw = + case cwLastSendTime cw of + Nothing -> cw + Just lastSend + | elapsedMs lastSend now > ssrIdleThreshold * rtoMs -> + let initialCwnd = fromIntegral (initialCwndPackets * cwMtu cw) + in cw + & #cwPhase + .~ SlowStart + & #cwCwnd + .~ initialCwnd + & #cwSsthresh + .~ cwCwnd cw + | otherwise -> cw + where + -- Idle for more than 2 RTOs triggers restart + ssrIdleThreshold = 2.0 + +-- | Bandwidth tracker using a sliding window with cached byte total. +data BandwidthTracker = BandwidthTracker + { btWindow :: !(Seq (MonoTime, Int)), + btWindowDurationMs :: !Double, + -- | Cached running total of bytes in window + btTotalBytes :: !Int + } + deriving (Show) + +makeFieldLabelsNoPrefix ''BandwidthTracker + +-- | Create a new bandwidth tracker. +newBandwidthTracker :: Double -> BandwidthTracker +newBandwidthTracker windowDurationMs = + BandwidthTracker + { btWindow = Seq.empty, + btWindowDurationMs = windowDurationMs, + btTotalBytes = 0 + } + +-- | Record bytes at the given time. +btRecord :: Int -> MonoTime -> BandwidthTracker -> BandwidthTracker +btRecord bytes now bt = + let bt' = bt & #btWindow %~ (Seq.|> (now, bytes)) & #btTotalBytes %~ (+ bytes) + in btCleanup now bt' + +-- | Get bytes per second. +btBytesPerSecond :: BandwidthTracker -> Double +btBytesPerSecond bt + | Seq.null (btWindow bt) = 0.0 + | otherwise = + let elapsedSecs = btWindowDurationMs bt / 1000.0 + in if elapsedSecs > 0 + then fromIntegral (btTotalBytes bt) / elapsedSecs + else 0.0 + +-- | Clean up old entries, maintaining the cached byte total. +btCleanup :: MonoTime -> BandwidthTracker -> BandwidthTracker +btCleanup now bt = + let cutoffMs = btWindowDurationMs bt + isStale (t, _) = elapsedMs t now >= cutoffMs + -- Drop stale entries from the front, subtracting their bytes + (stale, recent) = Seq.spanl isStale (btWindow bt) + evictedBytes = sum $ snd <$> stale + in bt & #btWindow .~ recent & #btTotalBytes %~ subtract evictedBytes + +-- | Pack multiple messages into batched packets. +-- Wire format: [u8 count][u16 len][data]... +batchMessages :: [BS.ByteString] -> Int -> [BS.ByteString] +batchMessages messages maxSize = reverse $ go messages [] [] batchHeaderSize 0 + where + -- Both accumulators are built in reverse (cons) for O(1) prepend. + go [] currentBatch batchesAcc _ msgCount + | msgCount > 0 = finalizeBatch msgCount currentBatch : batchesAcc + | otherwise = batchesAcc + go (msg : rest) currentBatch batchesAcc currentSize msgCount = + let msgWireSize = batchLengthSize + BS.length msg + shouldFinalize = + (currentSize + msgWireSize > maxSize && msgCount > 0) + || msgCount >= fromIntegral maxBatchMessages + in if shouldFinalize + then + go + (msg : rest) + [] + (finalizeBatch msgCount currentBatch : batchesAcc) + batchHeaderSize + 0 + else + go + rest + (msg : currentBatch) + batchesAcc + (currentSize + msgWireSize) + (msgCount + 1) + + finalizeBatch :: Int -> [BS.ByteString] -> BS.ByteString + finalizeBatch count msgsRev = + let msgs = reverse msgsRev + countByte = BSB.word8 (fromIntegral count) + msgBuilders = map (\m -> BSB.word16BE (fromIntegral (BS.length m)) <> BSB.byteString m) msgs + in BSL.toStrict $ BSB.toLazyByteString $ countByte <> mconcat msgBuilders + +-- | Unbatch a batched packet into individual messages. +unbatchMessages :: BS.ByteString -> Maybe [BS.ByteString] +unbatchMessages dat + | BS.null dat = Nothing + | otherwise = + let msgCount = fromIntegral (BS.index dat 0) :: Int + in readMessages msgCount 1 [] + where + readMessages 0 _ acc = Just (reverse acc) + readMessages n offset acc + | offset + 2 > BS.length dat = Nothing + | otherwise = + let len = + fromIntegral (BS.index dat offset) * 256 + + fromIntegral (BS.index dat (offset + 1)) + msgStart = offset + 2 + in if msgStart + len > BS.length dat + then Nothing + else + let msg = BS.take len (BS.drop msgStart dat) + in readMessages (n - 1) (msgStart + len) (msg : acc) diff --git a/src/GBNet/Connection.hs b/src/GBNet/Connection.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Connection.hs @@ -0,0 +1,759 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedLabels #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- Module : GBNet.Connection +-- Description : Connection state machine for reliable UDP +-- +-- Manages handshake, channels, reliability tracking, and congestion control. +module GBNet.Connection + ( -- * Connection state + ConnectionState (..), + DisconnectReason (..), + disconnectReasonCode, + parseDisconnectReason, + + -- * Errors + ConnectionError (..), + + -- * Outgoing packet + OutgoingPacket (..), + + -- * Connection (opaque) + Connection, + newConnection, + + -- * State queries + connectionState, + isConnected, + connectionStats, + connRemoteSeq, + connLocalSeq, + connClientSalt, + + -- * Operations + connect, + disconnect, + sendMessage, + receiveMessage, + receiveIncomingPayload, + + -- * Tick update + updateTick, + + -- * Packet handling + drainSendQueue, + createHeader, + processIncomingHeader, + + -- * Time tracking + touchRecvTime, + touchSendTime, + recordBytesSent, + recordBytesReceived, + + -- * State transitions + markConnected, + + -- * Reset + resetConnection, + + -- * Channel info + channelCount, + ) +where + +import Control.DeepSeq (NFData (..), rwhnf) +import Data.Bits (shiftR, (.&.)) +import qualified Data.ByteString as BS +import Data.Foldable (toList) +import Data.IntMap.Strict (IntMap) +import qualified Data.IntMap.Strict as IntMap +import Data.List (foldl', sortBy) +import Data.Ord (Down (..), comparing) +import Data.Sequence (Seq) +import qualified Data.Sequence as Seq +import Data.Word (Word32, Word64, Word8) +import GBNet.Channel (Channel, ChannelError (..), ChannelMessage (..)) +import qualified GBNet.Channel as Channel +import GBNet.Config (NetworkConfig (..)) +import GBNet.Congestion + ( BandwidthTracker, + CongestionController, + CongestionWindow, + btBytesPerSecond, + btRecord, + ccCanSend, + ccCongestionLevel, + ccDeductBudget, + ccRefillBudget, + ccUpdate, + cwCanSend, + cwCanSendPaced, + cwCongestionLevel, + cwOnAck, + cwOnLoss, + cwOnSend, + cwSlowStartRestart, + cwUpdatePacing, + newBandwidthTracker, + newCongestionController, + newCongestionWindow, + ) +import GBNet.Packet (PacketHeader (..), PacketType (..)) +import GBNet.Reliability (MonoTime, ReliableEndpoint, elapsedMs) +import qualified GBNet.Reliability as Rel +import GBNet.Stats + ( CongestionLevel (..), + NetworkStats (..), + assessConnectionQuality, + defaultNetworkStats, + ) +import GBNet.Types (ChannelId (..), SequenceNum (..), channelIdToInt) +import Optics ((%), (%~), (&), (.~), (?~)) +import Optics.TH (makeFieldLabelsNoPrefix) + +-- | Bandwidth tracking window duration in milliseconds. +bandwidthWindowMs :: Double +bandwidthWindowMs = 1000.0 + +-- | States of the connection state machine. +data ConnectionState + = Disconnected + | Connecting + | Connected + | Disconnecting + deriving (Eq, Show) + +-- | Typed disconnect reason decoded from wire code. +data DisconnectReason + = ReasonTimeout + | ReasonRequested + | ReasonKicked + | ReasonServerFull + | ReasonProtocolMismatch + | ReasonUnknown Word8 + deriving (Eq, Show) + +-- | Disconnect reason wire codes. +disconnectReasonCode :: DisconnectReason -> Word8 +disconnectReasonCode ReasonTimeout = 0 +disconnectReasonCode ReasonRequested = 1 +disconnectReasonCode ReasonKicked = 2 +disconnectReasonCode ReasonServerFull = 3 +disconnectReasonCode ReasonProtocolMismatch = 4 +disconnectReasonCode (ReasonUnknown code) = code + +-- | Parse disconnect reason from wire code. +parseDisconnectReason :: Word8 -> DisconnectReason +parseDisconnectReason 0 = ReasonTimeout +parseDisconnectReason 1 = ReasonRequested +parseDisconnectReason 2 = ReasonKicked +parseDisconnectReason 3 = ReasonServerFull +parseDisconnectReason 4 = ReasonProtocolMismatch +parseDisconnectReason code = ReasonUnknown code + +-- | Errors that can occur during connection operations. +data ConnectionError + = ErrNotConnected + | ErrAlreadyConnected + | ErrConnectionDenied Word8 + | ErrTimeout + | ErrProtocolMismatch + | ErrInvalidPacket + | ErrInvalidChannel ChannelId + | ErrChannelError ChannelError + | ErrMessageTooLarge + deriving (Eq, Show) + +instance NFData ConnectionError where rnf = rwhnf + +-- | Packet waiting to be sent. +data OutgoingPacket = OutgoingPacket + { opHeader :: !PacketHeader, + opType :: !PacketType, + opPayload :: !BS.ByteString + } + deriving (Show) + +instance NFData OutgoingPacket where rnf = rwhnf + +makeFieldLabelsNoPrefix ''OutgoingPacket + +-- | A single peer connection managing handshake, channels, reliability, and congestion. +data Connection = Connection + { connConfig :: !NetworkConfig, + connState :: !ConnectionState, + -- | Client-side salt for connection validation. + connClientSalt :: !Word64, + connServerSalt :: !Word64, + -- Timing + connLastSendTime :: !MonoTime, + connLastRecvTime :: !MonoTime, + connStartTime :: !(Maybe MonoTime), + connRequestTime :: !(Maybe MonoTime), + connRetryCount :: !Int, + -- | Local outgoing sequence number, incremented per packet. + connLocalSeq :: !SequenceNum, + -- Subsystems + connReliability :: !ReliableEndpoint, + connChannels :: !(IntMap Channel), + connChannelPriority :: ![Int], + connCongestion :: !CongestionController, + connCwnd :: !(Maybe CongestionWindow), + connBandwidthUp :: !BandwidthTracker, + connBandwidthDown :: !BandwidthTracker, + -- Queues + connSendQueue :: !(Seq OutgoingPacket), + -- Stats + connStats :: !NetworkStats, + -- Disconnect tracking + connDisconnectTime :: !(Maybe MonoTime), + connDisconnectRetries :: !Int, + -- Flags + connPendingAck :: !Bool, + connDataSentThisTick :: !Bool + } + deriving (Show) + +instance NFData Connection where rnf = rwhnf + +makeFieldLabelsNoPrefix ''Connection + +-- | Modify a single channel by index, leaving others unchanged. +modifyChannel :: Int -> (Channel -> Channel) -> Connection -> Connection +modifyChannel idx f conn = conn & #connChannels %~ IntMap.adjust f idx +{-# INLINE modifyChannel #-} + +-- | Create a new connection. +newConnection :: NetworkConfig -> Word64 -> MonoTime -> Connection +newConnection config clientSalt now = + let numChannels = ncMaxChannels config + defaultCfg = ncDefaultChannelConfig config + -- Pad custom configs with defaults, then take exactly numChannels + configs = take numChannels $ ncChannelConfigs config ++ repeat defaultCfg + channels = IntMap.fromList $ zip [0 ..] $ zipWith Channel.newChannel (map ChannelId [0 ..]) configs + -- Sort channel indices by priority (highest first) + -- Zip indices with configs, sort by priority, extract indices + priorityOrder = + map fst $ sortBy (comparing (Down . Channel.ccPriority . snd)) $ zip [0 ..] configs + congestion = + newCongestionController + (ncSendRate config) + (ncCongestionBadLossThreshold config) + (ncCongestionGoodRttThreshold config) + (ncCongestionRecoveryTimeMs config) + cwnd = + if ncUseCwndCongestion config + then Just (newCongestionWindow (ncMtu config)) + else Nothing + in Connection + { connConfig = config, + connState = Disconnected, + connClientSalt = clientSalt, + connServerSalt = 0, + connLastSendTime = now, + connLastRecvTime = now, + connStartTime = Nothing, + connRequestTime = Nothing, + connRetryCount = 0, + connLocalSeq = 0, + connReliability = Rel.newReliableEndpoint (ncPacketBufferSize config), + connChannels = channels, + connChannelPriority = priorityOrder, + connCongestion = congestion, + connCwnd = cwnd, + connBandwidthUp = newBandwidthTracker bandwidthWindowMs, + connBandwidthDown = newBandwidthTracker bandwidthWindowMs, + connSendQueue = Seq.empty, + connStats = defaultNetworkStats, + connDisconnectTime = Nothing, + connDisconnectRetries = 0, + connPendingAck = False, + connDataSentThisTick = False + } + +-- | Get connection state. +connectionState :: Connection -> ConnectionState +connectionState = connState +{-# INLINE connectionState #-} + +-- | Check if connected. +isConnected :: Connection -> Bool +isConnected conn = connState conn == Connected +{-# INLINE isConnected #-} + +-- | Get connection stats. +connectionStats :: Connection -> NetworkStats +connectionStats = connStats + +-- | Get remote sequence number (delegates to reliability layer). +connRemoteSeq :: Connection -> SequenceNum +connRemoteSeq = Rel.reRemoteSequence . connReliability +{-# INLINE connRemoteSeq #-} + +-- | Get number of channels. +channelCount :: Connection -> Word8 +channelCount = fromIntegral . IntMap.size . connChannels + +-- | Initiate connection. +connect :: MonoTime -> Connection -> Either ConnectionError Connection +connect now conn + | connState conn /= Disconnected = Left ErrAlreadyConnected + | otherwise = + Right $ + sendConnectionRequest $ + conn + & #connState + .~ Connecting + & #connRequestTime + ?~ now + & #connRetryCount + .~ 0 + +-- | Initiate disconnect. +disconnect :: DisconnectReason -> MonoTime -> Connection -> Connection +disconnect reason now conn + | connState conn == Disconnected = conn + | otherwise = + let header = createHeaderInternal conn + pkt = + OutgoingPacket + { opHeader = header, + opType = Disconnect, + opPayload = BS.singleton (disconnectReasonCode reason) + } + in conn + & #connState + .~ Disconnecting + & #connDisconnectTime + ?~ now + & #connDisconnectRetries + .~ 0 + & #connSendQueue + %~ (Seq.|> pkt) + & #connLocalSeq + %~ (+ 1) + +-- | Send a message on a channel. +sendMessage :: ChannelId -> BS.ByteString -> MonoTime -> Connection -> Either ConnectionError Connection +sendMessage channelId payload now conn + | connState conn /= Connected = Left ErrNotConnected + | otherwise = + case IntMap.lookup idx (connChannels conn) of + Nothing -> Left (ErrInvalidChannel channelId) + Just channel -> + case Channel.channelSend payload now channel of + Left chErr -> Left (ErrChannelError chErr) + Right (_msgSeq, channel') -> + Right $ modifyChannel idx (const channel') conn + where + idx = channelIdToInt channelId + +-- | Receive all pending messages from a channel. +receiveMessage :: ChannelId -> Connection -> ([BS.ByteString], Connection) +receiveMessage channelId conn = + case IntMap.lookup idx (connChannels conn) of + Nothing -> ([], conn) + Just channel -> + let (msgs, channel') = Channel.channelReceive channel + in (msgs, modifyChannel idx (const channel') conn) + where + idx = channelIdToInt channelId + +-- | Route an incoming payload through the appropriate channel. +receiveIncomingPayload :: ChannelId -> SequenceNum -> BS.ByteString -> MonoTime -> Connection -> Connection +receiveIncomingPayload channelId chSeq payload now conn = + case IntMap.lookup idx (connChannels conn) of + Nothing -> conn + Just _ -> modifyChannel idx (Channel.onMessageReceived chSeq payload now) conn + where + idx = channelIdToInt channelId + +-- | Create a packet header (increments local sequence). +createHeader :: Connection -> (PacketHeader, Connection) +createHeader conn = + let header = createHeaderInternal conn + in (header, conn & #connLocalSeq %~ (+ 1)) + +-- | Internal header creation without state update. +createHeaderInternal :: Connection -> PacketHeader +createHeaderInternal conn = + let (ackSeq, ackBits64) = Rel.getAckInfo (connReliability conn) + ackBits32 = fromIntegral ackBits64 :: Word32 -- Truncate to 32-bit wire format + in PacketHeader + { packetType = Payload, -- Will be overwritten + sequenceNum = connLocalSeq conn, + ack = ackSeq, + ackBitfield = ackBits32 + } + +-- | Send connection request packet. +sendConnectionRequest :: Connection -> Connection +sendConnectionRequest conn = + let header = + PacketHeader + { packetType = ConnectionRequest, + sequenceNum = 0, + ack = 0, + ackBitfield = 0 + } + pkt = + OutgoingPacket + { opHeader = header, + opType = ConnectionRequest, + opPayload = BS.empty + } + in conn & #connSendQueue %~ (Seq.|> pkt) + +-- | Process incoming packet header for reliability. +processIncomingHeader :: PacketHeader -> MonoTime -> Connection -> Connection +processIncomingHeader header now conn0 = + let -- Record received packet for ACK generation + conn1 = + conn0 + & #connReliability + %~ Rel.onPacketsReceived [sequenceNum header] + & #connPendingAck + .~ True + -- Process ACKs (extend 32-bit wire format to 64-bit) + ackBits64 = fromIntegral (ackBitfield header) :: Word64 + (ackResult, rel') = Rel.processAcks (ack header) ackBits64 now (connReliability conn1) + ackedPairs = Rel.arAcked ackResult + fastRetransmits = Rel.arFastRetransmit ackResult + conn2 = conn1 & #connReliability .~ rel' + -- Feed ack/loss info to cwnd + conn3 = case connCwnd conn2 of + Just cwVal -> + let hasLoss = not (null fastRetransmits) + ackedBytes = length ackedPairs * ncMtu (connConfig conn2) + cwVal' + | hasLoss = cwOnLoss cwVal + | ackedBytes > 0 = cwOnAck ackedBytes cwVal + | otherwise = cwVal + in conn2 & #connCwnd ?~ cwVal' + Nothing -> conn2 + -- Acknowledge messages on channels + conn4 = foldl' (\c (chId, chSeq) -> modifyChannel (channelIdToInt chId) (Channel.acknowledgeMessage chSeq) c) conn3 ackedPairs + in conn4 + +-- | Update connection state (called each tick). +updateTick :: MonoTime -> Connection -> Either ConnectionError Connection +updateTick now conn = case connState conn of + Disconnected -> Right conn + Connecting -> updateConnecting now conn + Connected -> updateConnected now conn + Disconnecting -> updateDisconnecting now conn + +-- | Update while connecting. +updateConnecting :: MonoTime -> Connection -> Either ConnectionError Connection +updateConnecting now conn = + case connRequestTime conn of + Nothing -> Right conn + Just reqTime + | elapsed <= timeoutMs -> Right conn + | retries > maxRetries -> Left ErrTimeout + | otherwise -> + Right $ + sendConnectionRequest $ + conn + & #connRetryCount + .~ retries + & #connRequestTime + ?~ now + where + elapsed = elapsedMs reqTime now + timeoutMs = ncConnectionRequestTimeoutMs (connConfig conn) + retries = connRetryCount conn + 1 + maxRetries = ncConnectionRequestMaxRetries (connConfig conn) + +-- | Update while connected. +updateConnected :: MonoTime -> Connection -> Either ConnectionError Connection +updateConnected now conn + | timeSinceRecv > timeoutMs = Left ErrTimeout + | otherwise = Right $ updateConnectedPure now conn + where + timeSinceRecv = elapsedMs (connLastRecvTime conn) now + timeoutMs = ncConnectionTimeoutMs (connConfig conn) + +-- | Pure connected update (after timeout check passes). +updateConnectedPure :: MonoTime -> Connection -> Connection +updateConnectedPure now conn0 = + let -- Update congestion + cfg = connConfig conn0 + cong' = ccRefillBudget (ncMtu cfg) $ ccUpdate (nsPacketLoss (connStats conn0)) (nsRtt (connStats conn0)) now (connCongestion conn0) + conn1 = conn0 & #connCongestion .~ cong' + -- Update cwnd pacing and check for slow start restart + conn2 = case connCwnd conn1 of + Just cwVal -> + let rto = Rel.rtoMs (connReliability conn1) + cwPaced = cwUpdatePacing rto $ cwSlowStartRestart rto now cwVal + in conn1 & #connCwnd ?~ cwPaced + Nothing -> conn1 + -- Send keepalive if needed + timeSinceSend = elapsedMs (connLastSendTime conn2) now + keepaliveMs = ncKeepaliveIntervalMs cfg + conn3 + | timeSinceSend > keepaliveMs = sendKeepalive conn2 + | otherwise = conn2 + -- Process channel outgoing messages + conn4 = processChannelOutput now conn3 + -- Update channels + conn5 = conn4 & #connChannels %~ IntMap.map (Channel.channelUpdate now) + -- Send AckOnly if needed + conn6 + | connPendingAck conn5 && not (connDataSentThisTick conn5) = sendAckOnly conn5 + | otherwise = conn5 + -- Compute congestion level (worst of binary and window-based) + binaryLevel = ccCongestionLevel (connCongestion conn6) + windowLevel = maybe CongestionNone cwCongestionLevel (connCwnd conn6) + congLevel = max binaryLevel windowLevel + -- Update stats + rel' = connReliability conn6 + conn7 = + conn6 + & #connStats + % #nsRtt + .~ Rel.srttMs rel' + & #connStats + % #nsPacketLoss + .~ Rel.packetLossPercent rel' + & #connStats + % #nsBandwidthUp + .~ btBytesPerSecond (connBandwidthUp conn6) + & #connStats + % #nsBandwidthDown + .~ btBytesPerSecond (connBandwidthDown conn6) + & #connStats + % #nsConnectionQuality + .~ assessConnectionQuality (Rel.srttMs rel') (Rel.packetLossPercent rel' * 100) + & #connStats + % #nsCongestionLevel + .~ congLevel + & #connPendingAck + .~ False + in conn7 + +-- | Process outgoing messages from channels. +processChannelOutput :: MonoTime -> Connection -> Connection +processChannelOutput now conn = + let conn' = conn & #connDataSentThisTick .~ False + in foldl' (processChannelIdx now) conn' (connChannelPriority conn') + +processChannelIdx :: MonoTime -> Connection -> Int -> Connection +processChannelIdx now conn chIdx = + if IntMap.member chIdx (connChannels conn) + then processChannelMessages now conn chIdx + else conn + +processChannelMessages :: MonoTime -> Connection -> Int -> Connection +processChannelMessages now conn chIdx = + let mtu = ncMtu (connConfig conn) + canSendCong = ccCanSend 0 mtu (connCongestion conn) + canSendCwnd = case connCwnd conn of + Just cw -> cwCanSend mtu cw && cwCanSendPaced now cw + Nothing -> True + in if not canSendCong || not canSendCwnd + then conn + else case IntMap.lookup chIdx (connChannels conn) of + Nothing -> conn + Just channel -> + case Channel.getOutgoingMessage channel of + Nothing -> conn + Just (msg, channel') -> + let msgData = cmData msg + msgSeqRaw = unSequenceNum (cmSequence msg) + -- Prepend payload header: channel (3 bits) + is_fragment (1 bit) + -- Format: [headerByte:1][channelSeqHi:1][channelSeqLo:1][msgData:N] + headerByte = fromIntegral chIdx .&. 0x07 -- Not a fragment + seqHi = fromIntegral (msgSeqRaw `shiftR` 8) :: Word8 + seqLo = fromIntegral (msgSeqRaw .&. 0xFF) :: Word8 + wireData = BS.cons headerByte $ BS.cons seqHi $ BS.cons seqLo msgData + header = createHeaderInternal conn + pkt = + OutgoingPacket + { opHeader = header {packetType = Payload}, + opType = Payload, + opPayload = wireData + } + cong' = ccDeductBudget (BS.length wireData) (connCongestion conn) + cwnd' = fmap (cwOnSend (BS.length wireData) now) (connCwnd conn) + rel' = + if Channel.channelIsReliable channel + then + Rel.onPacketSent + (connLocalSeq conn) + now + (ChannelId (fromIntegral chIdx)) + (cmSequence msg) + (BS.length wireData) + (connReliability conn) + else connReliability conn + conn' = + conn + & #connChannels + %~ IntMap.insert chIdx channel' + & #connSendQueue + %~ (Seq.|> pkt) + & #connLocalSeq + %~ (+ 1) + & #connCongestion + .~ cong' + & #connCwnd + .~ cwnd' + & #connReliability + .~ rel' + & #connDataSentThisTick + .~ True + in processChannelMessages now conn' chIdx + +-- | Enqueue an empty keepalive/ack-only packet. +-- Both keepalive and ack-only use the same wire format (Keepalive with empty payload). +sendKeepalive, sendAckOnly :: Connection -> Connection +sendKeepalive = enqueueEmptyPacket +sendAckOnly = enqueueEmptyPacket + +-- | Shared implementation for keepalive and ack-only packets. +enqueueEmptyPacket :: Connection -> Connection +enqueueEmptyPacket conn = + let header = createHeaderInternal conn + pkt = + OutgoingPacket + { opHeader = header {packetType = Keepalive}, + opType = Keepalive, + opPayload = BS.empty + } + in conn + & #connSendQueue + %~ (Seq.|> pkt) + & #connLocalSeq + %~ (+ 1) + +-- | Update while disconnecting. +updateDisconnecting :: MonoTime -> Connection -> Either ConnectionError Connection +updateDisconnecting now conn = + case connDisconnectTime conn of + Nothing -> Right conn + Just discTime + | elapsed <= timeoutMs -> Right conn + | retries >= maxRetries -> Right $ resetConnection $ conn & #connState .~ Disconnected + | otherwise -> + let header = createHeaderInternal conn + pkt = + OutgoingPacket + { opHeader = header, + opType = Disconnect, + opPayload = BS.singleton (disconnectReasonCode ReasonRequested) + } + in Right $ + conn + & #connDisconnectRetries + .~ (retries + 1) + & #connDisconnectTime + ?~ now + & #connSendQueue + %~ (Seq.|> pkt) + & #connLocalSeq + %~ (+ 1) + where + elapsed = elapsedMs discTime now + timeoutMs = ncDisconnectRetryTimeoutMs (connConfig conn) + retries = connDisconnectRetries conn + maxRetries = ncDisconnectRetries (connConfig conn) + +-- | Reset connection state. +resetConnection :: Connection -> Connection +resetConnection conn = + let config = connConfig conn + in conn + & #connStartTime + .~ Nothing + & #connRequestTime + .~ Nothing + & #connLocalSeq + .~ 0 + & #connSendQueue + .~ Seq.empty + & #connDisconnectTime + .~ Nothing + & #connDisconnectRetries + .~ 0 + & #connPendingAck + .~ False + & #connDataSentThisTick + .~ False + & #connChannels + .~ IntMap.map Channel.resetChannel (connChannels conn) + & #connCongestion + .~ newCongestionController + (ncSendRate config) + (ncCongestionBadLossThreshold config) + (ncCongestionGoodRttThreshold config) + (ncCongestionRecoveryTimeMs config) + & #connBandwidthUp + .~ newBandwidthTracker bandwidthWindowMs + & #connBandwidthDown + .~ newBandwidthTracker bandwidthWindowMs + +-- | Drain send queue. +drainSendQueue :: Connection -> ([OutgoingPacket], Connection) +drainSendQueue conn = + (toList (connSendQueue conn), conn & #connSendQueue .~ Seq.empty) + +-- | Update last receive time. +touchRecvTime :: MonoTime -> Connection -> Connection +touchRecvTime now conn = conn & #connLastRecvTime .~ now +{-# INLINE touchRecvTime #-} + +-- | Update last send time. +touchSendTime :: MonoTime -> Connection -> Connection +touchSendTime now conn = conn & #connLastSendTime .~ now +{-# INLINE touchSendTime #-} + +-- | Mark connection as established (Connected state). +-- Used after handshake completes. +markConnected :: MonoTime -> Connection -> Connection +markConnected now conn = + conn + & #connState + .~ Connected + & #connStartTime + ?~ now + & #connLocalSeq + .~ 0 + +-- | Record bytes sent for bandwidth tracking. +recordBytesSent :: Int -> MonoTime -> Connection -> Connection +recordBytesSent bytes now conn = + let bw = btRecord bytes now (connBandwidthUp conn) + in conn + & #connBandwidthUp + .~ bw + & #connStats + % #nsPacketsSent + %~ (+ 1) + & #connStats + % #nsBytesSent + %~ (+ fromIntegral bytes) + & #connLastSendTime + .~ now + +-- | Record bytes received for bandwidth tracking. +recordBytesReceived :: Int -> MonoTime -> Connection -> Connection +recordBytesReceived bytes now conn = + let bw = btRecord bytes now (connBandwidthDown conn) + in conn + & #connBandwidthDown + .~ bw + & #connStats + % #nsPacketsReceived + %~ (+ 1) + & #connStats + % #nsBytesReceived + %~ (+ fromIntegral bytes) diff --git a/src/GBNet/Fragment.hs b/src/GBNet/Fragment.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Fragment.hs @@ -0,0 +1,419 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedLabels #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- Module : GBNet.Fragment +-- Description : Message fragmentation, reassembly, and MTU discovery +-- +-- Splits large messages into fragments that fit within MTU, reassembles +-- incoming fragments, and discovers path MTU via binary search probing. +module GBNet.Fragment + ( -- * Constants + fragmentHeaderSize, + maxFragmentCount, + defaultProbeTimeoutMs, + defaultMaxProbeAttempts, + minMtu, + maxMtu, + mtuConvergenceThreshold, + + -- * Fragment header + FragmentHeader (..), + serializeFragmentHeader, + deserializeFragmentHeader, + + -- * Fragmentation + FragmentError (..), + fragmentMessage, + + -- * Reassembly + FragmentBuffer, + FragmentAssembler (..), + newFragmentAssembler, + processFragment, + cleanupFragments, + + -- * MTU discovery + MtuState, + MtuDiscovery (..), + newMtuDiscovery, + defaultMtuDiscovery, + nextProbe, + onProbeSuccess, + onProbeTimeout, + checkProbeTimeout, + discoveredMtu, + mtuIsComplete, + ) +where + +import Control.DeepSeq (NFData (..), rwhnf) +import Data.Bits (shiftL, shiftR, (.|.)) +import qualified Data.ByteString as BS +import Data.ByteString.Internal (unsafeCreate) +import qualified Data.ByteString.Unsafe as BSU +import Data.List (minimumBy) +import Data.Map.Strict (Map) +import qualified Data.Map.Strict as Map +import Data.Ord (comparing) +import Data.Word (Word32, Word8) +import Foreign.Storable (pokeByteOff) +import GBNet.Reliability (MonoTime, elapsedMs) +import GBNet.Types (MessageId (..)) +import Optics ((%~), (&), (.~), (?~)) +import Optics.TH (makeFieldLabelsNoPrefix) + +-- Constants + +-- | Fragment header size: message_id (4) + fragment_index (1) + fragment_count (1) = 6 bytes. +fragmentHeaderSize :: Int +fragmentHeaderSize = 6 + +-- | Maximum fragments per message. +maxFragmentCount :: Int +maxFragmentCount = 255 + +-- | Default probe timeout in milliseconds. +defaultProbeTimeoutMs :: Double +defaultProbeTimeoutMs = 500.0 + +-- | Default maximum probe attempts. +defaultMaxProbeAttempts :: Int +defaultMaxProbeAttempts = 10 + +-- | Minimum safe MTU. +minMtu :: Int +minMtu = 576 + +-- | Maximum typical MTU. +maxMtu :: Int +maxMtu = 1500 + +-- | MTU convergence threshold (stop when range is this small). +mtuConvergenceThreshold :: Int +mtuConvergenceThreshold = 1 + +-- | Fragment header. +data FragmentHeader = FragmentHeader + { fhMessageId :: !MessageId, + fhFragmentIndex :: !Word8, + fhFragmentCount :: !Word8 + } + deriving (Eq, Show) + +instance NFData FragmentHeader where + rnf (FragmentHeader mid idx cnt) = rnf mid `seq` rnf idx `seq` rnf cnt + +makeFieldLabelsNoPrefix ''FragmentHeader + +-- | Serialize fragment header to bytes. +-- Uses zero-allocation direct memory writes. +serializeFragmentHeader :: FragmentHeader -> BS.ByteString +serializeFragmentHeader !hdr = unsafeCreate fragmentHeaderSize $ \ptr -> do + let !msgId = unMessageId (fhMessageId hdr) + pokeByteOff ptr 0 (fromIntegral (msgId `shiftR` 24) :: Word8) + pokeByteOff ptr 1 (fromIntegral (msgId `shiftR` 16) :: Word8) + pokeByteOff ptr 2 (fromIntegral (msgId `shiftR` 8) :: Word8) + pokeByteOff ptr 3 (fromIntegral msgId :: Word8) + pokeByteOff ptr 4 (fhFragmentIndex hdr) + pokeByteOff ptr 5 (fhFragmentCount hdr) +{-# INLINE serializeFragmentHeader #-} + +-- | Deserialize fragment header from bytes. +-- Uses direct memory access for speed. +deserializeFragmentHeader :: BS.ByteString -> Maybe FragmentHeader +deserializeFragmentHeader !bs + | BS.length bs < fragmentHeaderSize = Nothing + | otherwise = + let !b0 = fromIntegral (BSU.unsafeIndex bs 0) :: Word32 + !b1 = fromIntegral (BSU.unsafeIndex bs 1) :: Word32 + !b2 = fromIntegral (BSU.unsafeIndex bs 2) :: Word32 + !b3 = fromIntegral (BSU.unsafeIndex bs 3) :: Word32 + !msgId = (b0 `shiftL` 24) .|. (b1 `shiftL` 16) .|. (b2 `shiftL` 8) .|. b3 + in Just + FragmentHeader + { fhMessageId = MessageId msgId, + fhFragmentIndex = BSU.unsafeIndex bs 4, + fhFragmentCount = BSU.unsafeIndex bs 5 + } +{-# INLINE deserializeFragmentHeader #-} + +-- | Fragmentation errors. +data FragmentError + = TooManyFragments + deriving (Eq, Show) + +instance NFData FragmentError where rnf = rwhnf + +-- | Split a message into fragments. +fragmentMessage :: MessageId -> BS.ByteString -> Int -> Either FragmentError [BS.ByteString] +fragmentMessage messageId dat maxFragmentPayload + | BS.null dat || maxFragmentPayload <= 0 = Right [] + | fragCount > maxFragmentCount = Left TooManyFragments + | otherwise = Right $ map makeFragment [0 .. fragCount - 1] + where + fragCount = (BS.length dat + maxFragmentPayload - 1) `div` maxFragmentPayload + + makeFragment :: Int -> BS.ByteString + makeFragment i = + let start = i * maxFragmentPayload + end = min ((i + 1) * maxFragmentPayload) (BS.length dat) + header = + FragmentHeader + { fhMessageId = messageId, + fhFragmentIndex = fromIntegral i, + fhFragmentCount = fromIntegral fragCount + } + in serializeFragmentHeader header <> BS.take (end - start) (BS.drop start dat) + +-- | Buffer for reassembling a single fragmented message. +data FragmentBuffer = FragmentBuffer + { fbFragments :: !(Map Word8 BS.ByteString), + fbFragmentCount :: !Word8, + fbCreatedAt :: !MonoTime, + fbTotalSize :: !Int + } + deriving (Show) + +makeFieldLabelsNoPrefix ''FragmentBuffer + +-- | Create a new fragment buffer. +newFragmentBuffer :: Word8 -> MonoTime -> FragmentBuffer +newFragmentBuffer count now = + FragmentBuffer + { fbFragments = Map.empty, + fbFragmentCount = count, + fbCreatedAt = now, + fbTotalSize = 0 + } + +-- | Insert a fragment. Returns (isComplete, updatedBuffer). +insertFragment :: Word8 -> BS.ByteString -> FragmentBuffer -> (Bool, FragmentBuffer) +insertFragment idx dat buf + | idx >= fbFragmentCount buf = (False, buf) + | Map.member idx (fbFragments buf) = (isComplete buf, buf) -- Already have this fragment + | otherwise = + let buf' = + buf + & #fbFragments + %~ Map.insert idx dat + & #fbTotalSize + %~ (+ BS.length dat) + in (isComplete buf', buf') + +-- | Check if all fragments received. +isComplete :: FragmentBuffer -> Bool +isComplete buf = Map.size (fbFragments buf) == fromIntegral (fbFragmentCount buf) + +-- | Assemble complete message from fragments. +assembleFragments :: FragmentBuffer -> Maybe BS.ByteString +assembleFragments buf + | not (isComplete buf) = Nothing + | otherwise = + let indices = [0 .. fbFragmentCount buf - 1] + fragments = mapM (`Map.lookup` fbFragments buf) indices + in BS.concat <$> fragments + +-- | Fragment reassembler managing multiple in-progress messages. +data FragmentAssembler = FragmentAssembler + { faBuffers :: !(Map MessageId FragmentBuffer), + faTimeoutMs :: !Double, + faMaxBufferSize :: !Int, + faCurrentBufferSize :: !Int + } + deriving (Show) + +instance NFData FragmentAssembler where rnf = rwhnf + +makeFieldLabelsNoPrefix ''FragmentAssembler + +-- | Create a new fragment assembler. +newFragmentAssembler :: Double -> Int -> FragmentAssembler +newFragmentAssembler timeoutMs maxSize = + FragmentAssembler + { faBuffers = Map.empty, + faTimeoutMs = timeoutMs, + faMaxBufferSize = maxSize, + faCurrentBufferSize = 0 + } + +-- | Process an incoming fragment. Returns reassembled message if complete. +processFragment :: BS.ByteString -> MonoTime -> FragmentAssembler -> (Maybe BS.ByteString, FragmentAssembler) +processFragment dat now asm0 = + let asm1 = cleanupFragments now asm0 + in case deserializeFragmentHeader dat of + Nothing -> (Nothing, asm1) + Just header -> + let fragData = BS.drop fragmentHeaderSize dat + fragSize = BS.length fragData + msgId = fhMessageId header + in case Map.lookup msgId (faBuffers asm1) of + Just existing + | fbFragmentCount existing /= fhFragmentCount header -> + (Nothing, asm1) + _ -> + let -- Enforce memory limit + asm2 + | faCurrentBufferSize asm1 + fragSize > faMaxBufferSize asm1 = + expireOldest asm1 + | otherwise = asm1 + -- Get or create buffer + buf = case Map.lookup msgId (faBuffers asm2) of + Just b -> b + Nothing -> newFragmentBuffer (fhFragmentCount header) now + -- Insert fragment + (complete, buf') = insertFragment (fhFragmentIndex header) fragData buf + asm3 = + asm2 + & #faBuffers + %~ Map.insert msgId buf' + & #faCurrentBufferSize + %~ (+ fragSize) + in if complete + then + let result = assembleFragments buf' + asm4 = + asm3 + & #faBuffers + %~ Map.delete msgId + & #faCurrentBufferSize + %~ subtract (fbTotalSize buf') + in (result, asm4) + else (Nothing, asm3) + +-- | Remove expired incomplete fragment buffers. +cleanupFragments :: MonoTime -> FragmentAssembler -> FragmentAssembler +cleanupFragments now asm = + let timeout = faTimeoutMs asm + (expired, kept) = Map.partition (\buf -> elapsedMs (fbCreatedAt buf) now >= timeout) (faBuffers asm) + removedSize = sum $ map fbTotalSize $ Map.elems expired + in asm + & #faBuffers + .~ kept + & #faCurrentBufferSize + %~ subtract removedSize + +-- | Expire the oldest buffer to make room. +expireOldest :: FragmentAssembler -> FragmentAssembler +expireOldest asm = + case findOldest (faBuffers asm) of + Nothing -> asm + Just (oldestId, oldestBuf) -> + asm + & #faBuffers + %~ Map.delete oldestId + & #faCurrentBufferSize + %~ subtract (fbTotalSize oldestBuf) + where + findOldest :: Map MessageId FragmentBuffer -> Maybe (MessageId, FragmentBuffer) + findOldest m = + case Map.toList m of + [] -> Nothing + xs -> Just $ minimumBy (comparing (fbCreatedAt . snd)) xs + +-- | MTU discovery state. +data MtuState + = MtuProbing + | MtuComplete + deriving (Eq, Show) + +-- | MTU discovery via binary search. +data MtuDiscovery = MtuDiscovery + { mdMinMtu :: !Int, + mdMaxMtu :: !Int, + mdCurrentProbe :: !Int, + mdDiscoveredMtu :: !Int, + mdState :: !MtuState, + mdProbeTimeoutMs :: !Double, + mdLastProbeTime :: !(Maybe MonoTime), + mdAttempts :: !Int, + mdMaxAttempts :: !Int + } + deriving (Show) + +makeFieldLabelsNoPrefix ''MtuDiscovery + +-- | Create MTU discovery with custom bounds. +newMtuDiscovery :: Int -> Int -> MtuDiscovery +newMtuDiscovery minM maxM = + let initialProbe = (minM + maxM) `div` 2 + in MtuDiscovery + { mdMinMtu = minM, + mdMaxMtu = maxM, + mdCurrentProbe = initialProbe, + mdDiscoveredMtu = minM, -- Start with safe default + mdState = MtuProbing, + mdProbeTimeoutMs = defaultProbeTimeoutMs, + mdLastProbeTime = Nothing, + mdAttempts = 0, + mdMaxAttempts = defaultMaxProbeAttempts + } + +-- | Create MTU discovery with default bounds. +defaultMtuDiscovery :: MtuDiscovery +defaultMtuDiscovery = newMtuDiscovery minMtu maxMtu + +-- | Get next probe size, or Nothing if complete. +nextProbe :: MonoTime -> MtuDiscovery -> (Maybe Int, MtuDiscovery) +nextProbe now md + | mdState md == MtuComplete || mdAttempts md >= mdMaxAttempts md = + (Nothing, md & #mdState .~ MtuComplete) + | mdMaxMtu md - mdMinMtu md <= mtuConvergenceThreshold = + (Nothing, md & #mdState .~ MtuComplete) + | otherwise = + case mdLastProbeTime md of + Just lastTime + | elapsedMs lastTime now < mdProbeTimeoutMs md -> + (Nothing, md) + _ -> + let probe = (mdMinMtu md + mdMaxMtu md) `div` 2 + md' = + md + & #mdCurrentProbe + .~ probe + & #mdLastProbeTime + ?~ now + & #mdAttempts + %~ (+ 1) + in (Just probe, md') + +-- | Called when probe succeeded (ack received). +onProbeSuccess :: Int -> MtuDiscovery -> MtuDiscovery +onProbeSuccess size md + | size >= mdMinMtu md = + md + & #mdDiscoveredMtu + .~ size + & #mdMinMtu + .~ size + | otherwise = md + +-- | Called when probe timed out (too large). +onProbeTimeout :: MtuDiscovery -> MtuDiscovery +onProbeTimeout md = md & #mdMaxMtu .~ mdCurrentProbe md + +-- | Check if current probe timed out. +checkProbeTimeout :: MonoTime -> MtuDiscovery -> MtuDiscovery +checkProbeTimeout now md + | mdState md == MtuComplete = md + | otherwise = + case mdLastProbeTime md of + Just lastTime + | elapsedMs lastTime now >= mdProbeTimeoutMs md -> + onProbeTimeout md + _ -> md + +-- | Get discovered MTU. +discoveredMtu :: MtuDiscovery -> Int +discoveredMtu = mdDiscoveredMtu + +-- | Check if MTU discovery is complete. +mtuIsComplete :: MtuDiscovery -> Bool +mtuIsComplete md = mdState md == MtuComplete diff --git a/src/GBNet/Net.hs b/src/GBNet/Net.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Net.hs @@ -0,0 +1,195 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedLabels #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- Module : GBNet.Net +-- Description : NetT monad transformer for network operations +-- +-- Provides a monad transformer that carries network state (socket, local address) +-- and implements MonadNetwork. This allows polymorphic network code that can +-- run against real IO or pure test backends. +-- +-- The IO backend uses a dedicated receive thread that blocks efficiently on +-- the socket via GHC's IO manager (epoll\/kqueue), delivering packets through +-- an STM 'TQueue' with zero polling overhead. +module GBNet.Net + ( -- * Network state + NetState (..), + newNetState, + + -- * NetT monad transformer + NetT (..), + runNetT, + evalNetT, + execNetT, + + -- * State access + getNetState, + getsNetState, + modifyNetState, + ) +where + +import Control.Concurrent (ThreadId, forkIO, killThread) +import Control.Concurrent.STM (TQueue, atomically, newTQueueIO, tryReadTQueue, writeTQueue) +import Control.Exception (IOException, handle) +import Control.Monad.IO.Class (MonadIO (..)) +import Control.Monad.State.Strict (StateT (..), evalStateT, execStateT, get, gets, modify') +import qualified Data.ByteString as BS +import GBNet.Class (MonadNetwork (..), MonadTime (..), NetError (..), getMonoTimeIO) +import GBNet.Security (validateAndStripCrc32) +import GBNet.Socket + ( UdpSocket (..), + closeSocket, + maxUdpPacketSize, + socketSendTo, + ) +import Network.Socket (SockAddr, Socket) +import qualified Network.Socket.ByteString as NSB +import Optics ((%), (%~), (&), (.~), (?~)) +import Optics.TH (makeFieldLabelsNoPrefix) + +-- | Network state carried by NetT. +-- +-- Contains a 'TQueue' fed by a dedicated receive thread, eliminating +-- the need for polling timeouts on the socket. +data NetState = NetState + { -- | The UDP socket (used for sending) + nsSocket :: !UdpSocket, + -- | Local address this peer is bound to + nsLocalAddr :: !SockAddr, + -- | Queue of received packets from the receive thread + nsRecvQueue :: !(TQueue (BS.ByteString, SockAddr)), + -- | Receive thread handle (for cleanup) + nsRecvThread :: !ThreadId + } + +makeFieldLabelsNoPrefix ''NetState + +-- | Create network state and start the dedicated receive thread. +-- +-- The receive thread blocks on the socket via GHC's IO manager +-- (epoll on Linux, kqueue on macOS), consuming zero CPU when idle. +-- Received packets are delivered through an STM 'TQueue'. +newNetState :: UdpSocket -> SockAddr -> IO NetState +newNetState sock addr = do + queue <- newTQueueIO + tid <- forkIO $ recvLoop (usSocket sock) queue + pure + NetState + { nsSocket = sock, + nsLocalAddr = addr, + nsRecvQueue = queue, + nsRecvThread = tid + } + +-- | Dedicated receive thread. +-- +-- Blocks on @recvFrom@ which GHC's runtime handles efficiently via the +-- IO manager — the green thread is parked with no CPU cost until data +-- arrives. Exits cleanly when the socket is closed. +recvLoop :: Socket -> TQueue (BS.ByteString, SockAddr) -> IO () +recvLoop sock queue = go + where + go = + handle (\(_ :: IOException) -> pure ()) $ do + (dat, addr) <- NSB.recvFrom sock maxUdpPacketSize + atomically $ writeTQueue queue (dat, addr) + go + +-- | Network monad transformer. +-- +-- Wraps StateT to carry network state. When @m@ is IO, this provides real +-- UDP networking via the 'MonadNetwork' instance defined below. +-- For testing, use @TestNet@ from "GBNet.TestNet" instead. +newtype NetT m a = NetT {unNetT :: StateT NetState m a} + deriving (Functor, Applicative, Monad) + +-- | Run a NetT computation with initial state, returning result and final state. +runNetT :: NetT m a -> NetState -> m (a, NetState) +runNetT (NetT m) = runStateT m + +-- | Run a NetT computation, discarding final state. +evalNetT :: (Monad m) => NetT m a -> NetState -> m a +evalNetT (NetT m) = evalStateT m + +-- | Run a NetT computation, returning only final state. +execNetT :: (Monad m) => NetT m a -> NetState -> m NetState +execNetT (NetT m) = execStateT m + +-- | Get the current network state. +getNetState :: (Monad m) => NetT m NetState +getNetState = NetT get + +-- | Get a field from network state. +getsNetState :: (Monad m) => (NetState -> a) -> NetT m a +getsNetState = NetT . gets + +-- | Modify network state. +modifyNetState :: (Monad m) => (NetState -> NetState) -> NetT m () +modifyNetState = NetT . modify' + +-- | MonadIO instance for NetT. +instance (MonadIO m) => MonadIO (NetT m) where + liftIO = NetT . liftIO + +-- | MonadTime instance for NetT IO. +instance MonadTime (NetT IO) where + getMonoTime = liftIO getMonoTimeIO + +-- | MonadNetwork instance for NetT IO. +-- +-- Receives are non-blocking reads from the TQueue (fed by the dedicated +-- receive thread). Sends go directly through the socket. +instance MonadNetwork (NetT IO) where + netSend toAddr bytes = do + st <- getNetState + now <- liftIO getMonoTimeIO + result <- liftIO $ socketSendTo bytes toAddr now (nsSocket st) + case result of + Left err -> pure $ Left (NetSendFailed (show err)) + Right (_, sock') -> do + modifyNetState $ #nsSocket .~ sock' + pure $ Right () + + netRecv = do + st <- getNetState + result <- liftIO $ atomically $ tryReadTQueue (nsRecvQueue st) + case result of + Nothing -> pure Nothing + Just (dat, addr) -> do + -- Update receive stats on the socket + now <- liftIO getMonoTimeIO + let len = BS.length dat + modifyNetState $ + #nsSocket + % #usStats + %~ ( \stats -> + stats + & #ssBytesReceived + %~ (+ fromIntegral len) + & #ssPacketsReceived + %~ (+ 1) + & #ssLastReceiveTime + ?~ now + ) + -- Validate CRC before returning + case validateAndStripCrc32 dat of + Nothing -> do + -- Increment CRC drop counter + modifyNetState $ #nsSocket % #usStats % #ssCrcDrops %~ (+ 1) + pure Nothing + Just validated -> pure $ Just (validated, addr) + + netClose = do + st <- getNetState + liftIO $ killThread (nsRecvThread st) + liftIO $ closeSocket (nsSocket st) diff --git a/src/GBNet/Net/IO.hs b/src/GBNet/Net/IO.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Net/IO.hs @@ -0,0 +1,34 @@ +-- | +-- Module : GBNet.Net.IO +-- Description : IO-based network initialization +-- +-- Provides initialization helpers for creating NetState with real sockets. +module GBNet.Net.IO + ( -- * Socket operations + initNetState, + ) +where + +import GBNet.Net (NetState, newNetState) +import GBNet.Socket + ( SocketError (..), + closeSocket, + newUdpSocket, + socketLocalAddr, + ) +import Network.Socket (SockAddr) + +-- | Initialize NetState by creating a socket bound to the given address. +-- Returns the socket and local address for constructing NetState. +initNetState :: SockAddr -> IO (Either SocketError NetState) +initNetState bindAddr = do + socketResult <- newUdpSocket bindAddr + case socketResult of + Left err -> pure $ Left err + Right sock -> do + addrResult <- socketLocalAddr sock + case addrResult of + Left err -> do + closeSocket sock + pure $ Left err + Right localAddr -> Right <$> newNetState sock localAddr diff --git a/src/GBNet/Packet.hs b/src/GBNet/Packet.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Packet.hs @@ -0,0 +1,195 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- Module : GBNet.Packet +-- Description : Packet types and header definitions for the wire protocol +-- +-- Defines 'PacketType' and 'PacketHeader' for the gbnet wire format. +-- Header is 68 bits: 4-bit type + 16-bit sequence + 16-bit ack + 32-bit ack bitfield. +module GBNet.Packet + ( -- * Types + PacketType (..), + PacketHeader (..), + Packet (..), + + -- * Constants + packetHeaderBitSize, + packetHeaderByteSize, + + -- * Serialization + serializePacket, + deserializePacket, + serializeHeader, + deserializeHeader, + ) +where + +import Control.DeepSeq (NFData (..), rwhnf) +import Data.Bits (shiftL, shiftR, (.&.), (.|.)) +import qualified Data.ByteString as BS +import Data.ByteString.Internal (unsafeCreate) +import qualified Data.ByteString.Unsafe as BSU +import Data.Word (Word16, Word32, Word8) +import Foreign.Storable (pokeByteOff) +import GBNet.Types (SequenceNum (..)) +import Optics.TH (makeFieldLabelsNoPrefix) + +-- | Header size in bits (4 + 16 + 16 + 32 = 68). +packetHeaderBitSize :: Int +packetHeaderBitSize = 68 + +-- | Packet type tag (4 bits on wire). +-- Additional fields are serialized after the header. +data PacketType + = -- | 0: Client initiates connection + ConnectionRequest + | -- | 1: Server accepts + ConnectionAccepted + | -- | 2: Server rejects (+ 8-bit reason in payload) + ConnectionDenied + | -- | 3: Normal game data (+ 3-bit channel, 1-bit is_fragment in payload) + Payload + | -- | 4: Graceful disconnect (+ 8-bit reason in payload) + Disconnect + | -- | 5: Keep connection alive + Keepalive + | -- | 6: Server challenge (+ 64-bit server_salt in payload) + ConnectionChallenge + | -- | 7: Client response (+ 64-bit client_salt in payload) + ConnectionResponse + deriving (Eq, Show, Enum, Bounded) + +instance NFData PacketType where rnf = rwhnf + +-- | Packet header (68 bits on wire). +data PacketHeader = PacketHeader + { -- | 4 bits + packetType :: !PacketType, + -- | 16 bits + sequenceNum :: !SequenceNum, + -- | 16 bits - most recent received sequence + ack :: !SequenceNum, + -- | 32 bits - preceding 32 acks + ackBitfield :: !Word32 + } + deriving (Eq, Show) + +instance NFData PacketHeader where + rnf (PacketHeader pt sn ak abf) = rnf pt `seq` rnf sn `seq` rnf ak `seq` rnf abf + +-- | Header size in bytes (68 bits = 9 bytes, rounded up). +packetHeaderByteSize :: Int +packetHeaderByteSize = (packetHeaderBitSize + 7) `div` 8 + +-- | A complete packet with header and payload. +data Packet = Packet + { pktHeader :: !PacketHeader, + pktPayload :: !BS.ByteString + } + deriving (Eq, Show) + +makeFieldLabelsNoPrefix ''PacketHeader +makeFieldLabelsNoPrefix ''Packet + +-- | Bit shift for packet type in first byte. +packetTypeBitShift :: Int +packetTypeBitShift = 4 + +-- | Serialize a packet header to bytes. +-- Uses optimized direct memory writes. +-- +-- Wire format (68 bits, MSB-first): +-- Byte 0: [type:4][seq_hi:4] +-- Byte 1: [seq_mid:8] +-- Byte 2: [seq_lo:4][ack_hi:4] +-- Byte 3: [ack_mid:8] +-- Byte 4: [ack_lo:4][abf_hi:4] +-- Bytes 5-7: [abf:24] +-- Byte 8: [abf_lo:4][pad:4] +serializeHeader :: PacketHeader -> BS.ByteString +serializeHeader !hdr = unsafeCreate packetHeaderByteSize $ \ptr -> do + let !pt = fromIntegral (fromEnum (packetType hdr)) :: Word8 + !(SequenceNum !sn) = sequenceNum hdr + !(SequenceNum !ak) = ack hdr + !abf = ackBitfield hdr + !byte0 = (pt `shiftL` packetTypeBitShift) .|. fromIntegral (sn `shiftR` 12) + !byte1 = fromIntegral (sn `shiftR` 4) :: Word8 + !byte2 = ((fromIntegral (sn .&. 0x0F) :: Word8) `shiftL` 4) .|. fromIntegral (ak `shiftR` 12) + !byte3 = fromIntegral (ak `shiftR` 4) :: Word8 + !byte4 = ((fromIntegral (ak .&. 0x0F) :: Word8) `shiftL` 4) .|. fromIntegral (abf `shiftR` 28) + !byte5 = fromIntegral (abf `shiftR` 20) :: Word8 + !byte6 = fromIntegral (abf `shiftR` 12) :: Word8 + !byte7 = fromIntegral (abf `shiftR` 4) :: Word8 + !byte8 = fromIntegral (abf .&. 0x0F) `shiftL` 4 :: Word8 + pokeByteOff ptr 0 byte0 + pokeByteOff ptr 1 byte1 + pokeByteOff ptr 2 byte2 + pokeByteOff ptr 3 byte3 + pokeByteOff ptr 4 byte4 + pokeByteOff ptr 5 byte5 + pokeByteOff ptr 6 byte6 + pokeByteOff ptr 7 byte7 + pokeByteOff ptr 8 byte8 +{-# INLINE serializeHeader #-} + +-- | Deserialize a packet header from bytes. +-- Uses optimized direct memory access. +deserializeHeader :: BS.ByteString -> Either String PacketHeader +deserializeHeader !bs + | BS.length bs < packetHeaderByteSize = Left "Header too short" + | otherwise = + let !b0 = BSU.unsafeIndex bs 0 + !b1 = BSU.unsafeIndex bs 1 + !b2 = BSU.unsafeIndex bs 2 + !b3 = BSU.unsafeIndex bs 3 + !b4 = BSU.unsafeIndex bs 4 + !b5 = BSU.unsafeIndex bs 5 + !b6 = BSU.unsafeIndex bs 6 + !b7 = BSU.unsafeIndex bs 7 + !b8 = BSU.unsafeIndex bs 8 + !ptVal = fromIntegral (b0 `shiftR` packetTypeBitShift) :: Int + !snHi = fromIntegral (b0 .&. 0x0F) :: Word16 + !snMid = fromIntegral b1 :: Word16 + !snLo = fromIntegral (b2 `shiftR` 4) :: Word16 + !sn = (snHi `shiftL` 12) .|. (snMid `shiftL` 4) .|. snLo + !akHi = fromIntegral (b2 .&. 0x0F) :: Word16 + !akMid = fromIntegral b3 :: Word16 + !akLo = fromIntegral (b4 `shiftR` 4) :: Word16 + !ak = (akHi `shiftL` 12) .|. (akMid `shiftL` 4) .|. akLo + !abf0 = fromIntegral (b4 .&. 0x0F) :: Word32 + !abf1 = fromIntegral b5 :: Word32 + !abf2 = fromIntegral b6 :: Word32 + !abf3 = fromIntegral b7 :: Word32 + !abf4 = fromIntegral (b8 `shiftR` 4) :: Word32 + !abf = (abf0 `shiftL` 28) .|. (abf1 `shiftL` 20) .|. (abf2 `shiftL` 12) .|. (abf3 `shiftL` 4) .|. abf4 + in if ptVal > fromEnum (maxBound :: PacketType) + then Left $ "Invalid packet type: " ++ show ptVal + else + Right + PacketHeader + { packetType = toEnum ptVal, + sequenceNum = SequenceNum sn, + ack = SequenceNum ak, + ackBitfield = abf + } +{-# INLINE deserializeHeader #-} + +-- | Serialize a complete packet (header + payload) to bytes. +serializePacket :: Packet -> BS.ByteString +serializePacket pkt = + let headerBytes = serializeHeader (pktHeader pkt) + in headerBytes <> pktPayload pkt + +-- | Deserialize a complete packet from bytes. +deserializePacket :: BS.ByteString -> Either String Packet +deserializePacket bs = do + header <- deserializeHeader bs + let payload = BS.drop packetHeaderByteSize bs + Right Packet {pktHeader = header, pktPayload = payload} diff --git a/src/GBNet/Peer.hs b/src/GBNet/Peer.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Peer.hs @@ -0,0 +1,559 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedLabels #-} + +-- | +-- Module : GBNet.Peer +-- Description : Unified peer networking API +-- +-- NetPeer provides a symmetric abstraction for game networking. +-- A peer can accept incoming connections (server-like), initiate +-- outgoing connections (client-like), or do both (P2P/mesh). +-- +-- Polymorphic game loop pattern: +-- +-- @ +-- gameLoop peer = do +-- (events, peer') <- peerTick [(channel, msg)] peer +-- -- handle events +-- gameLoop peer' +-- @ +module GBNet.Peer + ( -- * Peer identifier + PeerId (..), + peerIdFromAddr, + + -- * Connection direction + ConnectionDirection (..), + + -- * Events + PeerEvent (..), + + -- * Pure processing types + IncomingPacket (..), + RawPacket (..), + PeerResult (..), + + -- * Net peer + NetPeer (..), + newPeer, + newPeerState, + + -- * Connection management + peerConnect, + peerDisconnect, + + -- * Pure processing + peerProcess, + + -- * Polymorphic IO helpers + peerRecvAllM, + peerSendAllM, + peerShutdownM, + peerTick, + + -- * Internal (used by pure processing) + drainPeerSendQueue, + drainAllConnectionQueues, + + -- * Sending + peerSend, + peerBroadcast, + + -- * Pending connection (opaque) + PendingConnection, + + -- * Queries + peerCount, + peerIsConnected, + peerStats, + peerLocalAddr, + peerConnectedIds, + ) +where + +import qualified Data.ByteString as BS +import Data.Either (fromRight) +import Data.List (foldl') +import qualified Data.Map.Strict as Map +import GBNet.Class (MonadNetwork (..), MonadTime (..), MonoTime) +import GBNet.Config (NetworkConfig (..)) +import GBNet.Connection + ( ConnectionError (..), + DisconnectReason (..), + OutgoingPacket (..), + processIncomingHeader, + receiveIncomingPayload, + ) +import qualified GBNet.Connection as Conn +import GBNet.Fragment (newFragmentAssembler, processFragment) +import GBNet.Packet + ( Packet (..), + PacketHeader (..), + PacketType (..), + deserializePacket, + serializePacket, + ) +import GBNet.Peer.Handshake + ( handleConnectionAccepted, + handleConnectionChallenge, + handleConnectionRequest, + handleConnectionResponse, + handleDisconnect, + ) +import GBNet.Peer.Internal +import GBNet.Peer.Migration + ( findMigrationCandidate, + migrationCooldownMs, + ) +import GBNet.Peer.Protocol + ( decodeDenyReason, + decodePayloadHeader, + denyToDisconnectReason, + minPayloadSize, + ) +import qualified GBNet.Peer.Protocol as Proto +import GBNet.Reliability (elapsedMs) +import GBNet.Security (appendCrc32) +import GBNet.Socket + ( SocketError (..), + UdpSocket, + newUdpSocket, + socketLocalAddr, + ) +import GBNet.Stats (NetworkStats) +import GBNet.Types (ChannelId (..)) +import GBNet.Util (nextRandom) +import Network.Socket (SockAddr) +import Optics ((%~), (&), (.~)) + +-- | Create a new peer bound to the given address. +-- Returns the peer and socket. The socket is also stored in the peer for +-- backward compatibility, but new code should use the polymorphic API. +newPeer :: + SockAddr -> + NetworkConfig -> + MonoTime -> + IO (Either SocketError (NetPeer, UdpSocket)) +newPeer addr config now = do + socketResult <- newUdpSocket addr + case socketResult of + Left err -> return $ Left err + Right sock -> do + localAddrResult <- socketLocalAddr sock + let localAddr = case localAddrResult of + Left _ -> addr -- Fallback to bind address + Right a -> a + let peer = newPeerState sock localAddr config now + return $ Right (peer, sock) + +-- ----------------------------------------------------------------------------- +-- Connection management +-- ----------------------------------------------------------------------------- + +-- | Initiate an outgoing connection to a peer (pure). +peerConnect :: PeerId -> MonoTime -> NetPeer -> NetPeer +peerConnect peerId now peer + | Map.member peerId (npConnections peer) = peer -- Already connected + | Map.member peerId (npPending peer) = peer -- Already pending + | otherwise = + -- Generate client salt + let (salt, rng') = nextRandom (npRngState peer) + pending = + PendingConnection + { pcDirection = Outbound, + pcServerSalt = 0, -- Will be filled when we receive challenge + pcClientSalt = salt, + pcCreatedAt = now, + pcRetryCount = 0, + pcLastRetry = now + } + peer' = + peer + & #npPending + %~ Map.insert peerId pending + & #npRngState + .~ rng' + in -- Queue connection request + queueControlPacket ConnectionRequest BS.empty peerId peer' + +-- | Disconnect a specific peer (pure). +-- Transitions the connection to Disconnecting state for graceful shutdown +-- with retries, rather than removing it immediately. +peerDisconnect :: PeerId -> MonoTime -> NetPeer -> NetPeer +peerDisconnect peerId now peer = + case Map.lookup peerId (npConnections peer) of + Nothing -> peer + Just _ -> withConnection peerId (Conn.disconnect ReasonRequested now) peer + +-- ----------------------------------------------------------------------------- +-- Pure processing +-- ----------------------------------------------------------------------------- + +-- | Pure packet processing function. +-- Given the current time and a list of incoming packets, returns the updated +-- peer state, events that occurred, and packets to send. +-- +-- This is the core of the game loop - it's completely pure and deterministic. +-- Use 'peerRecvAllM' to get incoming packets and 'peerSendAllM' to send outgoing. +-- +-- See 'peerTick' for a convenience wrapper that combines receive, process, and send. +peerProcess :: MonoTime -> [IncomingPacket] -> NetPeer -> PeerResult +peerProcess now packets peer0 = + let -- Process incoming packets + internalPackets = map (\ip -> (ipFrom ip, ipData ip)) packets + (events1, peer1) = processPacketsPure internalPackets now peer0 + -- Update all connections and collect messages + (events2, peer2) = updateConnections now peer1 + -- Drain connection send queues into peer send queue + peer3 = drainAllConnectionQueues now peer2 + -- Retry pending outbound connections + peer4 = retryPendingConnectionsPure now peer3 + -- Cleanup expired pending connections + (events3, peer5) = cleanupPending now peer4 + -- Drain the peer's send queue + (outgoing, peer6) = drainPeerSendQueue peer5 + in PeerResult peer6 (events1 ++ events2 ++ events3) outgoing + +-- | Drain send queues from all connections into the peer's send queue. +-- Single-pass over the connection map via foldlWithKey'. +drainAllConnectionQueues :: MonoTime -> NetPeer -> NetPeer +drainAllConnectionQueues _now peer = + Map.foldlWithKey' drainOne (peer & #npConnections .~ Map.empty) (npConnections peer) + where + drainOne p peerId conn = + let (connPackets, conn') = Conn.drainSendQueue conn + rawPackets = map (outgoingToRaw peerId) connPackets + p' = p & #npConnections %~ Map.insert peerId conn' + in foldl' (flip queueRawPacket) p' rawPackets + + outgoingToRaw peerId (OutgoingPacket hdr ptype payload) = + let header = hdr {packetType = ptype} + pkt = Packet {pktHeader = header, pktPayload = payload} + raw = appendCrc32 (serializePacket pkt) + in RawPacket peerId raw + +-- ----------------------------------------------------------------------------- +-- Polymorphic IO helpers +-- ----------------------------------------------------------------------------- + +-- | Receive all available packets (polymorphic, non-blocking). +-- Returns immediately if no data is available. +peerRecvAllM :: (MonadNetwork m) => m [IncomingPacket] +peerRecvAllM = go [] + where + go acc = do + result <- netRecv + case result of + Nothing -> pure (reverse acc) + Just (dat, addr) -> + let pkt = IncomingPacket (PeerId addr) dat + in go (pkt : acc) + +-- | Send all outgoing packets (polymorphic). +peerSendAllM :: (MonadNetwork m) => [RawPacket] -> m () +peerSendAllM = mapM_ sendOne + where + sendOne (RawPacket pid dat) = netSend (unPeerId pid) dat + +-- | Shutdown the peer (polymorphic). +-- Disconnects all connections through the state machine and closes the network. +peerShutdownM :: (MonadNetwork m) => NetPeer -> m () +peerShutdownM peer = do + now <- getMonoTime + let peerIds = Map.keys (npConnections peer) + -- Disconnect each connection through the proper state machine + peer' = foldr (\pid p -> withConnection pid (Conn.disconnect ReasonRequested now) p) peer peerIds + (outgoing, _) = drainPeerSendQueue (drainAllConnectionQueues now peer') + peerSendAllM outgoing + netClose + +-- | Convenient single-function tick for game loops (polymorphic). +-- +-- Combines receive, process, queue messages, and send into one call. +-- Takes a list of (channel, message) pairs to send and returns events. +-- +-- @ +-- gameLoop peer = do +-- (events, peer') <- peerTick [(ch, msg)] peer +-- -- handle events +-- gameLoop peer' +-- @ +peerTick :: + (MonadNetwork m) => + [(ChannelId, BS.ByteString)] -> + NetPeer -> + m ([PeerEvent], NetPeer) +peerTick messages peer = do + now <- getMonoTime + -- 1. Receive all available packets + packets <- peerRecvAllM + -- 2. Queue messages to all connections + let peer1 = foldl' (queueMessage now) peer messages + -- 3. Process packets (pure) + let result = peerProcess now packets peer1 + peer2 = prPeer result + events = prEvents result + outgoing = prOutgoing result + -- 4. Send all outgoing packets + peerSendAllM outgoing + pure (events, peer2) + where + queueMessage now p (ch, msg) = peerBroadcast ch msg Nothing now p + +-- ----------------------------------------------------------------------------- +-- Packet handling (pure) +-- ----------------------------------------------------------------------------- + +-- | Process received packets (pure). +processPacketsPure :: + [(PeerId, BS.ByteString)] -> + MonoTime -> + NetPeer -> + ([PeerEvent], NetPeer) +processPacketsPure packets now peer = foldl' go ([], peer) packets + where + go (evts, p) (pid, dat) = + let (evts', p') = handlePacket pid dat now p + in (evts ++ evts', p') + +-- | Handle a single received packet (pure). +handlePacket :: + PeerId -> + BS.ByteString -> + MonoTime -> + NetPeer -> + ([PeerEvent], NetPeer) +handlePacket peerId dat now peer = + case parsePacket dat of + Nothing -> ([], peer) + Just pkt -> handlePacketByType peerId pkt now (packetType (pktHeader pkt)) peer + +-- | Parse packet from raw data. +parsePacket :: BS.ByteString -> Maybe Packet +parsePacket dat = + case deserializePacket dat of + Left _ -> Nothing + Right pkt -> Just pkt + +-- | Handle a packet by its type (pure). +handlePacketByType :: + PeerId -> + Packet -> + MonoTime -> + PacketType -> + NetPeer -> + ([PeerEvent], NetPeer) +handlePacketByType peerId pkt now ptype peer = case ptype of + ConnectionRequest -> handleConnectionRequest peerId now peer + ConnectionChallenge -> handleConnectionChallenge peerId pkt now peer + ConnectionResponse -> handleConnectionResponse peerId pkt now peer + ConnectionAccepted -> handleConnectionAccepted peerId now peer + ConnectionDenied -> + let reason = decodeDenyReason (pktPayload pkt) + peer' = removePending peerId peer + in ([PeerDisconnected peerId (denyToDisconnectReason reason)], peer') + Disconnect -> handleDisconnect peerId peer + Payload -> + if Map.member peerId (npConnections peer) + then handlePayload peerId pkt now peer + else handleMigration peerId pkt now peer + Keepalive -> + let peer' = + withConnection + peerId + (Conn.touchRecvTime now . processIncomingHeader (pktHeader pkt) now) + peer + in ([], peer') + +-- | Handle payload packet (pure). +-- Routes messages through the channel system for proper ordering/dedup. +handlePayload :: PeerId -> Packet -> MonoTime -> NetPeer -> ([PeerEvent], NetPeer) +handlePayload peerId pkt now peer = + case Map.lookup peerId (npConnections peer) of + Nothing -> ([], peer) + Just conn -> + let conn' = Conn.touchRecvTime now $ processIncomingHeader (pktHeader pkt) now conn + payload = pktPayload pkt + in case BS.uncons payload of + Nothing -> + ([], peer & #npConnections %~ Map.insert peerId conn') + Just (headerByte, rest) -> + let (channel, isFragment) = decodePayloadHeader headerByte + in if isFragment + then + let peer' = peer & #npConnections %~ Map.insert peerId conn' + in handleFragment peerId channel rest now peer' + else + let finalConn + | BS.length payload < minPayloadSize = conn' + | otherwise = case Proto.decodeChannelSeq rest of + Nothing -> conn' + Just (chSeq, msgData) -> + receiveIncomingPayload channel chSeq msgData now conn' + in ([], peer & #npConnections %~ Map.insert peerId finalConn) + +-- | Handle a fragment, reassembling if complete (pure). +-- After reassembly, routes through the channel system for ordering/dedup. +handleFragment :: PeerId -> ChannelId -> BS.ByteString -> MonoTime -> NetPeer -> ([PeerEvent], NetPeer) +handleFragment peerId channel fragData now peer = + let assemblers = npFragmentAssemblers peer + assembler = + Map.findWithDefault + (newFragmentAssembler fragmentTimeoutMs fragmentMaxBufferSize) + peerId + assemblers + (maybeComplete, assembler') = processFragment fragData now assembler + peer' = peer & #npFragmentAssemblers .~ Map.insert peerId assembler' assemblers + in case maybeComplete of + Nothing -> ([], peer') + Just completeData -> + case Proto.decodeChannelSeq completeData of + Nothing -> ([], peer') + Just (chSeq, msgData) -> + ([], withConnection peerId (receiveIncomingPayload channel chSeq msgData now) peer') + +-- | Try to migrate an existing connection to a new address, or ignore the packet (pure). +handleMigration :: PeerId -> Packet -> MonoTime -> NetPeer -> ([PeerEvent], NetPeer) +handleMigration newPeerId pkt now peer = + if not (ncEnableConnectionMigration (npConfig peer)) + then ([], peer) + else case findMigrationCandidate pkt now peer of + Nothing -> ([], peer) + Just (oldPeerId, conn, migrationToken) -> + case Map.lookup migrationToken (npMigrationCooldowns peer) of + Just lastMigration + | elapsedMs lastMigration now < migrationCooldownMs -> + ([], peer) -- Still in cooldown + _ -> + let peer' = + peer + & #npConnections + %~ (Map.insert newPeerId conn . Map.delete oldPeerId) + & #npMigrationCooldowns + %~ Map.insert migrationToken now + & #npFragmentAssemblers + %~ ( \fa -> case Map.lookup oldPeerId fa of + Nothing -> fa + Just asm -> Map.insert newPeerId asm $ Map.delete oldPeerId fa + ) + event = PeerMigrated oldPeerId newPeerId + (payloadEvents, peer'') = handlePayload newPeerId pkt now peer' + in (event : payloadEvents, peer'') + +-- | Update all connections and collect messages/disconnects (pure). +-- Uses reverse accumulator to avoid O(n^2) list appending. +updateConnections :: MonoTime -> NetPeer -> ([PeerEvent], NetPeer) +updateConnections now peer = + let conns = npConnections peer + (revEvents, conns', disconnectedIds) = Map.foldlWithKey' updateOne ([], Map.empty, []) conns + peer' = foldl' (flip cleanupPeer) (peer & #npConnections .~ conns') disconnectedIds + in (reverse revEvents, peer') + where + updateOne (revEvts, connsAcc, discs) peerId conn = + case Conn.updateTick now conn of + Left _err -> + -- Connection timed out + (PeerDisconnected peerId ReasonTimeout : revEvts, connsAcc, peerId : discs) + Right conn' + | Conn.connectionState conn' == Conn.Disconnected -> + -- Graceful disconnect complete + (PeerDisconnected peerId ReasonRequested : revEvts, connsAcc, peerId : discs) + | otherwise -> + -- Collect messages from all channels + let (msgs, conn'') = collectMessages peerId conn' + in (prependReversed msgs revEvts, Map.insert peerId conn'' connsAcc, discs) + + collectMessages peerId conn = + let numChannels = Conn.channelCount conn + in collectFromChannels peerId 0 numChannels conn [] + + collectFromChannels peerId ch maxCh conn revAcc + | ch >= maxCh = (reverse revAcc, conn) + | otherwise = + let chId = ChannelId ch + (msgs, conn') = Conn.receiveMessage chId conn + evts = map (PeerMessage peerId chId) msgs + in collectFromChannels peerId (ch + 1) maxCh conn' (prependReversed evts revAcc) + +-- | Prepend a list in reverse onto an accumulator. O(length xs). +-- Used for efficient reverse-accumulator pattern. +prependReversed :: [a] -> [a] -> [a] +prependReversed xs ys = foldl (flip (:)) ys xs + +-- | Retry pending outbound connections (pure). +retryPendingConnectionsPure :: MonoTime -> NetPeer -> NetPeer +retryPendingConnectionsPure now peer = + let outbound = Map.toList $ Map.filter (\p -> pcDirection p == Outbound) (npPending peer) + in foldl' (retryOne now) peer outbound + where + retryOne t p (peerId, pending) = + let elapsed = elapsedMs (pcLastRetry pending) t + retryInterval = ncConnectionRequestTimeoutMs (npConfig p) / fromIntegral (ncConnectionRequestMaxRetries (npConfig p) + 1) + in if elapsed > retryInterval && pcRetryCount pending < ncConnectionRequestMaxRetries (npConfig p) + then + let pending' = pending & #pcRetryCount %~ (+ 1) & #pcLastRetry .~ t + p' = p & #npPending %~ Map.insert peerId pending' + in queueControlPacket ConnectionRequest BS.empty peerId p' + else p + +-- | Cleanup expired pending connections (pure). +cleanupPending :: MonoTime -> NetPeer -> ([PeerEvent], NetPeer) +cleanupPending now peer = + let timeout = ncConnectionRequestTimeoutMs (npConfig peer) + (expired, kept) = Map.partition (\p -> elapsedMs (pcCreatedAt p) now > timeout) (npPending peer) + events = map (\(pid, _) -> PeerDisconnected pid ReasonTimeout) (Map.toList expired) + in (events, peer & #npPending .~ kept) + +-- | Send a message to a connected peer. +peerSend :: + PeerId -> + ChannelId -> + BS.ByteString -> + MonoTime -> + NetPeer -> + Either ConnectionError NetPeer +peerSend peerId channel dat now peer = + case Map.lookup peerId (npConnections peer) of + Nothing -> Left ErrNotConnected + Just conn -> + case Conn.sendMessage channel dat now conn of + Left err -> Left err + Right conn' -> + Right (peer & #npConnections %~ Map.insert peerId conn') + +-- | Broadcast a message to all connected peers. +-- This queues the message and drains connection queues so packets are ready to send. +peerBroadcast :: + ChannelId -> + BS.ByteString -> + Maybe PeerId -> + MonoTime -> + NetPeer -> + NetPeer +peerBroadcast channel dat except now peer = + let peerIds = filter (\p -> Just p /= except) $ Map.keys (npConnections peer) + -- Queue message to each connection's channel + peer' = foldl' (\p pid -> fromRight p (peerSend pid channel dat now p)) peer peerIds + in -- Drain connection queues to npSendQueue so packets are ready + drainAllConnectionQueues now peer' + +-- | Get number of connected peers. +peerCount :: NetPeer -> Int +peerCount = Map.size . npConnections + +-- | Check if a peer is connected. +peerIsConnected :: PeerId -> NetPeer -> Bool +peerIsConnected peerId peer = Map.member peerId (npConnections peer) + +-- | Get stats for a connected peer. +peerStats :: PeerId -> NetPeer -> Maybe NetworkStats +peerStats peerId peer = + Conn.connectionStats <$> Map.lookup peerId (npConnections peer) + +-- | Get the local address. +peerLocalAddr :: NetPeer -> SockAddr +peerLocalAddr = npLocalAddr + +-- | Get list of all connected peer IDs. +peerConnectedIds :: NetPeer -> [PeerId] +peerConnectedIds = Map.keys . npConnections diff --git a/src/GBNet/Peer/Handshake.hs b/src/GBNet/Peer/Handshake.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Peer/Handshake.hs @@ -0,0 +1,167 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE MultiWayIf #-} +{-# LANGUAGE OverloadedLabels #-} + +-- | +-- Module : GBNet.Peer.Handshake +-- Description : Connection handshake state machine +-- +-- Handles connection request, challenge, response, accepted, and disconnect +-- packets during the handshake protocol. +module GBNet.Peer.Handshake + ( handleConnectionRequest, + handleNewConnectionRequest, + handleConnectionChallenge, + handleConnectionResponse, + handleConnectionAccepted, + handleDisconnect, + ) +where + +import qualified Data.ByteString as BS +import qualified Data.Map.Strict as Map +import GBNet.Class (MonoTime) +import GBNet.Config (NetworkConfig (..)) +import GBNet.Connection + ( DisconnectReason (..), + newConnection, + ) +import qualified GBNet.Connection as Conn +import GBNet.Packet (Packet (..), PacketType (..)) +import GBNet.Peer.Internal +import GBNet.Peer.Protocol + ( DenyReason (..), + decodeSalt, + encodeDenyReason, + encodeSalt, + sockAddrToKey, + ) +import GBNet.Security (rateLimiterAllow) +import GBNet.Util (nextRandom) +import Optics ((%~), (&), (.~)) + +-- | Handle incoming connection request (pure). +handleConnectionRequest :: PeerId -> MonoTime -> NetPeer -> ([PeerEvent], NetPeer) +handleConnectionRequest peerId now peer = + case (Map.member peerId (npConnections peer), Map.lookup peerId (npPending peer)) of + (True, _) -> + -- Already connected, resend accept + ([], queueControlPacket ConnectionAccepted BS.empty peerId peer) + (_, Just p) -> + -- Already pending, resend challenge with stored salt + let saltPayload = encodeSalt (pcServerSalt p) + in ([], queueControlPacket ConnectionChallenge saltPayload peerId peer) + (False, Nothing) -> + handleNewConnectionRequest peerId now peer + +-- | Handle a genuinely new connection request after checking existing state (pure). +handleNewConnectionRequest :: PeerId -> MonoTime -> NetPeer -> ([PeerEvent], NetPeer) +handleNewConnectionRequest peerId now peer = + let addrKey = sockAddrToKey (unPeerId peerId) + (allowed, rl') = rateLimiterAllow addrKey now (npRateLimiter peer) + peer1 = peer & #npRateLimiter .~ rl' + pendingSize = Map.size (npPending peer1) + connSize = Map.size (npConnections peer1) + maxClients = ncMaxClients (npConfig peer1) + in if + | not allowed -> + ([], peer1 & #npRateLimitDrops %~ (+ 1)) + | pendingSize >= maxClients -> + ([], peer1 & #npRateLimitDrops %~ (+ 1)) + | connSize >= maxClients -> + let reason = encodeDenyReason DenyServerFull + in ([], queueControlPacket ConnectionDenied reason peerId peer1) + | otherwise -> + let (salt, rng') = nextRandom (npRngState peer1) + newPend = + PendingConnection + { pcDirection = Inbound, + pcServerSalt = salt, + pcClientSalt = 0, + pcCreatedAt = now, + pcRetryCount = 0, + pcLastRetry = now + } + peer2 = + peer1 + & #npPending + %~ Map.insert peerId newPend + & #npRngState + .~ rng' + saltPayload = encodeSalt salt + in ([], queueControlPacket ConnectionChallenge saltPayload peerId peer2) + +-- | Handle connection challenge (we're outbound, received their challenge) (pure). +handleConnectionChallenge :: PeerId -> Packet -> MonoTime -> NetPeer -> ([PeerEvent], NetPeer) +handleConnectionChallenge peerId pkt _now peer = + case Map.lookup peerId (npPending peer) of + Nothing -> ([], peer) + Just p + | pcDirection p /= Outbound -> ([], peer) + | otherwise -> + case decodeSalt (pktPayload pkt) of + Nothing -> ([], peer) + Just serverSalt -> + let p' = p & #pcServerSalt .~ serverSalt + peer' = peer & #npPending %~ Map.insert peerId p' + saltPayload = encodeSalt (pcClientSalt p') + in ([], queueControlPacket ConnectionResponse saltPayload peerId peer') + +-- | Handle connection response (we're inbound, received their response) (pure). +handleConnectionResponse :: PeerId -> Packet -> MonoTime -> NetPeer -> ([PeerEvent], NetPeer) +handleConnectionResponse peerId pkt now peer = + case Map.lookup peerId (npPending peer) of + Nothing -> ([], peer) + Just p + | pcDirection p /= Inbound -> ([], peer) + | otherwise -> + case decodeSalt (pktPayload pkt) of + Nothing -> ([], peer) + Just clientSalt + | clientSalt == 0 || clientSalt == pcServerSalt p -> + let reason = encodeDenyReason DenyInvalidChallenge + peer' = + queueControlPacket ConnectionDenied reason peerId $ + removePending peerId peer + in ([], peer') + | otherwise -> + let conn = + Conn.markConnected now $ + Conn.touchRecvTime now $ + newConnection (npConfig peer) clientSalt now + peer' = + peer + & #npConnections + %~ Map.insert peerId conn + & #npPending + %~ Map.delete peerId + in ([PeerConnected peerId Inbound], queueControlPacket ConnectionAccepted BS.empty peerId peer') + +-- | Handle connection accepted (we're outbound, they accepted) (pure). +handleConnectionAccepted :: PeerId -> MonoTime -> NetPeer -> ([PeerEvent], NetPeer) +handleConnectionAccepted peerId now peer = + case Map.lookup peerId (npPending peer) of + Nothing -> ([], peer) + Just p + | pcDirection p /= Outbound -> ([], peer) + | otherwise -> + let conn = + Conn.markConnected now $ + Conn.touchRecvTime now $ + newConnection (npConfig peer) (pcClientSalt p) now + peer' = + peer + & #npConnections + %~ Map.insert peerId conn + & #npPending + %~ Map.delete peerId + in ([PeerConnected peerId Outbound], peer') + +-- | Handle disconnect packet (pure). +handleDisconnect :: PeerId -> NetPeer -> ([PeerEvent], NetPeer) +handleDisconnect peerId peer = + if Map.member peerId (npConnections peer) + then + let peer' = cleanupPeer peerId (peer & #npConnections %~ Map.delete peerId) + in ([PeerDisconnected peerId ReasonRequested], peer') + else ([], removePending peerId peer) diff --git a/src/GBNet/Peer/Internal.hs b/src/GBNet/Peer/Internal.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Peer/Internal.hs @@ -0,0 +1,265 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedLabels #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeApplications #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- Module : GBNet.Peer.Internal +-- Description : Shared types, constants, and helpers for Peer sub-modules +-- +-- This internal module provides the core types ('NetPeer', 'PeerId', etc.), +-- named constants, and small helpers used by all Peer sub-modules. +-- It is not exposed to library consumers; import 'GBNet.Peer' instead. +module GBNet.Peer.Internal + ( -- * Peer identifier + PeerId (..), + peerIdFromAddr, + + -- * Connection direction + ConnectionDirection (..), + + -- * Events + PeerEvent (..), + + -- * Pure processing types + IncomingPacket (..), + RawPacket (..), + PeerResult (..), + + -- * Pending connection + PendingConnection (..), + + -- * Net peer + NetPeer (..), + newPeerState, + + -- * Constants + cookieSecretSize, + fragmentTimeoutMs, + fragmentMaxBufferSize, + + -- * Helpers + withConnection, + queueRawPacket, + drainPeerSendQueue, + queueControlPacket, + cleanupPeer, + removePending, + generateCookieSecret, + ) +where + +import qualified Data.ByteString as BS +import Data.Foldable (toList) +import Data.Map.Strict (Map) +import qualified Data.Map.Strict as Map +import Data.Sequence (Seq) +import qualified Data.Sequence as Seq +import Data.Word (Word64, Word8) +import GBNet.Class (MonoTime (..)) +import GBNet.Config (NetworkConfig (..)) +import GBNet.Connection + ( Connection, + DisconnectReason (..), + ) +import GBNet.Fragment (FragmentAssembler) +import GBNet.Packet + ( Packet (..), + PacketHeader (..), + PacketType, + serializePacket, + ) +import GBNet.Security (RateLimiter, appendCrc32, newRateLimiter) +import GBNet.Socket (UdpSocket) +import GBNet.Types (ChannelId (..)) +import GBNet.Util (nextRandom) +import Network.Socket (SockAddr) +import Optics ((%~), (&), (.~)) +import Optics.TH (makeFieldLabelsNoPrefix) + +-- | Peer identifier wrapping a socket address. +newtype PeerId = PeerId {unPeerId :: SockAddr} + deriving (Eq, Ord, Show) + +-- | Create a PeerId from a SockAddr. +peerIdFromAddr :: SockAddr -> PeerId +peerIdFromAddr = PeerId + +-- | Direction of connection establishment. +data ConnectionDirection + = -- | They connected to us + Inbound + | -- | We connected to them + Outbound + deriving (Eq, Show) + +-- | Events emitted by peer processing. +data PeerEvent + = -- | A peer connected + PeerConnected !PeerId !ConnectionDirection + | -- | A peer disconnected + PeerDisconnected !PeerId !DisconnectReason + | -- | Received a message from a peer + PeerMessage !PeerId !ChannelId !BS.ByteString + | -- | A peer's address changed (connection migration) + PeerMigrated !PeerId !PeerId -- old, new + deriving (Eq, Show) + +-- | An incoming packet from the network (CRC already validated and stripped). +data IncomingPacket = IncomingPacket + { ipFrom :: !PeerId, + -- | CRC-validated payload + ipData :: !BS.ByteString + } + deriving (Eq, Show) + +makeFieldLabelsNoPrefix ''IncomingPacket + +-- | An outgoing packet ready to send (with CRC appended). +data RawPacket = RawPacket + { rpTo :: !PeerId, + -- | Data with CRC appended + rpData :: !BS.ByteString + } + deriving (Eq, Show) + +makeFieldLabelsNoPrefix ''RawPacket + +-- | State of a pending connection (mid-handshake). +data PendingConnection = PendingConnection + { pcDirection :: !ConnectionDirection, + pcServerSalt :: !Word64, + pcClientSalt :: !Word64, + pcCreatedAt :: !MonoTime, + pcRetryCount :: !Int, + pcLastRetry :: !MonoTime + } + deriving (Show) + +makeFieldLabelsNoPrefix ''PendingConnection + +-- | Cookie secret size in bytes. +cookieSecretSize :: Int +cookieSecretSize = 32 + +-- | Fragment reassembly timeout in milliseconds. +fragmentTimeoutMs :: Double +fragmentTimeoutMs = 5000.0 + +-- | Fragment reassembly max buffer size. +fragmentMaxBufferSize :: Int +fragmentMaxBufferSize = 1024 * 1024 -- 1MB + +-- | A network peer that can accept and initiate connections. +data NetPeer = NetPeer + { npSocket :: !UdpSocket, + -- | Local address this peer is bound to (for polymorphic operations) + npLocalAddr :: !SockAddr, + npConnections :: !(Map PeerId Connection), + npPending :: !(Map PeerId PendingConnection), + npConfig :: !NetworkConfig, + npRateLimiter :: !RateLimiter, + npCookieSecret :: !BS.ByteString, + npRngState :: !Word64, + npFragmentAssemblers :: !(Map PeerId FragmentAssembler), + -- | Tracks last migration time per connection to rate-limit migrations + npMigrationCooldowns :: !(Map Word64 MonoTime), + -- | Queued outgoing control packets (for pure API) + npSendQueue :: !(Seq RawPacket), + -- | Number of packets dropped by rate limiting + npRateLimitDrops :: !Word64 + } + +makeFieldLabelsNoPrefix ''NetPeer + +-- | Result of pure packet processing. +data PeerResult = PeerResult + { -- | Updated peer state + prPeer :: !NetPeer, + -- | Events that occurred during processing + prEvents :: ![PeerEvent], + -- | Packets to send (call @peerSendAllM@ with these) + prOutgoing :: ![RawPacket] + } + +makeFieldLabelsNoPrefix ''PeerResult + +-- | Create peer state (pure). Used internally and for polymorphic backends. +newPeerState :: + UdpSocket -> + SockAddr -> + NetworkConfig -> + MonoTime -> + NetPeer +newPeerState sock localAddr config now = + let rng0 = unMonoTime now + (secret, rng1) = generateCookieSecret rng0 + in NetPeer + { npSocket = sock, + npLocalAddr = localAddr, + npConnections = Map.empty, + npPending = Map.empty, + npConfig = config, + npRateLimiter = newRateLimiter (ncRateLimitPerSecond config) now, + npCookieSecret = secret, + npRngState = rng1, + npFragmentAssemblers = Map.empty, + npMigrationCooldowns = Map.empty, + npSendQueue = Seq.empty, + npRateLimitDrops = 0 + } + +-- | Generate a pseudo-random cookie secret. +generateCookieSecret :: Word64 -> (BS.ByteString, Word64) +generateCookieSecret seed = go seed cookieSecretSize [] + where + go s 0 acc = (BS.pack (reverse acc), s) + go s n acc = + let (r, s') = nextRandom s + in go s' (n - 1) (fromIntegral @Word64 @Word8 r : acc) + +-- ----------------------------------------------------------------------------- +-- Helper functions for pure API +-- ----------------------------------------------------------------------------- + +-- | Update a connection by PeerId, returning unchanged peer if not found. +withConnection :: PeerId -> (Connection -> Connection) -> NetPeer -> NetPeer +withConnection pid f peer = peer & #npConnections %~ Map.adjust f pid +{-# INLINE withConnection #-} + +-- | Queue a raw packet for sending. +queueRawPacket :: RawPacket -> NetPeer -> NetPeer +queueRawPacket pkt peer = peer & #npSendQueue %~ (Seq.|> pkt) + +-- | Drain the peer's send queue. +drainPeerSendQueue :: NetPeer -> ([RawPacket], NetPeer) +drainPeerSendQueue peer = + (toList (npSendQueue peer), peer & #npSendQueue .~ Seq.empty) + +-- | Serialize and queue a control packet. +queueControlPacket :: PacketType -> BS.ByteString -> PeerId -> NetPeer -> NetPeer +queueControlPacket ptype payload pid peer = + let header = + PacketHeader + { packetType = ptype, + sequenceNum = 0, + ack = 0, + ackBitfield = 0 + } + pkt = Packet {pktHeader = header, pktPayload = payload} + raw = appendCrc32 (serializePacket pkt) + in queueRawPacket (RawPacket pid raw) peer + +-- | Clean up per-peer state (fragment assemblers, migration cooldowns). +cleanupPeer :: PeerId -> NetPeer -> NetPeer +cleanupPeer peerId peer = + peer & #npFragmentAssemblers %~ Map.delete peerId + +-- | Remove a peer from pending. +removePending :: PeerId -> NetPeer -> NetPeer +removePending peerId peer = peer & #npPending %~ Map.delete peerId diff --git a/src/GBNet/Peer/Migration.hs b/src/GBNet/Peer/Migration.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Peer/Migration.hs @@ -0,0 +1,49 @@ +{-# LANGUAGE DataKinds #-} + +-- | +-- Module : GBNet.Peer.Migration +-- Description : Pure migration detection helpers +-- +-- Finds migration candidates by matching incoming packet sequence numbers +-- against existing connections. Only depends on 'GBNet.Peer.Internal'. +module GBNet.Peer.Migration + ( findMigrationCandidate, + migrationTokenFor, + migrationCooldownMs, + ) +where + +import qualified Data.Map.Strict as Map +import Data.Word (Word64) +import GBNet.Class (MonoTime) +import GBNet.Config (NetworkConfig (..)) +import GBNet.Connection (Connection) +import qualified GBNet.Connection as Conn +import GBNet.Packet (Packet (..), PacketHeader (..)) +import GBNet.Peer.Internal (NetPeer (..), PeerId) + +-- | Migration cooldown in milliseconds. +migrationCooldownMs :: Double +migrationCooldownMs = 5000.0 + +-- | Find a connection that could match this packet (for migration). +-- Returns (oldPeerId, connection, migrationToken) if found. +findMigrationCandidate :: Packet -> MonoTime -> NetPeer -> Maybe (PeerId, Connection, Word64) +findMigrationCandidate pkt _now peer = + findCandidate (Map.toList (npConnections peer)) + where + incomingSeq = sequenceNum (pktHeader pkt) + maxDistance = ncMaxSequenceDistance (npConfig peer) + + findCandidate [] = Nothing + findCandidate ((pid, conn) : rest) = + let remoteSeq = Conn.connRemoteSeq conn + diff = abs (fromIntegral incomingSeq - fromIntegral remoteSeq :: Int) + in if diff <= fromIntegral maxDistance + then Just (pid, conn, migrationTokenFor conn) + else findCandidate rest + +-- | Generate a migration token for a connection. +-- This should be stable across the connection lifetime. +migrationTokenFor :: Connection -> Word64 +migrationTokenFor = Conn.connClientSalt diff --git a/src/GBNet/Peer/Protocol.hs b/src/GBNet/Peer/Protocol.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Peer/Protocol.hs @@ -0,0 +1,153 @@ +-- | +-- Module : GBNet.Peer.Protocol +-- Description : Pure encoding/decoding utilities for Peer protocol +-- +-- Salt encoding, deny reasons, payload header decoding, channel sequence +-- decoding, and FNV-1a hashing. No dependencies on other Peer sub-modules. +module GBNet.Peer.Protocol + ( -- * Salt encoding + encodeSalt, + decodeSalt, + + -- * Deny reasons + DenyReason (..), + encodeDenyReason, + decodeDenyReason, + denyToDisconnectReason, + + -- * Payload header + decodePayloadHeader, + payloadFragmentFlag, + payloadChannelMask, + minPayloadSize, + + -- * Channel sequence + decodeChannelSeq, + + -- * FNV-1a hashing + sockAddrToKey, + fnvMix, + fnvOffsetBasis, + fnvPrime, + ) +where + +import Data.Bits (shiftL, xor, (.&.), (.|.)) +import qualified Data.ByteString as BS +import Data.List (foldl') +import Data.Word (Word64, Word8) +import GBNet.Connection (DisconnectReason (..)) +import GBNet.Serialize (deserialize, serialize) +import GBNet.Types (ChannelId (..), SequenceNum (..)) +import Network.Socket (PortNumber, SockAddr (..)) + +-- | Encode a Word64 salt to bytes. +encodeSalt :: Word64 -> BS.ByteString +encodeSalt = serialize +{-# INLINE encodeSalt #-} + +-- | Decode a Word64 salt from bytes. +decodeSalt :: BS.ByteString -> Maybe Word64 +decodeSalt bs = case deserialize bs of + Left _ -> Nothing + Right v -> Just v +{-# INLINE decodeSalt #-} + +-- | Deny reason codes sent during handshake rejection. +data DenyReason + = DenyServerFull + | DenyInvalidChallenge + | DenyUnknown !Word8 + deriving (Eq, Show) + +-- | Encode a deny reason to bytes. +encodeDenyReason :: DenyReason -> BS.ByteString +encodeDenyReason DenyServerFull = BS.singleton 1 +encodeDenyReason DenyInvalidChallenge = BS.singleton 2 +encodeDenyReason (DenyUnknown code) = BS.singleton code + +-- | Decode a deny reason from bytes. +decodeDenyReason :: BS.ByteString -> DenyReason +decodeDenyReason bs = case BS.uncons bs of + Just (1, _) -> DenyServerFull + Just (2, _) -> DenyInvalidChallenge + Just (code, _) -> DenyUnknown code + Nothing -> DenyUnknown 0 + +-- | Convert a deny reason to a disconnect reason for event reporting. +denyToDisconnectReason :: DenyReason -> DisconnectReason +denyToDisconnectReason DenyServerFull = ReasonServerFull +denyToDisconnectReason DenyInvalidChallenge = ReasonProtocolMismatch +denyToDisconnectReason (DenyUnknown code) = ReasonUnknown code + +-- | Payload header: channel (3 bits) + is_fragment (1 bit) + reserved (4 bits) +-- Encoded in first byte: [is_fragment:1][reserved:4][channel:3] +-- Channel is low 3 bits, is_fragment is high bit (0x80) + +-- | Fragment flag bit in payload header. +payloadFragmentFlag :: Word8 +payloadFragmentFlag = 0x80 + +-- | Channel mask in payload header. +payloadChannelMask :: Word8 +payloadChannelMask = 0x07 + +-- | Decode payload header byte. +decodePayloadHeader :: Word8 -> (ChannelId, Bool) +decodePayloadHeader b = + let channel = ChannelId (b .&. payloadChannelMask) + isFragment = (b .&. payloadFragmentFlag) /= 0 + in (channel, isFragment) +{-# INLINE decodePayloadHeader #-} + +-- | Minimum size of a non-fragment payload: headerByte + 2 bytes channel sequence. +minPayloadSize :: Int +minPayloadSize = 3 + +-- | Decode channel sequence from payload bytes. +-- Returns (channelSeq, remaining data) or Nothing if too short. +decodeChannelSeq :: BS.ByteString -> Maybe (SequenceNum, BS.ByteString) +decodeChannelSeq bs + | BS.length bs < 2 = Nothing + | otherwise = + let chSeq = + SequenceNum $ + (fromIntegral (BS.index bs 0) `shiftL` 8) + .|. fromIntegral (BS.index bs 1) + in Just (chSeq, BS.drop 2 bs) +{-# INLINE decodeChannelSeq #-} + +-- | FNV-1a hash seed. +fnvOffsetBasis :: Word64 +fnvOffsetBasis = 14695981039346656037 + +-- | FNV-1a hash prime. +fnvPrime :: Word64 +fnvPrime = 1099511628211 + +-- | Convert SockAddr to a key for rate limiting using FNV-1a hash. +-- Hashes the raw address/port words directly, avoiding 'show'. +sockAddrToKey :: SockAddr -> Word64 +sockAddrToKey addr = case addr of + SockAddrInet port host -> + fnvMix (fnvMix fnvOffsetBasis (fromIntegral host)) (fromIntegral (portToWord port)) + SockAddrInet6 port _ (h1, h2, h3, h4) _ -> + foldl' + fnvMix + fnvOffsetBasis + [ fromIntegral (portToWord port), + fromIntegral h1, + fromIntegral h2, + fromIntegral h3, + fromIntegral h4 + ] + _ -> fnvMix fnvOffsetBasis 0 + where + portToWord :: PortNumber -> Word64 + portToWord = fromIntegral + {-# INLINE portToWord #-} + +-- | FNV-1a mix step: XOR then multiply. +fnvMix :: Word64 -> Word64 -> Word64 +fnvMix h val = (h `xor` val) * fnvPrime +{-# INLINE fnvMix #-} diff --git a/src/GBNet/Reliability.hs b/src/GBNet/Reliability.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Reliability.hs @@ -0,0 +1,775 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedLabels #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- Module : GBNet.Reliability +-- Description : Reliable packet delivery with RTT estimation and ACK tracking +-- +-- Jacobson\/Karels RTT estimation, adaptive RTO, fast retransmit on 3 NACKs, +-- bounded in-flight tracking, and rolling packet loss window. +module GBNet.Reliability + ( -- * Constants + initialRtoMillis, + ackBitsWindow, + rttAlpha, + rttBeta, + minRtoMs, + maxRtoMs, + lossWindowSize, + fastRetransmitThreshold, + defaultMaxSequenceDistance, + defaultMaxInFlight, + + -- * Monotonic time + MonoTime, + elapsedMs, + + -- * Loss window + LossWindow (..), + emptyLossWindow, + + -- * Sequence buffer + SBEntry (..), + SequenceBuffer, + newSequenceBuffer, + sbInsert, + sbInsertMany, + sbExists, + sbGet, + + -- * Received packet buffer (opaque) + ReceivedBuffer, + + -- * Sent packet buffer (opaque) + RingEntry, + SentPacketBuffer, + + -- * Sent packet record + SentPacketRecord (..), + + -- * Reliable endpoint + ReliableEndpoint (..), + AckResult (..), + newReliableEndpoint, + withMaxInFlight, + nextSequence, + onPacketSent, + onPacketsReceived, + processAcks, + getAckInfo, + updateRtt, + recordLossSample, + packetLossPercent, + isInFlight, + packetsInFlight, + rtoMs, + srttMs, + ) +where + +import Control.DeepSeq (NFData (..)) +import Data.Bits (complement, popCount, shiftL, (.&.), (.|.)) +import Data.List (foldl') +import qualified Data.Vector as V +import qualified Data.Vector.Mutable as MV +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM +import Data.Word (Word16, Word32, Word64, Word8) +import GBNet.Class (MonoTime (..)) +import GBNet.Types (ChannelId (..), SequenceNum (..)) +import GBNet.Util (sequenceDiff, sequenceGreaterThan) +import GBNet.ZeroCopy (zeroCopyMutate, zeroCopyMutate', zeroCopyMutateU') +import Optics ((%~), (&), (.~)) +import Optics.TH (makeFieldLabelsNoPrefix) + +-- Constants + +-- | Initial retransmission timeout in milliseconds (before any RTT samples). +initialRtoMillis :: Double +initialRtoMillis = 100.0 + +-- | Number of ACK bits tracked per packet (64-bit bitfield). +ackBitsWindow :: Word16 +ackBitsWindow = 64 + +-- | Jacobson\/Karels SRTT smoothing factor (1\/8). +rttAlpha :: Double +rttAlpha = 0.125 + +-- | Jacobson\/Karels RTTVAR smoothing factor (1\/4). +rttBeta :: Double +rttBeta = 0.25 + +-- | Minimum retransmission timeout in milliseconds. +minRtoMs :: Double +minRtoMs = 50.0 + +-- | Maximum retransmission timeout in milliseconds. +maxRtoMs :: Double +maxRtoMs = 2000.0 + +-- | Rolling loss window size (number of samples tracked). +lossWindowSize :: Int +lossWindowSize = 256 + +-- | Number of NACKs before triggering fast retransmit. +fastRetransmitThreshold :: Word8 +fastRetransmitThreshold = 3 + +-- | Default maximum sequence distance before treating packets as stale. +defaultMaxSequenceDistance :: Word16 +defaultMaxSequenceDistance = 32768 + +-- | Default maximum packets in flight. +defaultMaxInFlight :: Int +defaultMaxInFlight = 256 + +-- | 256-bit loss window using 4 Word64 values. +-- Zero-allocation updates via bit operations. +data LossWindow = LossWindow + { lwBits0 :: {-# UNPACK #-} !Word64, -- bits 0-63 + lwBits1 :: {-# UNPACK #-} !Word64, -- bits 64-127 + lwBits2 :: {-# UNPACK #-} !Word64, -- bits 128-191 + lwBits3 :: {-# UNPACK #-} !Word64 -- bits 192-255 + } + deriving (Show, Eq) + +makeFieldLabelsNoPrefix ''LossWindow + +instance NFData LossWindow where + rnf (LossWindow b0 b1 b2 b3) = rnf b0 `seq` rnf b1 `seq` rnf b2 `seq` rnf b3 + +-- | Empty loss window (all successful). +emptyLossWindow :: LossWindow +emptyLossWindow = LossWindow 0 0 0 0 + +-- | Set a bit in the loss window. Index must be 0-255. +lwSetBit :: Int -> Bool -> LossWindow -> LossWindow +lwSetBit idx val lw + | idx < 64 = lw & #lwBits0 %~ setBitWord64 (idx .&. 63) val + | idx < 128 = lw & #lwBits1 %~ setBitWord64 (idx .&. 63) val + | idx < 192 = lw & #lwBits2 %~ setBitWord64 (idx .&. 63) val + | otherwise = lw & #lwBits3 %~ setBitWord64 (idx .&. 63) val +{-# INLINE lwSetBit #-} + +-- | Count set bits (losses) in the first n entries. +lwCountLosses :: Int -> LossWindow -> Int +lwCountLosses n lw + | n <= 0 = 0 + | n >= lossWindowSize = totalBits + | n <= 64 = popCount (lwBits0 lw .&. maskN n) + | n <= 128 = popCount (lwBits0 lw) + popCount (lwBits1 lw .&. maskN (n - 64)) + | n <= 192 = popCount (lwBits0 lw) + popCount (lwBits1 lw) + popCount (lwBits2 lw .&. maskN (n - 128)) + | otherwise = popCount (lwBits0 lw) + popCount (lwBits1 lw) + popCount (lwBits2 lw) + popCount (lwBits3 lw .&. maskN (n - 192)) + where + totalBits = popCount (lwBits0 lw) + popCount (lwBits1 lw) + popCount (lwBits2 lw) + popCount (lwBits3 lw) + maskN k = (1 `shiftL` k) - 1 +{-# INLINE lwCountLosses #-} + +-- | Set or clear a bit in a Word64. +setBitWord64 :: Int -> Bool -> Word64 -> Word64 +setBitWord64 idx True w = w .|. (1 `shiftL` idx) +setBitWord64 idx False w = w .&. complement (1 `shiftL` idx) +{-# INLINE setBitWord64 #-} + +-- Monotonic time + +-- | Elapsed time in milliseconds. +elapsedMs :: MonoTime -> MonoTime -> Double +elapsedMs start now = fromIntegral (unMonoTime (now - start)) / 1e6 +{-# INLINE elapsedMs #-} + +-- SequenceBuffer + +-- | Ring buffer entry. Stores sequence number for validation. +data SBEntry a + = SBEmpty + | SBEntry !SequenceNum a + deriving (Show) + +instance (NFData a) => NFData (SBEntry a) where + rnf SBEmpty = () + rnf (SBEntry s a) = rnf s `seq` rnf a + +-- | High-performance circular buffer using Vector ring buffer. +-- O(1) insert, O(1) lookup, O(1) exists via zero-copy mutation. +-- Size must be power of 2 for fast bitwise modulo. +data SequenceBuffer a = SequenceBuffer + { sbEntries :: !(V.Vector (SBEntry a)), + sbSequence :: !SequenceNum, + sbSize :: !Int -- Keep for API compat, but use mask internally + } + deriving (Show) + +instance (NFData a) => NFData (SequenceBuffer a) where + rnf (SequenceBuffer e s sz) = rnf e `seq` rnf s `seq` rnf sz + +makeFieldLabelsNoPrefix ''SequenceBuffer + +-- | Create a new sequence buffer. Size is rounded up to power of 2. +newSequenceBuffer :: Int -> SequenceBuffer a +newSequenceBuffer requestedSize = + let size = nextPowerOf2 requestedSize + in SequenceBuffer + { sbEntries = V.replicate size SBEmpty, + sbSequence = 0, + sbSize = size + } + +-- | Round up to next power of 2. +nextPowerOf2 :: Int -> Int +nextPowerOf2 n + | n <= 1 = 1 + | otherwise = go 1 + where + go p + | p >= n = p + | otherwise = go (p * 2) +{-# INLINE nextPowerOf2 #-} + +-- | O(1) insert using zero-copy mutation. +sbInsert :: SequenceNum -> a -> SequenceBuffer a -> SequenceBuffer a +sbInsert seqNum val buf = + let !idx = seqToIndex seqNum buf + newHighest + | sequenceGreaterThan seqNum (sbSequence buf) = seqNum + | otherwise = sbSequence buf + entries' = zeroCopyMutate (sbEntries buf) $ \mv -> + MV.unsafeWrite mv idx (SBEntry seqNum val) + in buf & #sbEntries .~ entries' & #sbSequence .~ newHighest +{-# INLINE sbInsert #-} + +-- | Batched insert - O(n) with single thaw/freeze for n entries. +sbInsertMany :: [(SequenceNum, a)] -> SequenceBuffer a -> SequenceBuffer a +sbInsertMany [] buf = buf +sbInsertMany items buf = + let (entries', newHighest) = zeroCopyMutate' (sbEntries buf) $ \mv -> + let go _ !highest [] = return highest + go mv' !highest ((seqNum, val) : rest) = do + let !idx = seqToIndex seqNum buf + MV.unsafeWrite mv' idx (SBEntry seqNum val) + let !highest' = + if sequenceGreaterThan seqNum highest then seqNum else highest + go mv' highest' rest + in go mv (sbSequence buf) items + in buf & #sbEntries .~ entries' & #sbSequence .~ newHighest +{-# INLINE sbInsertMany #-} + +-- | O(1) existence check via direct index. +sbExists :: SequenceNum -> SequenceBuffer a -> Bool +sbExists seqNum buf = + case V.unsafeIndex (sbEntries buf) (seqToIndex seqNum buf) of + SBEntry storedSeq _ -> storedSeq == seqNum + SBEmpty -> False +{-# INLINE sbExists #-} + +-- | O(1) lookup via direct index. +sbGet :: SequenceNum -> SequenceBuffer a -> Maybe a +sbGet seqNum buf = + case V.unsafeIndex (sbEntries buf) (seqToIndex seqNum buf) of + SBEntry storedSeq v + | storedSeq == seqNum -> Just v + _ -> Nothing +{-# INLINE sbGet #-} + +-- | Fast index using bitwise AND (size must be power of 2). +seqToIndex :: SequenceNum -> SequenceBuffer a -> Int +seqToIndex (SequenceNum s) buf = fromIntegral s .&. (sbSize buf - 1) +{-# INLINE seqToIndex #-} + +-- ReceivedBuffer (optimized for packet deduplication) + +-- | Buffer size for received packets (must be power of 2 for fast modulo). +receivedBufferSize :: Int +receivedBufferSize = 256 +{-# INLINE receivedBufferSize #-} + +-- | Mask for fast modulo (size - 1). +receivedBufferMask :: Int +receivedBufferMask = receivedBufferSize - 1 +{-# INLINE receivedBufferMask #-} + +-- | High-performance received packet tracker using unboxed Word16 vector. +-- Stores sequence numbers directly - O(1) lookup via index, O(1) insert. +-- 512 bytes total (256 x Word16), cache-friendly contiguous memory. +data ReceivedBuffer = ReceivedBuffer + { rbSeqs :: !(VU.Vector Word16), -- Unboxed: contiguous like C array + rbHighest :: !SequenceNum -- Track highest for ack bit calculation + } + deriving (Show) + +makeFieldLabelsNoPrefix ''ReceivedBuffer + +instance NFData ReceivedBuffer where + rnf (ReceivedBuffer v h) = rnf v `seq` rnf h + +-- | Create empty received buffer (all slots set to maxBound as "empty" sentinel). +newReceivedBuffer :: ReceivedBuffer +newReceivedBuffer = + ReceivedBuffer + { rbSeqs = VU.replicate receivedBufferSize maxBound, + rbHighest = 0 + } + +-- | O(1) check if sequence number was received. Single array index + comparison. +rbExists :: SequenceNum -> ReceivedBuffer -> Bool +rbExists (SequenceNum s) (ReceivedBuffer v _) = + VU.unsafeIndex v (fromIntegral s .&. receivedBufferMask) == s +{-# INLINE rbExists #-} + +-- | Batched insert - O(n) for n sequence numbers via zero-copy mutation. +rbInsertMany :: [SequenceNum] -> ReceivedBuffer -> ReceivedBuffer +rbInsertMany [] buf = buf +rbInsertMany seqs buf = + let (v', newHighest) = zeroCopyMutateU' (rbSeqs buf) $ \mv -> + go mv (rbHighest buf) seqs + in ReceivedBuffer v' newHighest + where + go _ !highest [] = return highest + go mv !highest (sn@(SequenceNum s) : rest) = do + VUM.unsafeWrite mv (fromIntegral s .&. receivedBufferMask) s + let highest' = if sequenceGreaterThan sn highest then sn else highest + go mv highest' rest +{-# INLINE rbInsertMany #-} + +-- SentPacketRecord + +-- | Record of a sent packet, tracked for ACK processing and RTT estimation. +data SentPacketRecord = SentPacketRecord + { sprChannelId :: !ChannelId, + sprChannelSequence :: !SequenceNum, + sprSendTime :: !MonoTime, + sprSize :: !Int, + sprNackCount :: !Word8 + } + deriving (Show) + +makeFieldLabelsNoPrefix ''SentPacketRecord + +instance NFData SentPacketRecord where + rnf (SentPacketRecord cid cseq t s n) = + rnf cid `seq` rnf cseq `seq` rnf t `seq` rnf s `seq` rnf n + +-- | Ring buffer entry: stores sequence number for validation + record. +-- Nothing = slot empty, Just = occupied with given sequence. +data RingEntry = RingEmpty | RingEntry !SequenceNum !SentPacketRecord + deriving (Show) + +instance NFData RingEntry where + rnf RingEmpty = () + rnf (RingEntry s r) = rnf s `seq` rnf r + +-- | O(1) sent packet buffer using ring buffer indexed by sequence mod size. +-- All operations are O(1) except count and findOldest which are O(n). +data SentPacketBuffer = SentPacketBuffer + { spbEntries :: !(V.Vector RingEntry), + spbCount :: !Int -- Cached count for O(1) access + } + deriving (Show) + +makeFieldLabelsNoPrefix ''SentPacketBuffer + +instance NFData SentPacketBuffer where + rnf (SentPacketBuffer e c) = rnf e `seq` rnf c + +-- | Buffer size (fixed at 256 for Word16 sequence space efficiency). +ringBufferSize :: Int +ringBufferSize = 256 +{-# INLINE ringBufferSize #-} + +-- | Create empty sent packet buffer. +newSentPacketBuffer :: Int -> SentPacketBuffer +newSentPacketBuffer _ = + SentPacketBuffer + { spbEntries = V.replicate ringBufferSize RingEmpty, + spbCount = 0 + } + +-- | Convert sequence to ring index. O(1). +seqToIdx :: SequenceNum -> Int +seqToIdx (SequenceNum s) = fromIntegral s .&. (ringBufferSize - 1) +{-# INLINE seqToIdx #-} + +-- | O(1) lookup by sequence number. +spbLookup :: SequenceNum -> SentPacketBuffer -> Maybe SentPacketRecord +spbLookup seqNum buf = + case V.unsafeIndex (spbEntries buf) (seqToIdx seqNum) of + RingEntry storedSeq record + | storedSeq == seqNum -> Just record + _ -> Nothing +{-# INLINE spbLookup #-} + +-- | O(1) membership test. +spbMember :: SequenceNum -> SentPacketBuffer -> Bool +spbMember seqNum buf = + case V.unsafeIndex (spbEntries buf) (seqToIdx seqNum) of + RingEntry storedSeq _ + | storedSeq == seqNum -> True + _ -> False +{-# INLINE spbMember #-} + +-- | O(1) insert using zero-copy mutation. +spbInsert :: SequenceNum -> SentPacketRecord -> SentPacketBuffer -> SentPacketBuffer +spbInsert seqNum record buf = + let idx = seqToIdx seqNum + (entries', countDelta) = zeroCopyMutate' (spbEntries buf) $ \mvec -> do + old <- MV.unsafeRead mvec idx + MV.unsafeWrite mvec idx (RingEntry seqNum record) + return $! case old of + RingEmpty -> 1 :: Int + _ -> 0 + in buf & #spbEntries .~ entries' & #spbCount %~ (+ countDelta) +{-# INLINE spbInsert #-} + +-- | O(1) delete by sequence number. +spbDelete :: SequenceNum -> SentPacketBuffer -> SentPacketBuffer +spbDelete seqNum buf = + case V.unsafeIndex (spbEntries buf) idx of + RingEntry storedSeq _ + | storedSeq == seqNum -> + let entries' = zeroCopyMutate (spbEntries buf) $ \mvec -> + MV.unsafeWrite mvec idx RingEmpty + in buf & #spbEntries .~ entries' & #spbCount %~ subtract 1 + _ -> buf + where + idx = seqToIdx seqNum +{-# INLINE spbDelete #-} + +-- | O(n) find oldest entry (for eviction). Only called when buffer full. +spbFindOldest :: SentPacketBuffer -> Maybe (SequenceNum, SentPacketRecord) +spbFindOldest buf = V.foldl' older Nothing (spbEntries buf) + where + older Nothing (RingEntry s r) = Just (s, r) + older Nothing RingEmpty = Nothing + older acc@(Just (_, r1)) (RingEntry s r2) + | sprSendTime r2 < sprSendTime r1 = Just (s, r2) + | otherwise = acc + older acc RingEmpty = acc + +-- ReliableEndpoint + +-- | Result of processing ACKs. +data AckResult = AckResult + { arAcked :: ![(ChannelId, SequenceNum)], + arFastRetransmit :: ![(ChannelId, SequenceNum)] + } + deriving (Show, Eq) + +makeFieldLabelsNoPrefix ''AckResult + +instance NFData AckResult where + rnf (AckResult a f) = rnf a `seq` rnf f + +-- | Core reliability state: sequence tracking, RTT estimation, loss detection, +-- and in-flight packet management. +data ReliableEndpoint = ReliableEndpoint + { reLocalSequence :: !SequenceNum, + reRemoteSequence :: !SequenceNum, + reAckBits :: !Word64, + reSentPackets :: !SentPacketBuffer, + reReceivedPackets :: !ReceivedBuffer, + reMaxSequenceDistance :: !Word16, + reMaxInFlight :: !Int, + reSrtt :: !Double, + reRttvar :: !Double, + reRto :: !Double, + reHasRttSample :: !Bool, + reLossWindow :: !LossWindow, + reLossWindowIndex :: !Int, + reLossWindowCount :: !Int, + reTotalSent :: !Word64, + reTotalAcked :: !Word64, + reTotalLost :: !Word64, + rePacketsEvicted :: !Word64, + reBytesSent :: !Word64, + reBytesAcked :: !Word64 + } + deriving (Show) + +instance NFData ReliableEndpoint where + rnf (ReliableEndpoint ls rs ab sp rp msd mif srtt rv rto hrs lw lwi lwc ts ta tl pe bs ba) = + rnf ls `seq` + rnf rs `seq` + rnf ab `seq` + rnf sp `seq` + rnf rp `seq` + rnf msd `seq` + rnf mif `seq` + rnf srtt `seq` + rnf rv `seq` + rnf rto `seq` + rnf hrs `seq` + rnf lw `seq` + rnf lwi `seq` + rnf lwc `seq` + rnf ts `seq` + rnf ta `seq` + rnf tl `seq` + rnf pe `seq` + rnf bs `seq` + rnf ba + +makeFieldLabelsNoPrefix ''ReliableEndpoint + +-- | Create a new reliable endpoint with the given buffer size. +newReliableEndpoint :: Int -> ReliableEndpoint +newReliableEndpoint bufferSize = + ReliableEndpoint + { reLocalSequence = 0, + reRemoteSequence = 0, + reAckBits = 0, + reSentPackets = newSentPacketBuffer bufferSize, + reReceivedPackets = newReceivedBuffer, + reMaxSequenceDistance = defaultMaxSequenceDistance, + reMaxInFlight = defaultMaxInFlight, + reSrtt = 0.0, + reRttvar = 0.0, + reRto = initialRtoMillis, + reHasRttSample = False, + reLossWindow = emptyLossWindow, + reLossWindowIndex = 0, + reLossWindowCount = 0, + reTotalSent = 0, + reTotalAcked = 0, + reTotalLost = 0, + rePacketsEvicted = 0, + reBytesSent = 0, + reBytesAcked = 0 + } + +-- | Override the maximum in-flight packet count. +withMaxInFlight :: Int -> ReliableEndpoint -> ReliableEndpoint +withMaxInFlight maxFlight ep = ep & #reMaxInFlight .~ maxFlight + +-- | Allocate the next local sequence number. +nextSequence :: ReliableEndpoint -> (SequenceNum, ReliableEndpoint) +nextSequence ep = + let s = reLocalSequence ep + in (s, ep & #reLocalSequence .~ (s + 1)) + +-- | Record a sent packet for in-flight tracking and future ACK processing. +onPacketSent :: + SequenceNum -> + MonoTime -> + ChannelId -> + SequenceNum -> + Int -> + ReliableEndpoint -> + ReliableEndpoint +onPacketSent seqNum sendTime channelId channelSeq size ep = + let ep' = + if spbCount (reSentPackets ep) >= reMaxInFlight ep + then evictWorstInFlight ep + else ep + record = + SentPacketRecord + { sprChannelId = channelId, + sprChannelSequence = channelSeq, + sprSendTime = sendTime, + sprSize = size, + sprNackCount = 0 + } + in ep' + & #reSentPackets + %~ spbInsert seqNum record + & #reTotalSent + %~ (+ 1) + & #reBytesSent + %~ (+ fromIntegral size) + +evictWorstInFlight :: ReliableEndpoint -> ReliableEndpoint +evictWorstInFlight ep = + case spbFindOldest (reSentPackets ep) of + Nothing -> ep + Just (worstSeq, _) -> + ep + & #reSentPackets + %~ spbDelete worstSeq + & #rePacketsEvicted + %~ (+ 1) + +-- | Process received packets - processes multiple sequence numbers with ONE thaw/freeze. +-- This is the high-performance API for game loops processing many packets per frame. +onPacketsReceived :: [SequenceNum] -> ReliableEndpoint -> ReliableEndpoint +onPacketsReceived seqNums ep = + let validSeqs = filter isValid seqNums + newBuf = rbInsertMany validSeqs (reReceivedPackets ep) + (newRemote, newAckBits) = foldl' updateAckState (reRemoteSequence ep, reAckBits ep) validSeqs + in ep & #reReceivedPackets .~ newBuf & #reRemoteSequence .~ newRemote & #reAckBits .~ newAckBits + where + isValid sn = + let dist = fromIntegral (abs (sequenceDiff sn (reRemoteSequence ep))) :: Word32 + in dist <= fromIntegral (reMaxSequenceDistance ep) && not (rbExists sn (reReceivedPackets ep)) + + updateAckState (!remote, !bits) sn + | sequenceGreaterThan sn remote = + let d = fromIntegral (sequenceDiff sn remote) :: Word64 + newBits + | d < fromIntegral ackBitsWindow = (bits `shiftL` fromIntegral d) .|. (1 `shiftL` (fromIntegral d - 1)) + | otherwise = 0 + in (sn, newBits) + | otherwise = + let d = fromIntegral (sequenceDiff remote sn) :: Word64 + in if d > 0 && d <= fromIntegral ackBitsWindow + then (remote, bits .|. (1 `shiftL` (fromIntegral d - 1))) + else (remote, bits) + +-- | Mutation to apply to sent packet buffer. +data BufferMutation + = MutDelete !Int -- Delete at index + | MutNack !Int !SequenceNum !SentPacketRecord !Word8 -- idx, seqNum, record, newNack + +-- | Process ACKs from received packet. O(ackBitsWindow) = O(64). +-- Two-phase: pure lookup pass, then single batched mutation. +processAcks :: SequenceNum -> Word64 -> MonoTime -> ReliableEndpoint -> (AckResult, ReliableEndpoint) +processAcks ackSeq ackBitsVal now ep = + let buf = reSentPackets ep + -- Phase 1: Pure lookup, collect decisions + (acked, retrans, mutations, rttSum, rttCount, bytesAcked) = + foldl' (processOne buf) ([], [], [], 0.0, 0 :: Int, 0) [0 .. fromIntegral ackBitsWindow] + -- Phase 2: Apply all mutations in one ST block + buf' = applyMutations mutations buf + -- Update RTT with average of all samples + ep' = case rttCount of + 0 -> ep + _ -> updateRtt (rttSum / fromIntegral rttCount) ep + -- Record loss samples for each ack (success = not lost) + ep'' = foldl' (\e _ -> recordLossSample False e) ep' acked + ep''' = + ep'' + & #reSentPackets + .~ buf' + & #reTotalAcked + .~ (reTotalAcked ep + fromIntegral (length acked)) + & #reBytesAcked + .~ (reBytesAcked ep + fromIntegral bytesAcked) + in (AckResult acked retrans, ep''') + where + processOne buf (!acked, !retrans, !muts, !rttSum, !rttCnt, !bytes) i + | i == 0 = checkSeq ackSeq True buf acked retrans muts rttSum rttCnt bytes + | otherwise = + let seqToCheck = ackSeq - fromIntegral i + bitSet = (ackBitsVal .&. (1 `shiftL` (i - 1))) /= 0 + in checkSeq seqToCheck bitSet buf acked retrans muts rttSum rttCnt bytes + + checkSeq seqNum bitSet buf acked retrans muts rttSum rttCnt bytes = + case spbLookup seqNum buf of + Nothing -> (acked, retrans, muts, rttSum, rttCnt, bytes) + Just record + | bitSet -> + let idx = seqToIdx seqNum + pair = (sprChannelId record, sprChannelSequence record) + rtt = elapsedMs (sprSendTime record) now + in (pair : acked, retrans, MutDelete idx : muts, rttSum + rtt, rttCnt + 1, bytes + sprSize record) + | otherwise -> + let idx = seqToIdx seqNum + newNack = min 255 (sprNackCount record + 1) + retrans' = + if newNack == fastRetransmitThreshold + then (sprChannelId record, sprChannelSequence record) : retrans + else retrans + in (acked, retrans', MutNack idx seqNum record newNack : muts, rttSum, rttCnt, bytes) + +-- | Apply all mutations in one ST block via zero-copy mutation. +applyMutations :: [BufferMutation] -> SentPacketBuffer -> SentPacketBuffer +applyMutations [] buf = buf +applyMutations muts buf = + let (entries', deletions) = zeroCopyMutate' (spbEntries buf) $ \mvec -> + applyAll mvec muts 0 + in buf & #spbEntries .~ entries' & #spbCount %~ subtract deletions + where + applyAll _ [] !dels = return dels + applyAll mvec (MutDelete idx : rest) !dels = do + MV.unsafeWrite mvec idx RingEmpty + applyAll mvec rest (dels + 1) + applyAll mvec (MutNack idx seqNum record newNack : rest) !dels = do + MV.unsafeWrite mvec idx (RingEntry seqNum (record & #sprNackCount .~ newNack)) + applyAll mvec rest dels + +-- | Update SRTT and RTO using Jacobson\/Karels algorithm. +updateRtt :: Double -> ReliableEndpoint -> ReliableEndpoint +updateRtt sampleMs ep + | not (reHasRttSample ep) = + let newSrtt = sampleMs + newRttvar = sampleMs / 2.0 + newRto = clampRto (newSrtt + 4.0 * newRttvar) + in ep + & #reSrtt + .~ newSrtt + & #reRttvar + .~ newRttvar + & #reRto + .~ newRto + & #reHasRttSample + .~ True + | otherwise = + let newRttvar = (1.0 - rttBeta) * reRttvar ep + rttBeta * abs (sampleMs - reSrtt ep) + newSrtt = (1.0 - rttAlpha) * reSrtt ep + rttAlpha * sampleMs + newRto = clampRto (newSrtt + 4.0 * newRttvar) + in ep + & #reSrtt + .~ newSrtt + & #reRttvar + .~ newRttvar + & #reRto + .~ newRto + +clampRto :: Double -> Double +clampRto rto = max minRtoMs (min maxRtoMs rto) + +-- | Record a loss\/success sample in the rolling loss window. +recordLossSample :: Bool -> ReliableEndpoint -> ReliableEndpoint +recordLossSample lost ep = + let idx = reLossWindowIndex ep `mod` lossWindowSize + newWindow = lwSetBit idx lost (reLossWindow ep) + newCount = min lossWindowSize (reLossWindowCount ep + 1) + in ep + & #reLossWindow + .~ newWindow + & #reLossWindowIndex + %~ (+ 1) + & #reLossWindowCount + .~ newCount +{-# INLINE recordLossSample #-} + +-- | Get the current ACK sequence and bitfield for outgoing packet headers. +getAckInfo :: ReliableEndpoint -> (SequenceNum, Word64) +getAckInfo ep = (reRemoteSequence ep, reAckBits ep) +{-# INLINE getAckInfo #-} + +-- | Current retransmission timeout in milliseconds. +rtoMs :: ReliableEndpoint -> Double +rtoMs = reRto +{-# INLINE rtoMs #-} + +-- | Smoothed round-trip time in milliseconds. +srttMs :: ReliableEndpoint -> Double +srttMs = reSrtt +{-# INLINE srttMs #-} + +-- | Current packet loss as a fraction (0.0 to 1.0). +packetLossPercent :: ReliableEndpoint -> Double +packetLossPercent ep + | count == 0 = 0.0 + | otherwise = fromIntegral lost / fromIntegral count + where + count = reLossWindowCount ep + lost = lwCountLosses count (reLossWindow ep) +{-# INLINE packetLossPercent #-} + +-- | Check whether a sequence number is currently in flight. +isInFlight :: SequenceNum -> ReliableEndpoint -> Bool +isInFlight seqNum ep = spbMember seqNum (reSentPackets ep) +{-# INLINE isInFlight #-} + +-- | Number of packets currently in flight. +packetsInFlight :: ReliableEndpoint -> Int +packetsInFlight = spbCount . reSentPackets +{-# INLINE packetsInFlight #-} diff --git a/src/GBNet/Replication/Delta.hs b/src/GBNet/Replication/Delta.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Replication/Delta.hs @@ -0,0 +1,289 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedLabels #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- Module : GBNet.Replication.Delta +-- Description : Delta compression for bandwidth-efficient state replication +-- +-- 'DeltaTracker' manages per-sequence snapshots for delta encoding against +-- acknowledged baselines. 'BaselineManager' provides a ring buffer of +-- confirmed snapshots per connection. +-- +-- Delta compression only sends changed fields between states, dramatically +-- reducing bandwidth for entity replication. +module GBNet.Replication.Delta + ( -- * Types + BaselineSeq, + noBaseline, + + -- * NetworkDelta typeclass + NetworkDelta (..), + + -- * Delta tracker (sender side) + DeltaTracker, + newDeltaTracker, + deltaEncode, + deltaOnAck, + deltaReset, + deltaConfirmedSeq, + + -- * Baseline manager (receiver side) + BaselineManager, + newBaselineManager, + pushBaseline, + getBaseline, + baselineReset, + baselineCount, + baselineIsEmpty, + + -- * Decoding + deltaDecode, + ) +where + +import Control.DeepSeq (NFData (..)) +import qualified Data.ByteString as BS +import Data.Int (Int32) +import Data.Sequence (Seq) +import qualified Data.Sequence as Seq +import Data.Word (Word16) +import Foreign.Storable (Storable (..)) +import GBNet.Reliability (MonoTime, elapsedMs) +import GBNet.Serialize (deserialize, serialize) +import Optics ((&), (.~), (?~)) +import Optics.TH (makeFieldLabelsNoPrefix) + +-- | Sequence number used to identify baseline snapshots on the wire. +type BaselineSeq = Word16 + +-- | Sentinel value indicating no baseline is available (full state required). +noBaseline :: BaselineSeq +noBaseline = maxBound + +-- | Signed difference between two baseline sequence numbers with wraparound. +baselineSeqDiff :: BaselineSeq -> BaselineSeq -> Int32 +baselineSeqDiff s1 s2 = fromIntegral (s1 - s2 :: Word16) + +-- | Wire header size for baseline sequence (bytes). +baselineSeqSize :: Int +baselineSeqSize = 2 + +-- | Typeclass for delta-compressed network types. +-- +-- Implement this for types that benefit from delta compression. +-- The 'Delta' associated type contains only the changed fields. +-- +-- Example: +-- +-- @ +-- data PlayerState = PlayerState { pos :: Vec3, health :: Word8 } +-- data PlayerDelta = PlayerDelta { dPos :: Maybe Vec3, dHealth :: Maybe Word8 } +-- +-- instance NetworkDelta PlayerState where +-- type Delta PlayerState = PlayerDelta +-- diff current baseline = PlayerDelta +-- { dPos = if pos current /= pos baseline then Just (pos current) else Nothing +-- , dHealth = if health current /= health baseline then Just (health current) else Nothing +-- } +-- apply state delta = state +-- { pos = maybe (pos state) id (dPos delta) +-- , health = maybe (health state) id (dHealth delta) +-- } +-- @ +class NetworkDelta a where + -- | The delta type containing only changed fields. + type Delta a + + -- | Compute a delta from a baseline to the current state. + diff :: a -> a -> Delta a + + -- | Apply a delta to a state, updating changed fields. + apply :: a -> Delta a -> a + +-- | Tracks snapshots and encodes deltas against acknowledged baselines. +-- +-- Used on the sender side to: +-- 1. Store pending snapshots awaiting ACK +-- 2. Encode current state as delta against confirmed baseline +-- 3. Promote ACK'd snapshots to confirmed baseline +data DeltaTracker a = DeltaTracker + { -- | Unconfirmed snapshots awaiting ACK, ordered by sequence + dtPending :: !(Seq (BaselineSeq, a)), + -- | The most recently ACK-confirmed baseline + dtConfirmed :: !(Maybe (BaselineSeq, a)), + -- | Maximum pending snapshots before oldest are dropped + dtMaxPending :: !Int + } + +makeFieldLabelsNoPrefix ''DeltaTracker + +instance (NFData a) => NFData (DeltaTracker a) where + rnf (DeltaTracker p c m) = rnf p `seq` rnf c `seq` rnf m + +-- | Create a new delta tracker with the given maximum pending snapshot count. +newDeltaTracker :: Int -> DeltaTracker a +newDeltaTracker maxPending = + DeltaTracker + { dtPending = Seq.empty, + dtConfirmed = Nothing, + dtMaxPending = maxPending + } + +-- | Encode a snapshot as delta against confirmed baseline. +-- +-- Returns encoded bytes containing @[baseline_seq: u16][delta_payload]@. +-- Stores the snapshot as pending until ACK'd. +deltaEncode :: + (NetworkDelta a, Storable a, Storable (Delta a)) => + BaselineSeq -> + a -> + DeltaTracker a -> + (BS.ByteString, DeltaTracker a) +deltaEncode seq' current tracker = + let -- Encode based on whether we have a confirmed baseline + encoded = case dtConfirmed tracker of + Just (baseSeq, baseline) -> + let delta = diff current baseline + in serialize baseSeq <> serialize delta + Nothing -> + -- No baseline - write sentinel + full state + serialize noBaseline <> serialize current + + -- Store pending snapshot + pending = dtPending tracker + pending' = + if Seq.length pending >= dtMaxPending tracker + then Seq.drop 1 pending Seq.|> (seq', current) + else pending Seq.|> (seq', current) + + tracker' = tracker & #dtPending .~ pending' + in (encoded, tracker') + +-- | Called when a sequence is ACK'd. +-- +-- Promotes the matching snapshot to confirmed baseline and discards +-- older pending entries. +deltaOnAck :: BaselineSeq -> DeltaTracker a -> DeltaTracker a +deltaOnAck seq' tracker = + case Seq.findIndexL (\(s, _) -> s == seq') (dtPending tracker) of + Nothing -> tracker + Just idx -> + let (ackSeq, snapshot) = Seq.index (dtPending tracker) idx + -- Drop everything older than the acked position + pending' = + Seq.filter (\(s, _) -> baselineSeqDiff s ackSeq >= 0) $ + Seq.drop (idx + 1) (dtPending tracker) + in tracker + & #dtPending + .~ pending' + & #dtConfirmed + ?~ (ackSeq, snapshot) + +-- | Reset tracker state (e.g. on reconnect). +deltaReset :: DeltaTracker a -> DeltaTracker a +deltaReset tracker = + tracker + & #dtPending + .~ Seq.empty + & #dtConfirmed + .~ Nothing + +-- | Returns the confirmed baseline sequence, if any. +deltaConfirmedSeq :: DeltaTracker a -> Maybe BaselineSeq +deltaConfirmedSeq = fmap fst . dtConfirmed + +-- | Ring buffer of confirmed snapshots per connection. +-- +-- Used on the receiving side to look up baselines referenced in +-- incoming delta payloads. +data BaselineManager a = BaselineManager + { -- | Stored snapshots: (sequence, state, timestamp) + bmSnapshots :: !(Seq (BaselineSeq, a, MonoTime)), + -- | Maximum stored snapshots + bmMaxSnapshots :: !Int, + -- | Timeout in milliseconds before snapshots expire + bmTimeoutMs :: !Double + } + +makeFieldLabelsNoPrefix ''BaselineManager + +instance (NFData a) => NFData (BaselineManager a) where + rnf (BaselineManager s m t) = rnf s `seq` rnf m `seq` rnf t + +-- | Create a new baseline manager with the given capacity and expiration timeout. +newBaselineManager :: + -- | Maximum snapshots to store + Int -> + -- | Timeout in milliseconds + Double -> + BaselineManager a +newBaselineManager maxSnapshots timeoutMs = + BaselineManager + { bmSnapshots = Seq.empty, + bmMaxSnapshots = maxSnapshots, + bmTimeoutMs = timeoutMs + } + +-- | Store a confirmed snapshot at the given sequence. +pushBaseline :: BaselineSeq -> a -> MonoTime -> BaselineManager a -> BaselineManager a +pushBaseline seq' state now manager = + let evictExpired = Seq.filter (\(_, _, ts) -> elapsedMs ts now < bmTimeoutMs manager) + evictOldest s + | Seq.length s >= bmMaxSnapshots manager = Seq.drop 1 s + | otherwise = s + snapshots = evictOldest (evictExpired (bmSnapshots manager)) Seq.|> (seq', state, now) + in manager & #bmSnapshots .~ snapshots + +-- | Look up a baseline by sequence number. +getBaseline :: BaselineSeq -> BaselineManager a -> Maybe a +getBaseline seq' manager = + case Seq.findIndexR (\(s, _, _) -> s == seq') (bmSnapshots manager) of + Nothing -> Nothing + Just idx -> + let (_, state, _) = Seq.index (bmSnapshots manager) idx + in Just state + +-- | Clear all stored baselines. +baselineReset :: BaselineManager a -> BaselineManager a +baselineReset manager = manager & #bmSnapshots .~ Seq.empty + +-- | Number of stored baselines. +baselineCount :: BaselineManager a -> Int +baselineCount = Seq.length . bmSnapshots + +-- | Whether the baseline buffer is empty. +baselineIsEmpty :: BaselineManager a -> Bool +baselineIsEmpty = Seq.null . bmSnapshots + +-- | Decode a delta-encoded payload. +-- +-- Requires access to the baseline manager to look up the referenced baseline. +deltaDecode :: + (NetworkDelta a, Storable a, Storable (Delta a)) => + BS.ByteString -> + BaselineManager a -> + Either String a +deltaDecode dat baselines + | BS.length dat < baselineSeqSize = Left "Delta payload too short" + | otherwise = deserialize headerBytes >>= decodeWithBaseline + where + headerBytes = BS.take baselineSeqSize dat + payloadBytes = BS.drop baselineSeqSize dat + + decodeWithBaseline baseSeq + | baseSeq == noBaseline = deserialize payloadBytes + | otherwise = do + baseline <- lookupBaseline baseSeq + delta <- deserialize payloadBytes + pure $ apply baseline delta + + lookupBaseline baseSeq = + maybe (Left $ "Missing baseline for seq " ++ show baseSeq) Right $ + getBaseline baseSeq baselines diff --git a/src/GBNet/Replication/Interest.hs b/src/GBNet/Replication/Interest.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Replication/Interest.hs @@ -0,0 +1,134 @@ +-- | +-- Module : GBNet.Replication.Interest +-- Description : Interest management for area-of-interest filtering +-- +-- Determines which entities are relevant to a given observer, reducing +-- bandwidth by only replicating nearby or important entities. +-- +-- Use 'RadiusInterest' for sphere-based filtering or 'GridInterest' for +-- cell-based spatial partitioning. +module GBNet.Replication.Interest + ( -- * Position type + Position, + + -- * Interest manager typeclass + InterestManager (..), + + -- * Radius-based interest + RadiusInterest, + newRadiusInterest, + radiusInterestRadius, + + -- * Grid-based interest + GridInterest, + newGridInterest, + gridInterestCellSize, + ) +where + +import Control.DeepSeq (NFData (..)) + +-- | 3D position as (x, y, z). +type Position = (Float, Float, Float) + +-- | Typeclass for determining entity relevance to an observer. +-- +-- Implement this for custom interest filtering strategies. +class InterestManager a where + -- | Returns 'True' if the entity at @entityPos@ is relevant to the + -- observer at @observerPos@. + relevant :: a -> Position -> Position -> Bool + + -- | Returns a priority modifier for the entity relative to the observer. + -- Values > 1.0 increase priority, < 1.0 decrease it. + -- Default returns 1.0 (no modification). + priorityMod :: a -> Position -> Position -> Float + priorityMod _ _ _ = 1.0 + +-- | Radius-based interest: entities within @radius@ distance are relevant. +-- +-- Also provides distance-based priority: closer entities get higher priority. +data RadiusInterest = RadiusInterest + { riRadius :: !Float, + riRadiusSq :: !Float + } + deriving (Show) + +instance NFData RadiusInterest where + rnf (RadiusInterest r rs) = rnf r `seq` rnf rs + +-- | Create a radius-based interest manager. +newRadiusInterest :: Float -> RadiusInterest +newRadiusInterest radius = + RadiusInterest + { riRadius = radius, + riRadiusSq = radius * radius + } + +-- | Get the radius. +radiusInterestRadius :: RadiusInterest -> Float +radiusInterestRadius = riRadius + +instance InterestManager RadiusInterest where + relevant ri (ex, ey, ez) (ox, oy, oz) = + let dx = ex - ox + dy = ey - oy + dz = ez - oz + distSq = dx * dx + dy * dy + dz * dz + in distSq <= riRadiusSq ri + + priorityMod ri (ex, ey, ez) (ox, oy, oz) = + let dx = ex - ox + dy = ey - oy + dz = ez - oz + distSq = dx * dx + dy * dy + dz * dz + radiusSq = riRadiusSq ri + in if distSq >= radiusSq + then 0.0 + else -- Linear falloff: closer entities get higher priority + 1.0 - sqrt (distSq / radiusSq) + +-- | Grid-based area-of-interest: entities in the same or neighboring cells +-- are considered relevant. +-- +-- More efficient than radius checks for large numbers of entities when +-- combined with spatial hashing. +data GridInterest = GridInterest + { giCellSize :: !Float, + giInvCellSize :: !Float + } + deriving (Show) + +instance NFData GridInterest where + rnf (GridInterest c i) = rnf c `seq` rnf i + +-- | Create a grid-based interest manager. +newGridInterest :: Float -> GridInterest +newGridInterest cellSize = + GridInterest + { giCellSize = cellSize, + giInvCellSize = 1.0 / cellSize + } + +-- | Get the cell size. +gridInterestCellSize :: GridInterest -> Float +gridInterestCellSize = giCellSize + +-- | Grid cell coordinate (internal). +data GridCell = GridCell !Int !Int !Int + deriving (Eq, Show) + +-- | Convert position to grid cell. +toCell :: GridInterest -> Position -> GridCell +toCell gi (x, y, z) = + let inv = giInvCellSize gi + in GridCell + (floor (x * inv)) + (floor (y * inv)) + (floor (z * inv)) + +instance InterestManager GridInterest where + relevant gi entityPos observerPos = + let GridCell ex ey ez = toCell gi entityPos + GridCell ox oy oz = toCell gi observerPos + in abs (ex - ox) <= 1 && abs (ey - oy) <= 1 && abs (ez - oz) <= 1 diff --git a/src/GBNet/Replication/Interpolation.hs b/src/GBNet/Replication/Interpolation.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Replication/Interpolation.hs @@ -0,0 +1,235 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedLabels #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- Module : GBNet.Replication.Interpolation +-- Description : Snapshot interpolation for smooth client-side rendering +-- +-- 'SnapshotBuffer' stores timestamped snapshots and interpolates between +-- them at a configurable playback delay, enabling smooth rendering despite +-- network jitter. +-- +-- Example usage: +-- +-- @ +-- -- On receiving server snapshot +-- let buffer' = pushSnapshot serverTimestamp playerState buffer +-- +-- -- On render tick +-- case sampleSnapshot renderTime buffer' of +-- Nothing -> renderLastKnown +-- Just interpolated -> render interpolated +-- @ +module GBNet.Replication.Interpolation + ( -- * Constants + defaultBufferDepth, + defaultPlaybackDelayMs, + + -- * Interpolatable typeclass + Interpolatable (..), + + -- * Snapshot buffer + TimestampedSnapshot (..), + SnapshotBuffer, + newSnapshotBuffer, + newSnapshotBufferWithConfig, + + -- * Operations + pushSnapshot, + sampleSnapshot, + snapshotReset, + + -- * Queries + snapshotCount, + snapshotIsEmpty, + snapshotReady, + snapshotPlaybackDelayMs, + + -- * Configuration + setPlaybackDelayMs, + ) +where + +import Data.Sequence (Seq, ViewL (..), ViewR (..), (|>)) +import qualified Data.Sequence as Seq +import Optics ((&), (.~)) +import Optics.TH (makeFieldLabelsNoPrefix) + +-- | Default number of snapshots to buffer before interpolation begins. +defaultBufferDepth :: Int +defaultBufferDepth = 3 + +-- | Default playback delay in milliseconds behind the latest received snapshot. +defaultPlaybackDelayMs :: Double +defaultPlaybackDelayMs = 100.0 + +-- | Typeclass for types that support linear interpolation between two states. +-- +-- Implement this for any game state that needs smooth interpolation: +-- +-- @ +-- instance Interpolatable Vec3 where +-- lerp (Vec3 x1 y1 z1) (Vec3 x2 y2 z2) t = +-- Vec3 (x1 + (x2 - x1) * t) +-- (y1 + (y2 - y1) * t) +-- (z1 + (z2 - z1) * t) +-- @ +class Interpolatable a where + -- | Linearly interpolate between @self@ and @other@ by factor @t@ in [0, 1]. + lerp :: a -> a -> Float -> a + +-- | A timestamped snapshot for interpolation (internal). +data TimestampedSnapshot a = TimestampedSnapshot + { tsTimestamp :: !Double, + tsState :: !a + } + deriving (Show) + +makeFieldLabelsNoPrefix ''TimestampedSnapshot + +-- | Ring buffer of timestamped snapshots with interpolation sampling. +data SnapshotBuffer a = SnapshotBuffer + { sbSnapshots :: !(Seq (TimestampedSnapshot a)), + sbBufferDepth :: !Int, + sbPlaybackDelayMs :: !Double + } + deriving (Show) + +makeFieldLabelsNoPrefix ''SnapshotBuffer + +-- | Create a snapshot buffer with default settings. +newSnapshotBuffer :: SnapshotBuffer a +newSnapshotBuffer = + SnapshotBuffer + { sbSnapshots = Seq.empty, + sbBufferDepth = defaultBufferDepth, + sbPlaybackDelayMs = defaultPlaybackDelayMs + } + +-- | Create a snapshot buffer with custom settings. +newSnapshotBufferWithConfig :: + -- | Buffer depth (number of snapshots before ready) + Int -> + -- | Playback delay in milliseconds + Double -> + SnapshotBuffer a +newSnapshotBufferWithConfig depth delay = + SnapshotBuffer + { sbSnapshots = Seq.empty, + sbBufferDepth = depth, + sbPlaybackDelayMs = delay + } + +-- | Push a new snapshot with its server timestamp (in milliseconds). +-- +-- Snapshots must arrive in order; out-of-order snapshots are dropped. +pushSnapshot :: Double -> a -> SnapshotBuffer a -> SnapshotBuffer a +pushSnapshot timestamp state buffer = + let snapshots = sbSnapshots buffer + in -- Check if newer than latest + case Seq.viewr snapshots of + EmptyR -> + -- First snapshot + buffer & #sbSnapshots .~ Seq.singleton (TimestampedSnapshot timestamp state) + _ :> last' -> + if timestamp <= tsTimestamp last' + then buffer -- Drop out-of-order + else + let snapshots' = snapshots |> TimestampedSnapshot timestamp state + -- Keep buffer bounded + maxEntries = sbBufferDepth buffer * 2 + trimmed = + if Seq.length snapshots' > maxEntries + then Seq.drop (Seq.length snapshots' - maxEntries) snapshots' + else snapshots' + in buffer & #sbSnapshots .~ trimmed + +-- | Sample an interpolated state at @renderTime@ (in milliseconds). +-- +-- Returns 'Nothing' if insufficient snapshots are buffered. +sampleSnapshot :: (Interpolatable a) => Double -> SnapshotBuffer a -> Maybe a +sampleSnapshot renderTime buffer = + let targetTime = renderTime - sbPlaybackDelayMs buffer + snapshots = sbSnapshots buffer + in if Seq.length snapshots < 2 + then Nothing + else findAndInterpolate targetTime snapshots + +-- | Find bracketing snapshots and interpolate. +findAndInterpolate :: (Interpolatable a) => Double -> Seq (TimestampedSnapshot a) -> Maybe a +findAndInterpolate targetTime snapshots = + case findBracket targetTime (Seq.viewl snapshots) of + Just (a, b) -> + let duration = tsTimestamp b - tsTimestamp a + in if duration <= 0.0 + then Just (tsState a) + else + let t = realToFrac ((targetTime - tsTimestamp a) / duration) + tClamped = max 0.0 (min 1.0 t) + in Just (lerp (tsState a) (tsState b) tClamped) + Nothing -> + -- Target outside range - return boundary + case (Seq.viewl snapshots, Seq.viewr snapshots) of + (first :< _, _ :> lastSnap) + | targetTime > tsTimestamp lastSnap -> Just (tsState lastSnap) + | targetTime < tsTimestamp first -> Just (tsState first) + _ -> Nothing + +-- | Find two snapshots bracketing the target time. +findBracket :: + Double -> + ViewL (TimestampedSnapshot a) -> + Maybe (TimestampedSnapshot a, TimestampedSnapshot a) +findBracket _ EmptyL = Nothing +findBracket _ (_ :< rest) | Seq.null rest = Nothing +findBracket targetTime (a :< rest) = + case Seq.viewl rest of + b :< _ -> + if targetTime >= tsTimestamp a && targetTime <= tsTimestamp b + then Just (a, b) + else findBracket targetTime (Seq.viewl rest) + EmptyL -> Nothing + +-- | Number of buffered snapshots. +snapshotCount :: SnapshotBuffer a -> Int +snapshotCount = Seq.length . sbSnapshots + +-- | Whether the buffer is empty. +snapshotIsEmpty :: SnapshotBuffer a -> Bool +snapshotIsEmpty = Seq.null . sbSnapshots + +-- | Whether enough snapshots are buffered to begin interpolation. +snapshotReady :: SnapshotBuffer a -> Bool +snapshotReady buffer = Seq.length (sbSnapshots buffer) >= sbBufferDepth buffer + +-- | Clear all buffered snapshots. +snapshotReset :: SnapshotBuffer a -> SnapshotBuffer a +snapshotReset buffer = buffer & #sbSnapshots .~ Seq.empty + +-- | Get the playback delay in milliseconds. +snapshotPlaybackDelayMs :: SnapshotBuffer a -> Double +snapshotPlaybackDelayMs = sbPlaybackDelayMs + +-- | Set the playback delay in milliseconds. +setPlaybackDelayMs :: Double -> SnapshotBuffer a -> SnapshotBuffer a +setPlaybackDelayMs delay buffer = buffer & #sbPlaybackDelayMs .~ delay + +-- Common instances + +instance Interpolatable Float where + lerp a b t = a + (b - a) * t + +instance Interpolatable Double where + lerp a b t = a + (b - a) * realToFrac t + +instance (Interpolatable a, Interpolatable b) => Interpolatable (a, b) where + lerp (a1, b1) (a2, b2) t = (lerp a1 a2 t, lerp b1 b2 t) + +instance (Interpolatable a, Interpolatable b, Interpolatable c) => Interpolatable (a, b, c) where + lerp (a1, b1, c1) (a2, b2, c2) t = (lerp a1 a2 t, lerp b1 b2 t, lerp c1 c2 t) diff --git a/src/GBNet/Replication/Priority.hs b/src/GBNet/Replication/Priority.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Replication/Priority.hs @@ -0,0 +1,174 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedLabels #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- Module : GBNet.Replication.Priority +-- Description : Priority accumulator for bandwidth-limited entity replication +-- +-- 'PriorityAccumulator' tracks per-entity priority that grows over time, +-- ensuring all entities eventually get sent even at low priority. +-- +-- Use this to fairly allocate bandwidth when replicating many entities: +-- +-- @ +-- let acc = newPriorityAccumulator +-- & register playerId 10.0 -- high priority (10 units/sec) +-- & register npcId 2.0 -- low priority (2 units/sec) +-- +-- -- Each tick, accumulate priority based on elapsed time +-- let acc' = accumulate 0.016 acc -- 16ms tick +-- +-- -- Drain entities that fit in budget +-- let (selected, acc'') = drainTop 1200 entitySize acc' +-- @ +module GBNet.Replication.Priority + ( -- * Priority accumulator + PriorityAccumulator, + newPriorityAccumulator, + + -- * Entity management + register, + unregister, + + -- * Priority operations + accumulate, + applyModifier, + drainTop, + + -- * Queries + priorityCount, + priorityIsEmpty, + getPriority, + ) +where + +import Data.List (sortBy) +import Data.Map.Strict (Map) +import qualified Data.Map.Strict as Map +import Data.Ord (Down (..), comparing) +import Optics ((%~), (&), (.~)) +import Optics.TH (makeFieldLabelsNoPrefix) + +-- | Per-entity priority tracking entry. +data PriorityEntry = PriorityEntry + { peBase :: !Float, + peAccumulated :: !Float + } + deriving (Show) + +makeFieldLabelsNoPrefix ''PriorityEntry + +-- | Accumulates priority per entity and drains the highest-priority entities +-- that fit within a byte budget. +-- +-- Type parameter @id@ is the entity identifier type (must be 'Ord'). +newtype PriorityAccumulator id = PriorityAccumulator + { paEntries :: Map id PriorityEntry + } + deriving (Show) + +-- | Create an empty priority accumulator. +newPriorityAccumulator :: PriorityAccumulator id +newPriorityAccumulator = PriorityAccumulator Map.empty + +-- | Register an entity with a base priority. +-- +-- Higher base priority = more frequent sends. +-- Priority accumulates as @base * dt@ each tick. +register :: (Ord id) => id -> Float -> PriorityAccumulator id -> PriorityAccumulator id +register entityId basePriority (PriorityAccumulator entries) = + PriorityAccumulator $ + Map.insert + entityId + PriorityEntry {peBase = basePriority, peAccumulated = 0.0} + entries + +-- | Remove an entity from tracking. +unregister :: (Ord id) => id -> PriorityAccumulator id -> PriorityAccumulator id +unregister entityId (PriorityAccumulator entries) = + PriorityAccumulator $ Map.delete entityId entries + +-- | Advance accumulated priority for all entities by @dt@ seconds. +-- +-- Call this once per tick with elapsed time. +accumulate :: Float -> PriorityAccumulator id -> PriorityAccumulator id +accumulate dt (PriorityAccumulator entries) = + PriorityAccumulator $ + Map.map + (\e -> e & #peAccumulated %~ (+ (peBase e * dt))) + entries + +-- | Apply a priority modifier to a specific entity. +-- +-- Use with interest management: closer entities get modifier > 1.0, +-- farther entities get modifier < 1.0. +applyModifier :: (Ord id) => id -> Float -> PriorityAccumulator id -> PriorityAccumulator id +applyModifier entityId modifier (PriorityAccumulator entries) = + PriorityAccumulator $ + Map.adjust + (\e -> e & #peAccumulated %~ (* modifier)) + entityId + entries + +-- | Drain the highest-priority entities that fit within @budgetBytes@. +-- +-- @sizeFunc@ returns the serialized size in bytes for a given entity ID. +-- Returns the selected entity IDs in priority order (highest first), +-- and the updated accumulator with those entities' priorities reset. +drainTop :: + (Ord id) => + -- | Byte budget + Int -> + -- | Size function: id -> bytes + (id -> Int) -> + PriorityAccumulator id -> + ([id], PriorityAccumulator id) +drainTop budgetBytes sizeFunc (PriorityAccumulator entries) = + let -- Sort by accumulated priority (descending) + sorted = + sortBy (comparing (Down . peAccumulated . snd)) $ + Map.toList entries + + -- Select entities that fit in budget + (selected, _remaining) = selectWithinBudget budgetBytes sizeFunc sorted [] + + -- Reset priority for selected entities + entries' = + foldr + (Map.adjust (#peAccumulated .~ 0.0)) + entries + selected + in (selected, PriorityAccumulator entries') + +-- | Helper to select entities within budget. +selectWithinBudget :: + Int -> + (id -> Int) -> + [(id, PriorityEntry)] -> + [id] -> + ([id], Int) +selectWithinBudget remaining _ [] acc = (reverse acc, remaining) +selectWithinBudget remaining sizeFunc ((eid, _) : rest) acc = + let size = sizeFunc eid + in if size > remaining + then selectWithinBudget remaining sizeFunc rest acc + else selectWithinBudget (remaining - size) sizeFunc rest (eid : acc) + +-- | Number of tracked entities. +priorityCount :: PriorityAccumulator id -> Int +priorityCount = Map.size . paEntries + +-- | Whether the accumulator is empty. +priorityIsEmpty :: PriorityAccumulator id -> Bool +priorityIsEmpty = Map.null . paEntries + +-- | Get the current accumulated priority for an entity. +getPriority :: (Ord id) => id -> PriorityAccumulator id -> Maybe Float +getPriority entityId (PriorityAccumulator entries) = + peAccumulated <$> Map.lookup entityId entries diff --git a/src/GBNet/Security.hs b/src/GBNet/Security.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Security.hs @@ -0,0 +1,255 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedLabels #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- Module : GBNet.Security +-- Description : CRC32C integrity, rate limiting, and connect tokens +-- +-- Packet integrity via CRC32C (Castagnoli), connection rate limiting, +-- and netcode.io-style connect token validation. +module GBNet.Security + ( -- * CRC32C + crc32c, + appendCrc32, + validateAndStripCrc32, + crc32Size, + + -- * Rate limiting + RateLimiter (..), + newRateLimiter, + rateLimiterAllow, + + -- * Connect tokens + ConnectToken (..), + newConnectToken, + isTokenExpired, + TokenError (..), + TokenValidator (..), + newTokenValidator, + validateToken, + ) +where + +import Data.Bits (shiftL, shiftR, (.|.)) +import qualified Data.ByteString as BS +import Data.ByteString.Internal (unsafeCreate) +import qualified Data.ByteString.Unsafe as BSU +import qualified Data.Digest.CRC32C as CRC +import Data.List (minimumBy) +import Data.Map.Strict (Map) +import qualified Data.Map.Strict as Map +import Data.Ord (comparing) +import Data.Word (Word32, Word64, Word8) +import Foreign.Storable (pokeByteOff) +import GBNet.Reliability (MonoTime, elapsedMs) +import Optics ((%~), (&), (.~)) +import Optics.TH (makeFieldLabelsNoPrefix) + +-- | CRC32 checksum size in bytes. +crc32Size :: Int +crc32Size = 4 + +-- | Compute CRC32C checksum (hardware-accelerated via SSE4.2\/ARMv8 CRC). +crc32c :: BS.ByteString -> Word32 +crc32c = CRC.crc32c +{-# INLINE crc32c #-} + +-- | Append CRC32C checksum to data (little-endian). +appendCrc32 :: BS.ByteString -> BS.ByteString +appendCrc32 dat = + let crc = crc32c dat + in dat <> word32ToLEBytes crc + +-- | Validate and strip CRC32C. Returns Nothing if corrupt. +validateAndStripCrc32 :: BS.ByteString -> Maybe BS.ByteString +validateAndStripCrc32 dat + | BS.length dat < crc32Size = Nothing + | otherwise = + let payloadLen = BS.length dat - crc32Size + payload = BS.take payloadLen dat + crcBytes = BS.drop payloadLen dat + expected = word32FromLEBytes crcBytes + actual = crc32c payload + in if actual == expected + then Just payload + else Nothing + +-- | Convert Word32 to little-endian bytes. +-- Uses zero-allocation direct memory writes. +word32ToLEBytes :: Word32 -> BS.ByteString +word32ToLEBytes !w = unsafeCreate crc32Size $ \ptr -> do + pokeByteOff ptr 0 (fromIntegral w :: Word8) + pokeByteOff ptr 1 (fromIntegral (w `shiftR` 8) :: Word8) + pokeByteOff ptr 2 (fromIntegral (w `shiftR` 16) :: Word8) + pokeByteOff ptr 3 (fromIntegral (w `shiftR` 24) :: Word8) +{-# INLINE word32ToLEBytes #-} + +-- | Convert little-endian bytes to Word32. +-- Uses direct memory access for speed. +-- Caller must ensure at least 4 bytes; 'validateAndStripCrc32' guarantees this. +word32FromLEBytes :: BS.ByteString -> Word32 +word32FromLEBytes !bs = + let !b0 = fromIntegral (BSU.unsafeIndex bs 0) :: Word32 + !b1 = fromIntegral (BSU.unsafeIndex bs 1) :: Word32 + !b2 = fromIntegral (BSU.unsafeIndex bs 2) :: Word32 + !b3 = fromIntegral (BSU.unsafeIndex bs 3) :: Word32 + in b0 .|. (b1 `shiftL` 8) .|. (b2 `shiftL` 16) .|. (b3 `shiftL` 24) +{-# INLINE word32FromLEBytes #-} + +-- | Cleanup interval — sweep stale entries every 5 seconds. +cleanupIntervalMs :: Double +cleanupIntervalMs = 5000.0 + +-- | Rate limiter for connection requests per source. +-- Self-cleaning: stale entries are pruned automatically during 'rateLimiterAllow'. +data RateLimiter = RateLimiter + { rlRequests :: !(Map Word64 [MonoTime]), + rlMaxRequestsPerSecond :: !Int, + rlWindowMs :: !Double, + rlLastCleanup :: !MonoTime + } + deriving (Show) + +makeFieldLabelsNoPrefix ''RateLimiter + +-- | Create a new rate limiter. +newRateLimiter :: Int -> MonoTime -> RateLimiter +newRateLimiter maxReqs now = + RateLimiter + { rlRequests = Map.empty, + rlMaxRequestsPerSecond = maxReqs, + rlWindowMs = 1000.0, + rlLastCleanup = now + } + +-- | Check if a request should be allowed. Returns (allowed, updatedLimiter). +-- Automatically prunes stale entries when the cleanup interval has elapsed. +rateLimiterAllow :: Word64 -> MonoTime -> RateLimiter -> (Bool, RateLimiter) +rateLimiterAllow addrKey now rl + | recentCount >= rlMaxRequestsPerSecond rl' = (False, rl' & #rlRequests %~ Map.insert addrKey recent) + | otherwise = (True, rl' & #rlRequests %~ Map.insert addrKey (now : recent)) + where + rl' = maybeCleanup now rl + window = rlWindowMs rl' + timestamps = Map.findWithDefault [] addrKey (rlRequests rl') + -- Filter and count in single pass + (recentCount, recent) = foldr countRecent (0, []) timestamps + countRecent t (n, acc) + | elapsedMs t now < window = (n + 1, t : acc) + | otherwise = (n, acc) + +-- | Sweep stale entries if enough time has passed since last cleanup. +maybeCleanup :: MonoTime -> RateLimiter -> RateLimiter +maybeCleanup now rl + | elapsedMs (rlLastCleanup rl) now < cleanupIntervalMs = rl + | otherwise = + let window = rlWindowMs rl + cleanup = filter (\t -> elapsedMs t now < window) + cleaned = Map.map cleanup (rlRequests rl) + nonEmpty = Map.filter (not . null) cleaned + in rl & #rlRequests .~ nonEmpty & #rlLastCleanup .~ now + +-- | Connect token for authentication. +data ConnectToken = ConnectToken + { ctClientId :: !Word64, + ctCreateTime :: !MonoTime, + ctExpireDurationMs :: !Double, + ctUserData :: !BS.ByteString + } + deriving (Show) + +makeFieldLabelsNoPrefix ''ConnectToken + +-- | Create a new connect token. +newConnectToken :: Word64 -> Double -> BS.ByteString -> MonoTime -> ConnectToken +newConnectToken clientId expireMs userData now = + ConnectToken + { ctClientId = clientId, + ctCreateTime = now, + ctExpireDurationMs = expireMs, + ctUserData = userData + } + +-- | Check if token is expired. +isTokenExpired :: MonoTime -> ConnectToken -> Bool +isTokenExpired now token = elapsedMs (ctCreateTime token) now > ctExpireDurationMs token + +-- | Token validation errors. +data TokenError + = TokenExpired + | TokenReplayed + | TokenInvalid + deriving (Eq, Show) + +-- | Server-side token validator. +data TokenValidator = TokenValidator + { tvUsedTokens :: !(Map Word64 MonoTime), + tvTokenLifetimeMs :: !Double, + tvMaxTrackedTokens :: !Int, + tvTokensEvicted :: !Int + } + deriving (Show) + +makeFieldLabelsNoPrefix ''TokenValidator + +-- | Create a new token validator. +newTokenValidator :: Double -> Int -> TokenValidator +newTokenValidator lifetimeMs maxTracked = + TokenValidator + { tvUsedTokens = Map.empty, + tvTokenLifetimeMs = lifetimeMs, + tvMaxTrackedTokens = maxTracked, + tvTokensEvicted = 0 + } + +-- | Validate a connect token. Returns (Result, updatedValidator). +validateToken :: ConnectToken -> MonoTime -> TokenValidator -> (Either TokenError Word64, TokenValidator) +validateToken token now tv + | isTokenExpired now token = (Left TokenExpired, tv) + | Map.member (ctClientId token) (tvUsedTokens tv) = (Left TokenReplayed, tv) + | otherwise = + let tv' = + tv & #tvUsedTokens %~ Map.insert (ctClientId token) now + tv'' = enforceLimit now tv' + in (Right (ctClientId token), tv'') + +-- | Enforce maximum tracked tokens limit. +enforceLimit :: MonoTime -> TokenValidator -> TokenValidator +enforceLimit now tv + | Map.size (tvUsedTokens tv) <= tvMaxTrackedTokens tv = tv + | otherwise = + -- First try cleanup + let cleaned = cleanupExpired now tv + in if Map.size (tvUsedTokens cleaned) <= tvMaxTrackedTokens cleaned + then cleaned + else evictOldest cleaned + +-- | Remove expired tokens. +cleanupExpired :: MonoTime -> TokenValidator -> TokenValidator +cleanupExpired now tv = + let lifetime = tvTokenLifetimeMs tv + keep (_clientId, created) = elapsedMs created now < lifetime + kept = Map.filterWithKey (curry keep) (tvUsedTokens tv) + in tv & #tvUsedTokens .~ kept + +-- | Evict oldest token. +evictOldest :: TokenValidator -> TokenValidator +evictOldest tv = + case findOldest (Map.toList (tvUsedTokens tv)) of + Nothing -> tv + Just (oldestId, _) -> + tv + & #tvUsedTokens + %~ Map.delete oldestId + & #tvTokensEvicted + %~ (+ 1) + where + findOldest [] = Nothing + findOldest xs = Just $ minimumBy (comparing snd) xs diff --git a/src/GBNet/Serialize.hs b/src/GBNet/Serialize.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Serialize.hs @@ -0,0 +1,188 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | +-- Module : GBNet.Serialize +-- Description : Zero-allocation serialization for Storable types +-- +-- Provides fast serialization for any type with a 'Storable' instance. +-- Use @deriveStorable@ from "GBNet.Serialize.TH" to generate instances, +-- then use 'serialize' and 'deserialize' for zero-copy conversion. +-- +-- Wire format is always little-endian for cross-platform compatibility. +-- On LE platforms (x86, ARM), the LE helpers compile to native operations +-- with zero overhead — GHC eliminates the dead branch at compile time. +-- +-- @ +-- data Vec3 = Vec3 !Float !Float !Float +-- deriveStorable ''Vec3 +-- +-- -- Pure, zero-allocation serialization: +-- let bytes = serialize (Vec3 1.0 2.0 3.0) +-- let Right v = deserialize bytes :: Either String Vec3 +-- @ +module GBNet.Serialize + ( -- * Serialization (pure, zero-copy) + serialize, + deserialize, + + -- * Re-exports for deriveStorable + ByteString, + Storable (..), + plusPtr, + castPtr, + + -- * Little-endian helpers (used by generated Storable instances) + pokeWord16LE, + pokeWord32LE, + pokeWord64LE, + pokeInt16LE, + pokeInt32LE, + pokeInt64LE, + pokeFloatLE, + pokeDoubleLE, + peekWord16LE, + peekWord32LE, + peekWord64LE, + peekInt16LE, + peekInt32LE, + peekInt64LE, + peekFloatLE, + peekDoubleLE, + ) +where + +import Data.ByteString (ByteString) +import qualified Data.ByteString as BS +import Data.ByteString.Internal (toForeignPtr, unsafeCreate) +import Data.Int (Int16, Int32, Int64) +import Data.Word (Word16, Word32, Word64, byteSwap16, byteSwap32, byteSwap64) +import Foreign.ForeignPtr (withForeignPtr) +import Foreign.Ptr (Ptr, castPtr, plusPtr) +import Foreign.Storable (Storable (..)) +import GHC.ByteOrder (ByteOrder (..), targetByteOrder) +import GHC.Float (castDoubleToWord64, castFloatToWord32, castWord32ToFloat, castWord64ToDouble) +import System.IO.Unsafe (unsafeDupablePerformIO) + +-- | Serialize any Storable type to a ByteString. +-- Pure interface, zero allocation beyond the result ByteString. +serialize :: (Storable a) => a -> ByteString +serialize !x = unsafeCreate (sizeOf x) $ \ptr -> + poke (castPtr ptr) x +{-# INLINE serialize #-} + +-- | Deserialize a ByteString to any Storable type. +-- Returns Left with error message if buffer is too short. +deserialize :: forall a. (Storable a) => ByteString -> Either String a +deserialize !bs + | BS.length bs < needed = Left "deserialize: buffer too short" + | otherwise = unsafeDupablePerformIO $ do + let (fptr, off, _) = toForeignPtr bs + val <- withForeignPtr fptr $ \ptr -> + peek (castPtr (ptr `plusPtr` off)) + return $ Right val + where + needed = sizeOf (undefined :: a) +{-# INLINE deserialize #-} + +-- --------------------------------------------------------------------------- +-- Little-endian poke helpers +-- On LE: compiles to native pokeByteOff (zero overhead). +-- On BE: byte-swaps before writing (single bswap instruction). +-- --------------------------------------------------------------------------- + +-- | Write a 'Word16' in little-endian byte order. +pokeWord16LE :: Ptr a -> Int -> Word16 -> IO () +pokeWord16LE ptr off w = case targetByteOrder of + LittleEndian -> pokeByteOff ptr off w + BigEndian -> pokeByteOff ptr off (byteSwap16 w) +{-# INLINE pokeWord16LE #-} + +-- | Write a 'Word32' in little-endian byte order. +pokeWord32LE :: Ptr a -> Int -> Word32 -> IO () +pokeWord32LE ptr off w = case targetByteOrder of + LittleEndian -> pokeByteOff ptr off w + BigEndian -> pokeByteOff ptr off (byteSwap32 w) +{-# INLINE pokeWord32LE #-} + +-- | Write a 'Word64' in little-endian byte order. +pokeWord64LE :: Ptr a -> Int -> Word64 -> IO () +pokeWord64LE ptr off w = case targetByteOrder of + LittleEndian -> pokeByteOff ptr off w + BigEndian -> pokeByteOff ptr off (byteSwap64 w) +{-# INLINE pokeWord64LE #-} + +-- | Write an 'Int16' in little-endian byte order. +pokeInt16LE :: Ptr a -> Int -> Int16 -> IO () +pokeInt16LE ptr off i = pokeWord16LE ptr off (fromIntegral i) +{-# INLINE pokeInt16LE #-} + +-- | Write an 'Int32' in little-endian byte order. +pokeInt32LE :: Ptr a -> Int -> Int32 -> IO () +pokeInt32LE ptr off i = pokeWord32LE ptr off (fromIntegral i) +{-# INLINE pokeInt32LE #-} + +-- | Write an 'Int64' in little-endian byte order. +pokeInt64LE :: Ptr a -> Int -> Int64 -> IO () +pokeInt64LE ptr off i = pokeWord64LE ptr off (fromIntegral i) +{-# INLINE pokeInt64LE #-} + +-- | Write a 'Float' in little-endian byte order. +pokeFloatLE :: Ptr a -> Int -> Float -> IO () +pokeFloatLE ptr off f = pokeWord32LE ptr off (castFloatToWord32 f) +{-# INLINE pokeFloatLE #-} + +-- | Write a 'Double' in little-endian byte order. +pokeDoubleLE :: Ptr a -> Int -> Double -> IO () +pokeDoubleLE ptr off d = pokeWord64LE ptr off (castDoubleToWord64 d) +{-# INLINE pokeDoubleLE #-} + +-- --------------------------------------------------------------------------- +-- Little-endian peek helpers +-- --------------------------------------------------------------------------- + +-- | Read a 'Word16' in little-endian byte order. +peekWord16LE :: Ptr a -> Int -> IO Word16 +peekWord16LE ptr off = case targetByteOrder of + LittleEndian -> peekByteOff ptr off + BigEndian -> byteSwap16 <$> peekByteOff ptr off +{-# INLINE peekWord16LE #-} + +-- | Read a 'Word32' in little-endian byte order. +peekWord32LE :: Ptr a -> Int -> IO Word32 +peekWord32LE ptr off = case targetByteOrder of + LittleEndian -> peekByteOff ptr off + BigEndian -> byteSwap32 <$> peekByteOff ptr off +{-# INLINE peekWord32LE #-} + +-- | Read a 'Word64' in little-endian byte order. +peekWord64LE :: Ptr a -> Int -> IO Word64 +peekWord64LE ptr off = case targetByteOrder of + LittleEndian -> peekByteOff ptr off + BigEndian -> byteSwap64 <$> peekByteOff ptr off +{-# INLINE peekWord64LE #-} + +-- | Read an 'Int16' in little-endian byte order. +peekInt16LE :: Ptr a -> Int -> IO Int16 +peekInt16LE ptr off = fromIntegral <$> peekWord16LE ptr off +{-# INLINE peekInt16LE #-} + +-- | Read an 'Int32' in little-endian byte order. +peekInt32LE :: Ptr a -> Int -> IO Int32 +peekInt32LE ptr off = fromIntegral <$> peekWord32LE ptr off +{-# INLINE peekInt32LE #-} + +-- | Read an 'Int64' in little-endian byte order. +peekInt64LE :: Ptr a -> Int -> IO Int64 +peekInt64LE ptr off = fromIntegral <$> peekWord64LE ptr off +{-# INLINE peekInt64LE #-} + +-- | Read a 'Float' in little-endian byte order. +peekFloatLE :: Ptr a -> Int -> IO Float +peekFloatLE ptr off = castWord32ToFloat <$> peekWord32LE ptr off +{-# INLINE peekFloatLE #-} + +-- | Read a 'Double' in little-endian byte order. +peekDoubleLE :: Ptr a -> Int -> IO Double +peekDoubleLE ptr off = castWord64ToDouble <$> peekWord64LE ptr off +{-# INLINE peekDoubleLE #-} diff --git a/src/GBNet/Serialize/TH.hs b/src/GBNet/Serialize/TH.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Serialize/TH.hs @@ -0,0 +1,306 @@ +{-# LANGUAGE TemplateHaskell #-} + +-- | +-- Module : GBNet.Serialize.TH +-- Description : Template Haskell for zero-allocation Storable instances +-- +-- Generates 'Storable' instances for product types, enabling zero-allocation +-- serialization via @serialize@. Wire format is always little-endian for +-- cross-platform compatibility. On LE platforms, the generated code compiles +-- to native memory operations with zero overhead. +-- +-- Usage: +-- +-- @ +-- {-# LANGUAGE TemplateHaskell #-} +-- import GBNet.Serialize.TH +-- import GBNet.Serialize +-- +-- data Vec3 = Vec3 !Float !Float !Float +-- deriveStorable ''Vec3 +-- +-- data Transform = Transform !Vec3 !Quaternion +-- deriveStorable ''Transform -- nested types just work +-- +-- -- Serialize any Storable type: +-- serialize (Vec3 1.0 2.0 3.0) -- :: ByteString +-- serialize (Transform pos rot) -- :: ByteString +-- @ +module GBNet.Serialize.TH + ( deriveStorable, + ) +where + +import Data.Int (Int16, Int32, Int64, Int8) +import Data.Word (Word16, Word32, Word64, Word8) +import Foreign.Ptr (castPtr, plusPtr) +import Foreign.Storable (Storable (..)) +import GBNet.Serialize + ( peekDoubleLE, + peekFloatLE, + peekInt16LE, + peekInt32LE, + peekInt64LE, + peekWord16LE, + peekWord32LE, + peekWord64LE, + pokeDoubleLE, + pokeFloatLE, + pokeInt16LE, + pokeInt32LE, + pokeInt64LE, + pokeWord16LE, + pokeWord32LE, + pokeWord64LE, + ) +import Language.Haskell.TH + +-- | Byte sizes and alignments of primitive types. +primitiveSizeAlign :: Type -> Maybe (Int, Int) +primitiveSizeAlign (ConT n) + | n == ''Word8 = Just (1, 1) + | n == ''Word16 = Just (2, 2) + | n == ''Word32 = Just (4, 4) + | n == ''Word64 = Just (8, 8) + | n == ''Int8 = Just (1, 1) + | n == ''Int16 = Just (2, 2) + | n == ''Int32 = Just (4, 4) + | n == ''Int64 = Just (8, 8) + | n == ''Float = Just (4, 4) + | n == ''Double = Just (8, 8) +primitiveSizeAlign _ = Nothing + +-- | Check if a type is a known primitive. +isPrimitive :: Type -> Bool +isPrimitive t = case primitiveSizeAlign t of + Just _ -> True + Nothing -> False + +-- | Get the LE poke function for a multi-byte primitive. +-- Returns Nothing for single-byte types (no endianness concern). +lePokeName :: Type -> Maybe Name +lePokeName (ConT n) + | n == ''Word16 = Just 'pokeWord16LE + | n == ''Word32 = Just 'pokeWord32LE + | n == ''Word64 = Just 'pokeWord64LE + | n == ''Int16 = Just 'pokeInt16LE + | n == ''Int32 = Just 'pokeInt32LE + | n == ''Int64 = Just 'pokeInt64LE + | n == ''Float = Just 'pokeFloatLE + | n == ''Double = Just 'pokeDoubleLE +lePokeName _ = Nothing + +-- | Get the LE peek function for a multi-byte primitive. +-- Returns Nothing for single-byte types (no endianness concern). +lePeekName :: Type -> Maybe Name +lePeekName (ConT n) + | n == ''Word16 = Just 'peekWord16LE + | n == ''Word32 = Just 'peekWord32LE + | n == ''Word64 = Just 'peekWord64LE + | n == ''Int16 = Just 'peekInt16LE + | n == ''Int32 = Just 'peekInt32LE + | n == ''Int64 = Just 'peekInt64LE + | n == ''Float = Just 'peekFloatLE + | n == ''Double = Just 'peekDoubleLE +lePeekName _ = Nothing + +-- | Derive a Storable instance for a product type. +-- Supports primitives and nested Storable types. +deriveStorable :: Name -> Q [Dec] +deriveStorable typeName = do + info <- reify typeName + case info of + TyConI (DataD _ _ _ _ [con] _) -> do + let (conName, fieldTypes) = conFieldTypes con + mkStorableInstance typeName conName fieldTypes + TyConI DataD {} -> + fail $ "deriveStorable: " ++ show typeName ++ " must have exactly one constructor (no sum types)" + TyConI (NewtypeD _ _ _ _ con _) -> do + let (conName, fieldTypes) = conFieldTypes con + mkStorableInstance typeName conName fieldTypes + _ -> fail $ "deriveStorable: " ++ show typeName ++ " is not a data type" + +-- | Generate Storable instance. +mkStorableInstance :: Name -> Name -> [Type] -> Q [Dec] +mkStorableInstance typeName conName fieldTypes = do + let fieldCount = length fieldTypes + + -- Generate sizeOf + sizeOfBody <- mkSizeOfBody fieldTypes + + -- Generate alignment (max of field alignments, minimum 1) + alignBody <- mkAlignmentBody fieldTypes + + -- Generate poke + pokeBody <- mkPokeBody conName fieldTypes fieldCount + + -- Generate peek + peekBody <- mkPeekBody conName fieldTypes + + -- Build the instance + let inst = + InstanceD + Nothing + [] -- No constraints for now; nested types must have Storable + (AppT (ConT ''Storable) (ConT typeName)) + [ FunD 'sizeOf [Clause [WildP] (NormalB sizeOfBody) []], + FunD 'alignment [Clause [WildP] (NormalB alignBody) []], + FunD 'poke pokeBody, + FunD 'peek peekBody + ] + return [inst] + +-- | Generate sizeOf body: sum of all field sizes. +mkSizeOfBody :: [Type] -> Q Exp +mkSizeOfBody [] = litE (integerL 0) +mkSizeOfBody fieldTypes = do + sizeExps <- mapM mkFieldSizeOf fieldTypes + return $ foldl1 (\a b -> InfixE (Just a) (VarE '(+)) (Just b)) sizeExps + +-- | Generate size expression for a single field. +mkFieldSizeOf :: Type -> Q Exp +mkFieldSizeOf t = case primitiveSizeAlign t of + Just (size, _) -> litE (integerL (fromIntegral size)) + Nothing -> [|sizeOf (undefined :: $(return t))|] + +-- | Generate alignment body: max of field alignments. +mkAlignmentBody :: [Type] -> Q Exp +mkAlignmentBody [] = litE (integerL 1) +mkAlignmentBody fieldTypes = do + alignExps <- mapM mkFieldAlignment fieldTypes + return $ foldl1 mkMaxApp alignExps + where + mkMaxApp a = AppE (AppE (VarE 'max) a) + +-- | Generate alignment expression for a single field. +mkFieldAlignment :: Type -> Q Exp +mkFieldAlignment t = case primitiveSizeAlign t of + Just (_, align) -> litE (integerL (fromIntegral align)) + Nothing -> [|alignment (undefined :: $(return t))|] + +-- | Generate poke body. +mkPokeBody :: Name -> [Type] -> Int -> Q [Clause] +mkPokeBody conName fieldTypes fieldCount = do + ptrName <- newName "ptr" + varNames <- mapM (\i -> newName ("f" ++ show i)) [0 .. fieldCount - 1] + + let pat = ConP conName [] (map VarP varNames) + + -- Build offset calculations and poke statements + stmts <- mkPokeStmts ptrName varNames fieldTypes + + let body = DoE Nothing stmts + return [Clause [VarP ptrName, pat] (NormalB body) []] + +-- | Generate poke statements with offset tracking. +mkPokeStmts :: Name -> [Name] -> [Type] -> Q [Stmt] +mkPokeStmts ptrName varNames fieldTypes = go 0 (zip varNames fieldTypes) + where + go _ [] = return [] + go offset ((varName, fieldType) : rest) = do + stmt <- mkPokeStmt ptrName varName fieldType offset + restStmts <- case primitiveSizeAlign fieldType of + Just (size, _) -> go (offset + size) rest + Nothing -> goAfterNested ptrName offset fieldType rest + return (stmt : restStmts) + + -- After a nested type, offset must be calculated dynamically + goAfterNested _ _ _ [] = return [] + goAfterNested ptr baseOffset prevType ((varName, fieldType) : rest) = do + stmt <- mkPokeStmtDynamic ptr varName fieldType baseOffset prevType + restStmts <- case primitiveSizeAlign fieldType of + Just (size, _) -> goAfterNested ptr (baseOffset + size) fieldType rest + Nothing -> goAfterNested ptr baseOffset fieldType rest + return (stmt : restStmts) + +-- | Generate a single poke statement for a field (LE-aware for multi-byte primitives). +mkPokeStmt :: Name -> Name -> Type -> Int -> Q Stmt +mkPokeStmt ptrName varName fieldType offset + | isPrimitive fieldType = case lePokeName fieldType of + Nothing -> + -- Single-byte type, no endianness concern + NoBindS <$> [|pokeByteOff $(varE ptrName) offset $(varE varName)|] + Just pokeFn -> + -- Multi-byte primitive, use LE helper + NoBindS <$> [|$(varE pokeFn) $(varE ptrName) offset $(varE varName)|] + | otherwise = + -- Nested Storable type: poke at offset using castPtr + NoBindS <$> [|poke (castPtr ($(varE ptrName) `plusPtr` offset)) $(varE varName)|] + +-- | Generate poke with dynamic offset (after nested type). +mkPokeStmtDynamic :: Name -> Name -> Type -> Int -> Type -> Q Stmt +mkPokeStmtDynamic ptrName varName fieldType baseOffset prevType = do + let offsetExpr = [|baseOffset + sizeOf (undefined :: $(return prevType))|] + if isPrimitive fieldType + then case lePokeName fieldType of + Nothing -> + NoBindS <$> [|pokeByteOff $(varE ptrName) ($offsetExpr) $(varE varName)|] + Just pokeFn -> + NoBindS <$> [|$(varE pokeFn) $(varE ptrName) ($offsetExpr) $(varE varName)|] + else NoBindS <$> [|poke (castPtr ($(varE ptrName) `plusPtr` ($offsetExpr))) $(varE varName)|] + +-- | Generate peek body. +mkPeekBody :: Name -> [Type] -> Q [Clause] +mkPeekBody conName fieldTypes = do + ptrName <- newName "ptr" + + -- Build the peek expression using Applicative + peekExpr <- mkPeekExpr ptrName conName fieldTypes + + return [Clause [VarP ptrName] (NormalB peekExpr) []] + +-- | Generate peek expression using Applicative style. +-- Builds: pure Constructor <*> peek1 <*> peek2 <*> ... +mkPeekExpr :: Name -> Name -> [Type] -> Q Exp +mkPeekExpr _ conName [] = [|return $(return (ConE conName))|] +mkPeekExpr ptrName conName fieldTypes = do + -- Start with: pure Constructor + let startAcc = AppE (VarE 'pure) (ConE conName) + go 0 fieldTypes startAcc + where + go _ [] acc = return acc + go offset (t : ts) acc = do + peekField <- mkPeekField ptrName t offset + -- acc <*> peekField + let newAcc = InfixE (Just acc) (VarE '(<*>)) (Just peekField) + case primitiveSizeAlign t of + Just (size, _) -> go (offset + size) ts newAcc + Nothing -> goAfterNested ptrName offset t ts newAcc + + goAfterNested _ _ _ [] acc = return acc + goAfterNested ptr baseOffset prevType (t : ts) acc = do + peekField <- mkPeekFieldDynamic ptr t baseOffset prevType + let newAcc = InfixE (Just acc) (VarE '(<*>)) (Just peekField) + case primitiveSizeAlign t of + Just (size, _) -> + goAfterNested ptr (baseOffset + size) t ts newAcc + Nothing -> + goAfterNested ptr baseOffset t ts newAcc + +-- | Generate peek for a single field (LE-aware for multi-byte primitives). +mkPeekField :: Name -> Type -> Int -> Q Exp +mkPeekField ptrName fieldType offset + | isPrimitive fieldType = case lePeekName fieldType of + Nothing -> [|peekByteOff $(varE ptrName) offset|] + Just peekFn -> [|$(varE peekFn) $(varE ptrName) offset|] + | otherwise = [|peek (castPtr ($(varE ptrName) `plusPtr` offset))|] + +-- | Generate peek with dynamic offset. +mkPeekFieldDynamic :: Name -> Type -> Int -> Type -> Q Exp +mkPeekFieldDynamic ptrName fieldType baseOffset prevType = do + let offsetExpr = [|baseOffset + sizeOf (undefined :: $(return prevType))|] + if isPrimitive fieldType + then case lePeekName fieldType of + Nothing -> [|peekByteOff $(varE ptrName) ($offsetExpr)|] + Just peekFn -> [|$(varE peekFn) $(varE ptrName) ($offsetExpr)|] + else [|peek (castPtr ($(varE ptrName) `plusPtr` ($offsetExpr)))|] + +-- | Extract constructor name and field types. +conFieldTypes :: Con -> (Name, [Type]) +conFieldTypes (NormalC name fields) = (name, map snd fields) +conFieldTypes (RecC name fields) = (name, map (\(_, _, t) -> t) fields) +conFieldTypes (InfixC (_, t1) name (_, t2)) = (name, [t1, t2]) +conFieldTypes (ForallC _ _ con) = conFieldTypes con +conFieldTypes (GadtC [name] fields _) = (name, map snd fields) +conFieldTypes (RecGadtC [name] fields _) = (name, map (\(_, _, t) -> t) fields) +conFieldTypes _ = error "deriveStorable: unsupported constructor form" diff --git a/src/GBNet/Simulator.hs b/src/GBNet/Simulator.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Simulator.hs @@ -0,0 +1,212 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedLabels #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- Module : GBNet.Simulator +-- Description : Network condition simulator for testing +-- +-- Simulates packet loss, latency, jitter, duplicates, reordering, +-- and bandwidth limiting for testing purposes. +module GBNet.Simulator + ( -- * Simulator + NetworkSimulator (..), + newNetworkSimulator, + simulatorProcessSend, + simulatorReceiveReady, + simulatorPendingCount, + simulatorConfig, + + -- * Delayed packet + DelayedPacket (..), + ) +where + +import Control.Monad (when) +import Control.Monad.State.Strict (State, runState) +import qualified Data.ByteString as BS +import Data.Foldable (toList) +import Data.Sequence (Seq) +import qualified Data.Sequence as Seq +import Data.Word (Word64) +import GBNet.Class (MonoTime (..)) +import GBNet.Config (SimulationConfig (..)) +import GBNet.Reliability (elapsedMs) +import GBNet.Util (nextRandom, randomDouble) +import Optics ((&), (.~)) +import Optics.State (use) +import Optics.State.Operators ((%=), (.=)) +import Optics.TH (makeFieldLabelsNoPrefix) + +-- | Maximum extra delay (ms) applied to out-of-order packets. +outOfOrderMaxDelayMs :: Double +outOfOrderMaxDelayMs = 50.0 + +-- | Packets with total delay below this threshold (ms) are delivered immediately. +immediateDeliveryThresholdMs :: Double +immediateDeliveryThresholdMs = 1.0 + +-- | Maximum extra delay (ms) applied to duplicate packets. +duplicateJitterMs :: Double +duplicateJitterMs = 20.0 + +-- | A packet delayed for later delivery. +data DelayedPacket = DelayedPacket + { dpData :: !BS.ByteString, + dpAddr :: !Word64, -- Address key (simplified) + dpDeliverAt :: !MonoTime + } + deriving (Show) + +makeFieldLabelsNoPrefix ''DelayedPacket + +-- | Network condition simulator. +data NetworkSimulator = NetworkSimulator + { nsConfig :: !SimulationConfig, + nsDelayedPackets :: !(Seq DelayedPacket), + nsTokenBucketTokens :: !Double, + nsLastTokenRefill :: !MonoTime, + nsRngState :: !Word64 -- Simple LCG RNG state + } + deriving (Show) + +makeFieldLabelsNoPrefix ''NetworkSimulator + +-- | Create a new network simulator. +newNetworkSimulator :: SimulationConfig -> MonoTime -> NetworkSimulator +newNetworkSimulator config now = + NetworkSimulator + { nsConfig = config, + nsDelayedPackets = Seq.empty, + nsTokenBucketTokens = 0.0, + nsLastTokenRefill = now, + nsRngState = unMonoTime now -- Use time as initial seed + } + +-- | Process an outgoing packet through the simulator. +-- Returns (immediate packets, updated simulator). +simulatorProcessSend :: + BS.ByteString -> + Word64 -> + MonoTime -> + NetworkSimulator -> + ([(BS.ByteString, Word64)], NetworkSimulator) +simulatorProcessSend dat addr now = + runState $ do + r1 <- nextRandomS + r2 <- nextRandomS + r3 <- nextRandomS + r4 <- nextRandomS + config <- use #nsConfig + + let lost = checkLoss config r1 + if lost + then pure [] + else do + overBandwidth <- checkBandwidthS (BS.length dat) now config + if overBandwidth + then pure [] + else do + let totalDelayMs = calculateDelay config r2 r3 r4 + deliverAt = now + MonoTime (round (totalDelayMs * 1e6)) + immediate <- scheduleOrDeliverS dat addr deliverAt totalDelayMs + scheduleDuplicateS dat addr now totalDelayMs config + pure immediate + +-- | Check if a packet should be dropped due to simulated loss. +checkLoss :: SimulationConfig -> Word64 -> Bool +checkLoss config rng = + let threshold = simPacketLoss config + in threshold > 0.0 && randomDouble rng < threshold + +-- | Check if bandwidth limit is exceeded, consuming tokens if available. +checkBandwidthS :: Int -> MonoTime -> SimulationConfig -> State NetworkSimulator Bool +checkBandwidthS packetSize now config + | simBandwidthLimitBytesPerSec config <= 0 = pure False + | otherwise = do + refillTokensS now + tokens <- use #nsTokenBucketTokens + let needed = fromIntegral packetSize + if tokens < needed + then pure True + else do + #nsTokenBucketTokens .= (tokens - needed) + pure False + +-- | Calculate total delay including latency, jitter, and out-of-order reordering. +calculateDelay :: SimulationConfig -> Word64 -> Word64 -> Word64 -> Double +calculateDelay config rJitter rOoOChance rOoODelay = + let baseLatency = fromIntegral (simLatencyMs config) :: Double + jitter + | simJitterMs config > 0 = randomDouble rJitter * fromIntegral (simJitterMs config) + | otherwise = 0.0 + outOfOrderChance = simOutOfOrderChance config + extraDelay + | outOfOrderChance > 0.0 && randomDouble rOoOChance < outOfOrderChance = + randomDouble rOoODelay * outOfOrderMaxDelayMs + | otherwise = 0.0 + in baseLatency + jitter + extraDelay + +-- | Deliver a packet immediately or schedule it for delayed delivery. +scheduleOrDeliverS :: + BS.ByteString -> Word64 -> MonoTime -> Double -> State NetworkSimulator [(BS.ByteString, Word64)] +scheduleOrDeliverS dat addr deliverAt totalDelayMs + | totalDelayMs < immediateDeliveryThresholdMs = pure [(dat, addr)] + | otherwise = do + let delayed = DelayedPacket {dpData = dat, dpAddr = addr, dpDeliverAt = deliverAt} + #nsDelayedPackets %= (Seq.|> delayed) + pure [] + +-- | Maybe schedule a duplicate packet with extra jitter. +scheduleDuplicateS :: BS.ByteString -> Word64 -> MonoTime -> Double -> SimulationConfig -> State NetworkSimulator () +scheduleDuplicateS dat addr now baseDelayMs config = do + r <- nextRandomS + let dupChance = simDuplicateChance config + when (dupChance > 0.0 && randomDouble r < dupChance) $ do + let dupDeliverAt = now + MonoTime (round ((baseDelayMs + randomDouble r * duplicateJitterMs) * 1e6)) + dupPacket = DelayedPacket {dpData = dat, dpAddr = addr, dpDeliverAt = dupDeliverAt} + #nsDelayedPackets %= (Seq.|> dupPacket) + +-- | Retrieve packets ready for delivery. +-- Scans all delayed packets since they may not be sorted by delivery time. +simulatorReceiveReady :: MonoTime -> NetworkSimulator -> ([(BS.ByteString, Word64)], NetworkSimulator) +simulatorReceiveReady now sim = + let (ready, notReady) = Seq.partition (\pkt -> dpDeliverAt pkt <= now) (nsDelayedPackets sim) + results = map (\pkt -> (dpData pkt, dpAddr pkt)) (toList ready) + in (results, sim & #nsDelayedPackets .~ notReady) + +-- | Get count of pending delayed packets. +simulatorPendingCount :: NetworkSimulator -> Int +simulatorPendingCount = Seq.length . nsDelayedPackets + +-- | Get the simulation configuration. +simulatorConfig :: NetworkSimulator -> SimulationConfig +simulatorConfig = nsConfig +{-# INLINE simulatorConfig #-} + +-- | Refill token bucket based on elapsed time (State version). +refillTokensS :: MonoTime -> State NetworkSimulator () +refillTokensS now = do + sim <- use #nsLastTokenRefill + config <- use #nsConfig + tokens <- use #nsTokenBucketTokens + let elapsedSecs = elapsedMs sim now / 1000.0 + refillRate = fromIntegral (simBandwidthLimitBytesPerSec config) + newTokens = tokens + elapsedSecs * refillRate + maxTokens = refillRate -- Cap at 1 second worth + cappedTokens = min newTokens maxTokens + #nsTokenBucketTokens .= cappedTokens + #nsLastTokenRefill .= now + +-- | Stateful wrapper around the shared 'nextRandom' from GBNet.Util. +nextRandomS :: State NetworkSimulator Word64 +nextRandomS = do + s <- use #nsRngState + let (output, next) = nextRandom s + #nsRngState .= next + pure output diff --git a/src/GBNet/Socket.hs b/src/GBNet/Socket.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Socket.hs @@ -0,0 +1,129 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedLabels #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- Module : GBNet.Socket +-- Description : Non-blocking UDP socket wrapper with statistics +-- +-- Platform-agnostic non-blocking UDP socket wrapper with per-socket +-- statistics tracking. +module GBNet.Socket + ( -- * Constants + maxUdpPacketSize, + + -- * Socket errors + SocketError (..), + + -- * UDP socket + UdpSocket (..), + newUdpSocket, + closeSocket, + socketLocalAddr, + socketSendTo, + ) +where + +import Control.Exception (IOException, catch) +import qualified Data.ByteString as BS +import GBNet.Reliability (MonoTime) +import GBNet.Stats (SocketStats (..), defaultSocketStats) +import Network.Socket (SockAddr, Socket) +import qualified Network.Socket as NS +import qualified Network.Socket.ByteString as NSB +import Optics ((%), (%~), (&), (?~)) +import Optics.TH (makeFieldLabelsNoPrefix) + +-- | Maximum size of a single UDP datagram. +maxUdpPacketSize :: Int +maxUdpPacketSize = 65536 + +-- | Errors that can occur during socket operations. +data SocketError + = SocketIoError !String + | SocketInvalidAddress + | SocketClosed + | SocketOther !String + deriving (Eq, Show) + +-- | Non-blocking UDP socket with per-socket statistics. +data UdpSocket = UdpSocket + { usSocket :: !Socket, + usStats :: !SocketStats + } + deriving (Show) + +makeFieldLabelsNoPrefix ''UdpSocket + +-- | Create a new UDP socket bound to the specified address. +-- Uses bracket-style cleanup to avoid leaking the socket FD if bind fails. +newUdpSocket :: SockAddr -> IO (Either SocketError UdpSocket) +newUdpSocket addr = do + sockResult <- tryIO $ NS.socket NS.AF_INET NS.Datagram NS.defaultProtocol + case sockResult of + Left err -> return $ Left (SocketIoError err) + Right sock -> do + bindResult <- tryIO $ do + NS.setSocketOption sock NS.ReuseAddr 1 + NS.bind sock addr + NS.withFdSocket sock NS.setNonBlockIfNeeded + case bindResult of + Left err -> do + NS.close sock -- Clean up on failure + return $ Left (SocketIoError err) + Right () -> + return $ + Right + UdpSocket + { usSocket = sock, + usStats = defaultSocketStats + } + +-- | Close the socket. +closeSocket :: UdpSocket -> IO () +closeSocket sock = NS.close (usSocket sock) + +-- | Get the local address this socket is bound to. +socketLocalAddr :: UdpSocket -> IO (Either SocketError SockAddr) +socketLocalAddr sock = do + result <- tryIO $ NS.getSocketName (usSocket sock) + return $ case result of + Left err -> Left (SocketIoError err) + Right addr -> Right addr + +-- | Send data to a specific address. +socketSendTo :: + BS.ByteString -> + SockAddr -> + MonoTime -> + UdpSocket -> + IO (Either SocketError (Int, UdpSocket)) +socketSendTo dat addr now sock = do + result <- tryIO $ NSB.sendTo (usSocket sock) dat addr + return $ case result of + Left err -> Left (SocketIoError err) + Right sent -> + let sock' = + sock + & #usStats + % #ssBytesSent + %~ (+ fromIntegral sent) + & #usStats + % #ssPacketsSent + %~ (+ 1) + & (#usStats % #ssLastSendTime) + ?~ now + in Right (sent, sock') + +-- | Try an IO action, catching IOExceptions. +tryIO :: IO a -> IO (Either String a) +tryIO action = + (Right <$> action) `catch` handler + where + handler :: IOException -> IO (Either String a) + handler e = return $ Left (show e) diff --git a/src/GBNet/Stats.hs b/src/GBNet/Stats.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Stats.hs @@ -0,0 +1,139 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- Module : GBNet.Stats +-- Description : Network statistics tracking +-- +-- Provides 'NetworkStats' for tracking connection health metrics. +module GBNet.Stats + ( -- * Connection quality + ConnectionQuality (..), + assessConnectionQuality, + + -- * Congestion level + CongestionLevel (..), + + -- * Network statistics + NetworkStats (..), + defaultNetworkStats, + + -- * Socket statistics + SocketStats (..), + defaultSocketStats, + ) +where + +import Data.Word (Word64) +import GBNet.Reliability (MonoTime) +import Optics.TH (makeFieldLabelsNoPrefix) + +-- | Connection quality assessment. +data ConnectionQuality + = QualityExcellent + | QualityGood + | QualityFair + | QualityPoor + | QualityBad + deriving (Eq, Show, Ord) + +-- Quality thresholds: (lossPercent, rttMs) +badLossThreshold, poorLossThreshold, fairLossThreshold, goodLossThreshold :: Double +badLossThreshold = 10.0 +poorLossThreshold = 5.0 +fairLossThreshold = 2.0 +goodLossThreshold = 0.5 + +badRttThreshold, poorRttThreshold, fairRttThreshold, goodRttThreshold :: Double +badRttThreshold = 500.0 +poorRttThreshold = 250.0 +fairRttThreshold = 150.0 +goodRttThreshold = 80.0 + +-- | Assess connection quality from RTT and packet loss. +assessConnectionQuality :: Double -> Double -> ConnectionQuality +assessConnectionQuality rttMs lossPercent + | lossPercent > badLossThreshold || rttMs > badRttThreshold = QualityBad + | lossPercent > poorLossThreshold || rttMs > poorRttThreshold = QualityPoor + | lossPercent > fairLossThreshold || rttMs > fairRttThreshold = QualityFair + | lossPercent > goodLossThreshold || rttMs > goodRttThreshold = QualityGood + | otherwise = QualityExcellent + +-- | Congestion pressure level reported to the application. +-- +-- Applications can use this to adapt their behavior: +-- +-- * 'CongestionNone' — send freely +-- * 'CongestionElevated' — consider reducing non-essential traffic +-- * 'CongestionHigh' — drop low-priority data, reduce send rate +-- * 'CongestionCritical' — sends are being suppressed, only send essential data +data CongestionLevel + = CongestionNone + | CongestionElevated + | CongestionHigh + | CongestionCritical + deriving (Eq, Show, Ord) + +-- | Network statistics for a connection. +data NetworkStats = NetworkStats + { nsPacketsSent :: !Word64, + nsPacketsReceived :: !Word64, + nsBytesSent :: !Word64, + nsBytesReceived :: !Word64, + nsRtt :: !Double, + nsPacketLoss :: !Double, + nsBandwidthUp :: !Double, + nsBandwidthDown :: !Double, + nsConnectionQuality :: !ConnectionQuality, + nsCongestionLevel :: !CongestionLevel + } + deriving (Show) + +-- | Default (zero) network statistics. +defaultNetworkStats :: NetworkStats +defaultNetworkStats = + NetworkStats + { nsPacketsSent = 0, + nsPacketsReceived = 0, + nsBytesSent = 0, + nsBytesReceived = 0, + nsRtt = 0.0, + nsPacketLoss = 0.0, + nsBandwidthUp = 0.0, + nsBandwidthDown = 0.0, + nsConnectionQuality = QualityExcellent, + nsCongestionLevel = CongestionNone + } + +-- | Socket-level statistics. +data SocketStats = SocketStats + { ssPacketsSent :: !Word64, + ssPacketsReceived :: !Word64, + ssBytesSent :: !Word64, + ssBytesReceived :: !Word64, + ssCrcDrops :: !Word64, + ssLastSendTime :: !(Maybe MonoTime), + ssLastReceiveTime :: !(Maybe MonoTime) + } + deriving (Show) + +-- | Default socket statistics. +defaultSocketStats :: SocketStats +defaultSocketStats = + SocketStats + { ssPacketsSent = 0, + ssPacketsReceived = 0, + ssBytesSent = 0, + ssBytesReceived = 0, + ssCrcDrops = 0, + ssLastSendTime = Nothing, + ssLastReceiveTime = Nothing + } + +makeFieldLabelsNoPrefix ''NetworkStats +makeFieldLabelsNoPrefix ''SocketStats diff --git a/src/GBNet/TestNet.hs b/src/GBNet/TestNet.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/TestNet.hs @@ -0,0 +1,282 @@ +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE DeriveFunctor #-} +{-# LANGUAGE FlexibleContexts #-} +{-# LANGUAGE FlexibleInstances #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} +{-# LANGUAGE MultiParamTypeClasses #-} +{-# LANGUAGE OverloadedLabels #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} +{-# LANGUAGE TypeOperators #-} +{-# LANGUAGE UndecidableInstances #-} + +-- | +-- Module : GBNet.TestNet +-- Description : Pure deterministic network for testing +-- +-- A simulated network that runs entirely in pure code. +-- Useful for: +-- - Unit testing connection logic +-- - Property-based testing +-- - Deterministic replay +-- - Network condition simulation +module GBNet.TestNet + ( -- * Test network monad + TestNet (..), + runTestNet, + + -- * Test network state + InFlightPacket (..), + TestNetConfig (..), + TestNetState (..), + initialTestNetState, + defaultTestNetConfig, + + -- * Operations + advanceTime, + simulateLatency, + simulateLoss, + getPendingPackets, + + -- * Multi-peer world simulation + TestWorld (..), + newTestWorld, + runPeerInWorld, + deliverPackets, + worldAdvanceTime, + ) +where + +import Control.Monad.State.Strict (MonadState, State, get, runState) +import Data.ByteString (ByteString) +import Data.Foldable (toList) +import Data.List (partition) +import qualified Data.Map.Strict as Map +import Data.Sequence (Seq) +import qualified Data.Sequence as Seq +import Data.Word (Word64) +import GBNet.Class (MonadNetwork (..), MonadTime (..), MonoTime (..), NetError (..)) +import GBNet.Security (validateAndStripCrc32) +import GBNet.Util (nextRandom, randomDouble) +import Network.Socket (SockAddr) +import Optics ((%), (%~), (&), (.~)) +import Optics.State (use) +import Optics.State.Operators ((%=), (.=)) +import Optics.TH (makeFieldLabelsNoPrefix) + +-- | A packet in transit. +data InFlightPacket = InFlightPacket + { ifpFrom :: !SockAddr, + ifpTo :: !SockAddr, + ifpData :: !ByteString, + ifpDeliverAt :: !MonoTime -- When this packet should be delivered + } + deriving (Show) + +makeFieldLabelsNoPrefix ''InFlightPacket + +-- | Nanoseconds per millisecond. +nsPerMs :: Word64 +nsPerMs = 1000000 + +-- | Test network configuration. +data TestNetConfig = TestNetConfig + { tncLatencyNs :: !Word64, -- Simulated one-way latency in nanoseconds + tncLossRate :: !Double, -- Packet loss probability (0.0 - 1.0) + tncJitterNs :: !Word64 -- Random jitter range in nanoseconds + } + deriving (Show) + +makeFieldLabelsNoPrefix ''TestNetConfig + +-- | Default test config (no latency, no loss). +defaultTestNetConfig :: TestNetConfig +defaultTestNetConfig = + TestNetConfig + { tncLatencyNs = 0, + tncLossRate = 0.0, + tncJitterNs = 0 + } + +-- | State of the test network. +data TestNetState = TestNetState + { tnsCurrentTime :: !MonoTime, + tnsLocalAddr :: !SockAddr, + tnsInFlight :: !(Seq InFlightPacket), -- Packets in transit + tnsInbox :: !(Seq (ByteString, SockAddr)), -- Delivered packets + tnsConfig :: !TestNetConfig, + tnsRng :: !Word64, + tnsClosed :: !Bool + } + deriving (Show) + +makeFieldLabelsNoPrefix ''TestNetState + +-- | Create initial test network state. +initialTestNetState :: SockAddr -> TestNetState +initialTestNetState localAddr = + TestNetState + { tnsCurrentTime = 0, + tnsLocalAddr = localAddr, + tnsInFlight = Seq.empty, + tnsInbox = Seq.empty, + tnsConfig = defaultTestNetConfig, + tnsRng = 42, -- Deterministic seed + tnsClosed = False + } + +-- | The test network monad. +newtype TestNet a = TestNet {unTestNet :: State TestNetState a} + deriving (Functor, Applicative, Monad, MonadState TestNetState) + +-- | Run a test network computation. +runTestNet :: TestNet a -> TestNetState -> (a, TestNetState) +runTestNet (TestNet m) = runState m + +instance MonadTime TestNet where + getMonoTime = use #tnsCurrentTime + +instance MonadNetwork TestNet where + netSend toAddr bytes = do + st <- get + if tnsClosed st + then pure (Left NetSocketClosed) + else do + let cfg = tnsConfig st + (r1, rng') = nextRandom (tnsRng st) + if randomDouble r1 < tncLossRate cfg + then do + #tnsRng .= rng' + pure (Right ()) + else do + let (r2, rng'') = nextRandom rng' + jitterRange = tncJitterNs cfg + jitter = if jitterRange == 0 then 0 else r2 `mod` (jitterRange + 1) + deliverAt = tnsCurrentTime st + MonoTime (tncLatencyNs cfg) + MonoTime jitter + pkt = + InFlightPacket + { ifpFrom = tnsLocalAddr st, + ifpTo = toAddr, + ifpData = bytes, + ifpDeliverAt = deliverAt + } + #tnsInFlight %= (Seq.|> pkt) + #tnsRng .= rng'' + pure (Right ()) + + netRecv = do + inbox <- use #tnsInbox + case Seq.viewl inbox of + Seq.EmptyL -> pure Nothing + (bytes, from) Seq.:< rest -> do + #tnsInbox .= rest + -- Validate and strip CRC, matching the IO backend behavior + case validateAndStripCrc32 bytes of + Nothing -> pure Nothing -- Drop corrupt packets + Just validated -> pure (Just (validated, from)) + + netClose = #tnsClosed .= True + +-- | Advance time and deliver packets that are ready. +advanceTime :: MonoTime -> TestNet () +advanceTime newTime = do + st <- get + let (ready, stillInFlight) = + Seq.partition (\p -> ifpDeliverAt p <= newTime) (tnsInFlight st) + -- Only deliver packets addressed to us + delivered = + (\p -> (ifpData p, ifpFrom p)) + <$> Seq.filter (\p -> ifpTo p == tnsLocalAddr st) ready + #tnsCurrentTime .= newTime + #tnsInFlight .= stillInFlight + #tnsInbox %= (Seq.>< delivered) + +-- | Configure simulated one-way latency (in milliseconds). +simulateLatency :: Word64 -> TestNet () +simulateLatency ms = #tnsConfig % #tncLatencyNs .= ms * nsPerMs + +-- | Configure simulated packet loss. +simulateLoss :: Double -> TestNet () +simulateLoss rate = #tnsConfig % #tncLossRate .= rate + +-- | Get all packets currently in flight (for debugging/testing). +getPendingPackets :: TestNet [InFlightPacket] +getPendingPackets = toList <$> use #tnsInFlight + +-- ----------------------------------------------------------------------------- +-- Multi-peer world simulation +-- ----------------------------------------------------------------------------- + +-- | A world containing multiple peers for integration testing. +-- Allows deterministic simulation of multi-peer scenarios. +data TestWorld = TestWorld + { -- | State for each peer, keyed by their address + twPeers :: !(Map.Map SockAddr TestNetState), + -- | Global simulation time + twGlobalTime :: !MonoTime + } + deriving (Show) + +makeFieldLabelsNoPrefix ''TestWorld + +-- | Create a new test world with no peers. +newTestWorld :: TestWorld +newTestWorld = + TestWorld + { twPeers = Map.empty, + twGlobalTime = 0 + } + +-- | Run a TestNet operation for a specific peer in the world. +-- If the peer doesn't exist, it's created with the given address. +runPeerInWorld :: SockAddr -> TestNet a -> TestWorld -> (a, TestWorld) +runPeerInWorld addr action world = + let peerState = + Map.findWithDefault + (initialTestNetState addr) {tnsCurrentTime = twGlobalTime world} + addr + (twPeers world) + (result, peerState') = runTestNet action peerState + world' = world & #twPeers %~ Map.insert addr peerState' + in (result, world') + +-- | Deliver all ready packets between peers. +-- Packets are moved from sender's outFlight to receiver's inbox +-- if the delivery time has passed. +deliverPackets :: TestWorld -> TestWorld +deliverPackets world = + let time = twGlobalTime world + -- Collect all packets from all peers + allPackets = + concatMap + (toList . tnsInFlight) + (Map.elems (twPeers world)) + -- Partition into ready and not ready + (ready, notReady) = partition (\p -> ifpDeliverAt p <= time) allPackets + -- Clear all peer outFlight, then redistribute + clearedPeers = Map.map (\ps -> ps & #tnsInFlight .~ Seq.empty) (twPeers world) + -- Put not-ready packets back to their senders + peersWithPending = foldr putBackPending clearedPeers notReady + -- Deliver ready packets to their destinations + peersWithDelivered = foldr deliverOne peersWithPending ready + in world & #twPeers .~ peersWithDelivered + where + putBackPending pkt = + Map.adjust + (\ps -> ps & #tnsInFlight %~ (Seq.|> pkt)) + (ifpFrom pkt) + + deliverOne pkt = + Map.adjust + (\ps -> ps & #tnsInbox %~ (Seq.|> (ifpData pkt, ifpFrom pkt))) + (ifpTo pkt) + +-- | Advance time for all peers in the world. +worldAdvanceTime :: MonoTime -> TestWorld -> TestWorld +worldAdvanceTime newTime world = + let world' = world & #twGlobalTime .~ newTime + updatedPeers = + Map.map + (\ps -> ps & #tnsCurrentTime .~ newTime) + (twPeers world') + in deliverPackets (world' & #twPeers .~ updatedPeers) diff --git a/src/GBNet/Types.hs b/src/GBNet/Types.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Types.hs @@ -0,0 +1,48 @@ +{-# LANGUAGE DerivingStrategies #-} +{-# LANGUAGE GeneralizedNewtypeDeriving #-} + +-- | +-- Module : GBNet.Types +-- Description : Domain-specific newtypes for type safety +-- +-- Wraps raw primitives in newtypes so the compiler catches type mix-ups +-- at compile time. Zero runtime cost via newtype erasure. +module GBNet.Types + ( -- * Channel identifier + ChannelId (..), + channelIdToInt, + + -- * Sequence number + SequenceNum (..), + + -- * Message identifier + MessageId (..), + ) +where + +import Control.DeepSeq (NFData) +import Data.Word (Word16, Word32, Word8) +import Foreign.Storable (Storable (..)) + +-- | Channel identifier (stored as Word8). +-- Not a quantity — no 'Num' instance. Construct via @ChannelId 0@. +newtype ChannelId = ChannelId {unChannelId :: Word8} + deriving stock (Eq, Ord, Show) + deriving newtype (Bounded, NFData, Storable) + +-- | Convert a 'ChannelId' to 'Int' for IntMap indexing. +channelIdToInt :: ChannelId -> Int +channelIdToInt = fromIntegral . unChannelId +{-# INLINE channelIdToInt #-} + +-- | Sequence number with Word16 wraparound arithmetic. +-- Derives 'Num' because @seq + 1@ is the canonical operation. +newtype SequenceNum = SequenceNum {unSequenceNum :: Word16} + deriving stock (Eq, Ord, Show) + deriving newtype (Bounded, Enum, Integral, NFData, Num, Real, Storable) + +-- | Fragment message identifier. +-- Derives 'Num' for increment patterns in fragmentation. +newtype MessageId = MessageId {unMessageId :: Word32} + deriving stock (Eq, Ord, Show) + deriving newtype (Bounded, NFData, Num, Storable) diff --git a/src/GBNet/Util.hs b/src/GBNet/Util.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/Util.hs @@ -0,0 +1,66 @@ +-- | Shared utilities: sequence number wraparound and deterministic RNG. +module GBNet.Util + ( sequenceHalfRange, + sequenceGreaterThan, + sequenceDiff, + + -- * Deterministic RNG (LCG + SplitMix output mixing) + nextRandom, + randomDouble, + ) +where + +import Data.Bits (shiftR, xor, (.&.)) +import Data.Int (Int32) +import Data.Word (Word16, Word64) +import GBNet.Types (SequenceNum (..)) + +-- | Half the Word16 sequence space, used for wraparound comparison. +sequenceHalfRange :: Word16 +sequenceHalfRange = 32768 + +-- | Compares two sequence numbers accounting for Word16 wraparound. +-- Returns True if s1 is "greater than" s2 in circular sequence space. +sequenceGreaterThan :: SequenceNum -> SequenceNum -> Bool +sequenceGreaterThan (SequenceNum s1) (SequenceNum s2) = + ((s1 > s2) && (s1 - s2 <= sequenceHalfRange)) + || ((s1 < s2) && (s2 - s1 > sequenceHalfRange)) +{-# INLINE sequenceGreaterThan #-} + +-- | Computes the signed difference between two sequence numbers, +-- accounting for Word16 wraparound. +sequenceDiff :: SequenceNum -> SequenceNum -> Int32 +sequenceDiff (SequenceNum s1) (SequenceNum s2) + | diff > half = diff - full + | diff < -half = diff + full + | otherwise = diff + where + half = fromIntegral sequenceHalfRange :: Int32 + full = half * 2 + diff = fromIntegral s1 - fromIntegral s2 :: Int32 +{-# INLINE sequenceDiff #-} + +-- | LCG constants. +lcgMultiplier, lcgIncrement :: Word64 +lcgMultiplier = 6364136223846793005 +lcgIncrement = 1442695040888963407 + +-- | SplitMix-style random number generator. +-- Returns (output, nextState). Pure, deterministic, no IO. +nextRandom :: Word64 -> (Word64, Word64) +nextRandom s = + let next = lcgMultiplier * s + lcgIncrement + -- SplitMix-style output mixing (state is not exposed) + z0 = next `xor` (next `shiftR` 30) + z1 = z0 * 0xBF58476D1CE4E5B9 + z2 = z1 `xor` (z1 `shiftR` 27) + z3 = z2 * 0x94D049BB133111EB + output = z3 `xor` (z3 `shiftR` 31) + in (output, next) +{-# INLINE nextRandom #-} + +-- | Convert random Word64 to double in [0, 1). +-- Uses lower 32 bits; SplitMix output mixing ensures all bits are well-distributed. +randomDouble :: Word64 -> Double +randomDouble w = fromIntegral (w .&. 0xFFFFFFFF) / 4294967296.0 +{-# INLINE randomDouble #-} diff --git a/src/GBNet/ZeroCopy.hs b/src/GBNet/ZeroCopy.hs new file mode 100644 --- /dev/null +++ b/src/GBNet/ZeroCopy.hs @@ -0,0 +1,94 @@ +{-# LANGUAGE RankNTypes #-} +{-# LANGUAGE ScopedTypeVariables #-} + +-- | +-- Module : GBNet.ZeroCopy +-- Description : Zero-copy vector mutation via ST +-- +-- Internal module providing zero-overhead in-place mutation of immutable +-- vectors. Safety is enforced through two mechanisms: +-- +-- 1. __ST rank-2 type__: the @forall s@ prevents the mutable vector +-- from escaping the callback — compiler-enforced, not convention. +-- +-- 2. __Module opacity__: this module is internal. All buffer types +-- (@SequenceBuffer@, @SentPacketBuffer@, etc.) are opaque, so +-- external code cannot access the raw vectors. +-- +-- The @unsafeThaw@\/@unsafeFreeze@ pair are pointer casts (zero memcpy), +-- giving C-level mutation performance with Haskell's type-level guarantees. +module GBNet.ZeroCopy + ( -- * Unboxed vectors + zeroCopyMutateU, + zeroCopyMutateU', + + -- * Boxed vectors + zeroCopyMutate, + zeroCopyMutate', + ) +where + +import Control.Monad.ST (ST, runST) +import qualified Data.Vector as V +import qualified Data.Vector.Mutable as MV +import Data.Vector.Unboxed (Unbox) +import qualified Data.Vector.Unboxed as VU +import qualified Data.Vector.Unboxed.Mutable as VUM + +-- | Zero-copy mutation of an unboxed vector. +-- +-- @unsafeThaw@\/@unsafeFreeze@ are pointer casts — zero allocation. +-- The @forall s@ prevents the mutable vector from escaping the callback. +zeroCopyMutateU :: + forall a. + (Unbox a) => + VU.Vector a -> + (forall s. VUM.MVector s a -> ST s ()) -> + VU.Vector a +zeroCopyMutateU vec action = runST $ do + mv <- VU.unsafeThaw vec + action mv + VU.unsafeFreeze mv +{-# INLINE zeroCopyMutateU #-} + +-- | Zero-copy mutation of an unboxed vector, returning a result. +zeroCopyMutateU' :: + forall a b. + (Unbox a) => + VU.Vector a -> + (forall s. VUM.MVector s a -> ST s b) -> + (VU.Vector a, b) +zeroCopyMutateU' vec action = runST $ do + mv <- VU.unsafeThaw vec + result <- action mv + v' <- VU.unsafeFreeze mv + return (v', result) +{-# INLINE zeroCopyMutateU' #-} + +-- | Zero-copy mutation of a boxed vector. +-- +-- @unsafeThaw@\/@unsafeFreeze@ are pointer casts — zero allocation. +-- The @forall s@ prevents the mutable vector from escaping the callback. +zeroCopyMutate :: + forall a. + V.Vector a -> + (forall s. MV.MVector s a -> ST s ()) -> + V.Vector a +zeroCopyMutate vec action = runST $ do + mv <- V.unsafeThaw vec + action mv + V.unsafeFreeze mv +{-# INLINE zeroCopyMutate #-} + +-- | Zero-copy mutation of a boxed vector, returning a result. +zeroCopyMutate' :: + forall a b. + V.Vector a -> + (forall s. MV.MVector s a -> ST s b) -> + (V.Vector a, b) +zeroCopyMutate' vec action = runST $ do + mv <- V.unsafeThaw vec + result <- action mv + v' <- V.unsafeFreeze mv + return (v', result) +{-# INLINE zeroCopyMutate' #-} diff --git a/test/Main.hs b/test/Main.hs new file mode 100644 --- /dev/null +++ b/test/Main.hs @@ -0,0 +1,1827 @@ +{-# LANGUAGE BangPatterns #-} +{-# LANGUAGE DataKinds #-} +{-# LANGUAGE OverloadedStrings #-} +{-# LANGUAGE ScopedTypeVariables #-} +{-# LANGUAGE TemplateHaskell #-} +{-# LANGUAGE TypeFamilies #-} + +module Main where + +import Data.Bits ((.&.)) +import Data.List (foldl') +import qualified Data.ByteString as BS +import Data.Word (Word16, Word32, Word64, Word8) +import GBNet.Channel +import GBNet.Class () +import GBNet.Config + ( ConfigError (..), + NetworkConfig (..), + SimulationConfig (..), + defaultNetworkConfig, + defaultSimulationConfig, + maxChannelCount, + validateConfig, + ) +import GBNet.Congestion +import GBNet.Connection + ( Connection (..), + ConnectionError (..), + ConnectionState (..), + DisconnectReason (..), + connect, + connectionState, + createHeader, + markConnected, + newConnection, + sendMessage, + ) +import GBNet.Fragment +import GBNet.Packet +import GBNet.Peer +import GBNet.Reliability +import GBNet.Replication.Delta +import GBNet.Replication.Interest +import GBNet.Replication.Interpolation +import GBNet.Replication.Priority +import GBNet.Security +import GBNet.Serialize (deserialize, serialize) +import GBNet.Serialize.TH (deriveStorable) +import GBNet.Simulator +import GBNet.Socket (UdpSocket (..)) +import GBNet.Stats (CongestionLevel (..), defaultSocketStats) +import GBNet.TestNet +import GBNet.Types (ChannelId (..), MessageId (..), SequenceNum (..)) +import GBNet.Util +import Network.Socket (SockAddr (..), tupleToHostAddress) +import qualified Network.Socket as NS +import Test.QuickCheck (Arbitrary (..), Property, elements, quickCheck, withMaxSuccess, (==>)) + +-------------------------------------------------------------------------------- +-- TH-derived test types (Storable) +-------------------------------------------------------------------------------- + +data Vec3 = Vec3 + { vecX :: !Float, + vecY :: !Float, + vecZ :: !Float + } + deriving (Eq, Show) + +deriveStorable ''Vec3 + +instance Arbitrary Vec3 where + arbitrary = Vec3 <$> arbitrary <*> arbitrary <*> arbitrary + +instance Arbitrary PacketType where + arbitrary = elements [minBound .. maxBound] + +instance Arbitrary SequenceNum where + arbitrary = SequenceNum <$> arbitrary + +instance Arbitrary PacketHeader where + arbitrary = + PacketHeader + <$> arbitrary + <*> arbitrary + <*> arbitrary + <*> arbitrary + +instance Arbitrary MessageId where + arbitrary = MessageId <$> arbitrary + +instance Arbitrary FragmentHeader where + arbitrary = + FragmentHeader + <$> arbitrary + <*> arbitrary + <*> arbitrary + +instance Arbitrary TestDeltaState where + arbitrary = TestDeltaState <$> arbitrary <*> arbitrary + +-------------------------------------------------------------------------------- +-- Delta test types (Storable-based) +-------------------------------------------------------------------------------- + +-- Simple test type for delta compression: a pair of Word8 values. +data TestDeltaState = TestDeltaState !Word8 !Word8 + deriving (Eq, Show) + +deriveStorable ''TestDeltaState + +data TestDeltaDelta = TestDeltaDelta !Word8 !Word8 !Word8 !Word8 + -- Flags + values: hasA, hasB, valA, valB + deriving (Eq, Show) + +deriveStorable ''TestDeltaDelta + +instance NetworkDelta TestDeltaState where + type Delta TestDeltaState = TestDeltaDelta + diff (TestDeltaState a1 b1) (TestDeltaState a2 b2) = + TestDeltaDelta + (if a1 /= a2 then 1 else 0) + (if b1 /= b2 then 1 else 0) + a1 + b1 + apply (TestDeltaState a b) (TestDeltaDelta hasA hasB valA valB) = + TestDeltaState + (if hasA /= 0 then valA else a) + (if hasB /= 0 then valB else b) + +-------------------------------------------------------------------------------- +-- Main +-------------------------------------------------------------------------------- + +main :: IO () +main = do + putStrLn "=== GB-Net Haskell Tests ===" + putStrLn "" + + -- Storable serialization + testStorableRoundTrip + testPacketHeaderRoundTrip + + -- Reliability module + testSequenceGreaterThan + testSequenceDiff + testSequenceAtBoundaries + testSequenceBufferOps + testSequenceBufferWraparound + testSequenceBufferCollision + testRttConvergence + testAdaptiveRto + testPacketLossTracking + testAckBitsNoFalseAck + testProcessAcksReturnsChannelInfo + testInFlightEviction + testFastRetransmit + + -- Congestion control + testBinaryCongestionModeTransition + testBinaryRateRecovery + testCwndSlowStart + testCwndLossHalvesCwnd + testCwndSlowStartRestart + testCongestionLevelBinary + testCongestionLevelWindow + testBatchAndUnbatch + + -- Integration: Pure peer API + testPeerHandshake + testPeerMessageDelivery + testPeerDisconnect + testPeerConnectionTimeout + testPeerMaxClients + + -- Integration: Full TestNet polymorphic lifecycle + testTestNetHandshake + testTestNetMessageRoundTrip + + -- Replication: Interest management + testRadiusInterestRelevant + testRadiusInterestPriority + testGridInterestRelevant + + -- Replication: Priority accumulator + testPriorityAccumulate + testPriorityDrain + testPriorityUnregister + + -- Replication: Snapshot interpolation + testSnapshotPushAndReady + testSnapshotInterpolation + testSnapshotOutOfOrder + + -- Property-based tests + testPropertyRoundTrips + + -- Adversarial: Malformed packet handling + testTruncatedPacket + testGarbagePayload + testEmptyPacket + + -- Integration: Connection migration + testConnectionMigration + + -- Channel delivery modes, errors, retransmit + testChannelSendBufferFull + testChannelSendOversized + testChannelUnreliableDelivery + testChannelReliableOrderedDelivery + testChannelReliableSequencedDropOld + testChannelRetransmit + + -- Fragment: split, reassemble, header roundtrip, cleanup, too-large + testFragmentSplitReassemble + testFragmentHeaderRoundTrip + testFragmentCleanupExpiry + testFragmentTooLarge + + -- Security: CRC32C, rate limiting, token validation + testCrc32Roundtrip + testCrc32RejectCorrupt + testRateLimiterAllow + testRateLimiterDeny + testTokenValidation + testTokenExpired + testTokenReplayed + + -- Connection state machine + testConnectionStateMachine + testConnectionSendReceive + + -- Config validation + testValidateConfigValid + testValidateConfigErrors + + -- Delta encode/decode + testDeltaEncodeDecodeTrivial + + -- Simulator + testSimulatorBasic + testSimulatorPeerDelivery + + putStrLn "" + putStrLn "All tests passed!" + +-------------------------------------------------------------------------------- +-- Helpers +-------------------------------------------------------------------------------- + +assertEqual :: (Eq a, Show a) => String -> a -> a -> IO () +assertEqual name expected actual = + if expected == actual + then putStrLn $ " PASS: " ++ name + else + error $ + " FAIL: " + ++ name + ++ " expected " + ++ show expected + ++ " got " + ++ show actual + +-------------------------------------------------------------------------------- +-- Storable serialization tests +-------------------------------------------------------------------------------- + +testStorableRoundTrip :: IO () +testStorableRoundTrip = do + putStrLn "Storable round-trip:" + let v = Vec3 1.0 (-2.5) 100.0 + bytes = serialize v + assertEqual "Vec3 size" 12 (BS.length bytes) + case deserialize bytes :: Either String Vec3 of + Left err -> error $ " FAIL: deserialize Vec3: " ++ err + Right v' -> assertEqual "Vec3 roundtrip" v v' + +testPacketHeaderRoundTrip :: IO () +testPacketHeaderRoundTrip = do + putStrLn "PacketHeader roundtrip:" + let headers = + [ PacketHeader Payload (SequenceNum 42) (SequenceNum 40) 0xDEADBEEF, + PacketHeader ConnectionRequest (SequenceNum 0) (SequenceNum 0) 0, + PacketHeader Keepalive (SequenceNum 65535) (SequenceNum 65535) 0xFFFFFFFF, + PacketHeader Disconnect (SequenceNum 12345) (SequenceNum 54321) 0x12345678 + ] + mapM_ testHeader headers + where + testHeader hdr = do + let bytes = serializeHeader hdr + assertEqual ("size for " ++ show (packetType hdr)) packetHeaderByteSize (BS.length bytes) + case deserializeHeader bytes of + Left err -> error $ "deserialize failed: " ++ err + Right hdr' -> do + assertEqual "roundtrip packetType" (packetType hdr) (packetType hdr') + assertEqual "roundtrip sequenceNum" (sequenceNum hdr) (sequenceNum hdr') + assertEqual "roundtrip ack" (ack hdr) (ack hdr') + assertEqual "roundtrip ackBitfield" (ackBitfield hdr) (ackBitfield hdr') + +-------------------------------------------------------------------------------- +-- Reliability module tests +-------------------------------------------------------------------------------- + +testSequenceGreaterThan :: IO () +testSequenceGreaterThan = do + putStrLn "sequenceGreaterThan:" + assertEqual "1 > 0" True (sequenceGreaterThan (SequenceNum 1) (SequenceNum 0)) + assertEqual "0 > 1" False (sequenceGreaterThan (SequenceNum 0) (SequenceNum 1)) + assertEqual "100 > 50" True (sequenceGreaterThan (SequenceNum 100) (SequenceNum 50)) + assertEqual "50 > 100" False (sequenceGreaterThan (SequenceNum 50) (SequenceNum 100)) + -- Wraparound + assertEqual "0 > 65535" True (sequenceGreaterThan (SequenceNum 0) (SequenceNum 65535)) + assertEqual "65535 > 0" False (sequenceGreaterThan (SequenceNum 65535) (SequenceNum 0)) + assertEqual "1 > 65534" True (sequenceGreaterThan (SequenceNum 1) (SequenceNum 65534)) + assertEqual "100 > 65500" True (sequenceGreaterThan (SequenceNum 100) (SequenceNum 65500)) + +testSequenceDiff :: IO () +testSequenceDiff = do + putStrLn "sequenceDiff:" + assertEqual "diff(5,3)" 2 (sequenceDiff (SequenceNum 5) (SequenceNum 3)) + assertEqual "diff(3,5)" (-2) (sequenceDiff (SequenceNum 3) (SequenceNum 5)) + assertEqual "diff(100,100)" 0 (sequenceDiff (SequenceNum 100) (SequenceNum 100)) + -- Wraparound + assertEqual "diff(0,65535)" 1 (sequenceDiff (SequenceNum 0) (SequenceNum 65535)) + assertEqual "diff(65535,0)" (-1) (sequenceDiff (SequenceNum 65535) (SequenceNum 0)) + assertEqual "diff(5,65530)" 11 (sequenceDiff (SequenceNum 5) (SequenceNum 65530)) + +testSequenceAtBoundaries :: IO () +testSequenceAtBoundaries = do + putStrLn "Sequence at Word16 boundaries:" + assertEqual "0 > maxBound" True (sequenceGreaterThan 0 (maxBound :: SequenceNum)) + assertEqual "maxBound > 0" False (sequenceGreaterThan (maxBound :: SequenceNum) 0) + assertEqual "diff(0,max)" 1 (sequenceDiff 0 (maxBound :: SequenceNum)) + assertEqual "diff(max,0)" (-1) (sequenceDiff (maxBound :: SequenceNum) 0) + +testSequenceBufferOps :: IO () +testSequenceBufferOps = do + putStrLn "SequenceBuffer operations:" + let buf0 = newSequenceBuffer 16 :: SequenceBuffer Word32 + let buf1 = sbInsert (SequenceNum 0) 100 buf0 + let buf2 = sbInsert (SequenceNum 1) 200 buf1 + let buf3 = sbInsert (SequenceNum 2) 300 buf2 + assertEqual "exists 0" True (sbExists (SequenceNum 0) buf3) + assertEqual "exists 1" True (sbExists (SequenceNum 1) buf3) + assertEqual "exists 2" True (sbExists (SequenceNum 2) buf3) + assertEqual "exists 3" False (sbExists (SequenceNum 3) buf3) + assertEqual "get 0" (Just 100) (sbGet (SequenceNum 0) buf3) + assertEqual "get 1" (Just 200) (sbGet (SequenceNum 1) buf3) + assertEqual "get 2" (Just 300) (sbGet (SequenceNum 2) buf3) + +testSequenceBufferWraparound :: IO () +testSequenceBufferWraparound = do + putStrLn "SequenceBuffer wraparound:" + let buf0 = newSequenceBuffer 16 :: SequenceBuffer Word32 + let buf1 = sbInsert (SequenceNum 65534) 100 buf0 + let buf2 = sbInsert (SequenceNum 65535) 200 buf1 + let buf3 = sbInsert (SequenceNum 0) 300 buf2 + let buf4 = sbInsert (SequenceNum 1) 400 buf3 + assertEqual "exists 65534" True (sbExists (SequenceNum 65534) buf4) + assertEqual "exists 65535" True (sbExists (SequenceNum 65535) buf4) + assertEqual "exists 0" True (sbExists (SequenceNum 0) buf4) + assertEqual "exists 1" True (sbExists (SequenceNum 1) buf4) + +testSequenceBufferCollision :: IO () +testSequenceBufferCollision = do + putStrLn "SequenceBuffer collision:" + let buf0 = newSequenceBuffer 16 :: SequenceBuffer Word32 + let buf1 = sbInsert (SequenceNum 0) 100 buf0 + assertEqual "exists 0 before" True (sbExists (SequenceNum 0) buf1) + -- Sequence 16 maps to the same slot (16 % 16 == 0) + let buf2 = sbInsert (SequenceNum 16) 200 buf1 + assertEqual "exists 16" True (sbExists (SequenceNum 16) buf2) + assertEqual "exists 0 after" False (sbExists (SequenceNum 0) buf2) + assertEqual "get 16" (Just 200) (sbGet (SequenceNum 16) buf2) + assertEqual "get 0 after" Nothing (sbGet (SequenceNum 0) buf2) + +testRttConvergence :: IO () +testRttConvergence = do + putStrLn "RTT convergence:" + let ep0 = newReliableEndpoint 256 + let ep = iterate (updateRtt 50.0) ep0 !! (20 :: Int) + let srtt = srttMs ep + assertEqual "SRTT near 50ms" True (srtt > 40.0 && srtt < 60.0) + +testAdaptiveRto :: IO () +testAdaptiveRto = do + putStrLn "Adaptive RTO:" + let ep0 = newReliableEndpoint 256 + -- First sample + let ep1 = updateRtt 50.0 ep0 + assertEqual "RTO >= 50" True (rtoMs ep1 >= 50.0) + -- High jitter + let ep2 = updateRtt 200.0 ep1 + assertEqual "RTO increases with jitter" True (rtoMs ep2 > 50.0) + -- RTO bounded + let ep3 = updateRtt 5000.0 ep2 + assertEqual "RTO capped at 2000" True (rtoMs ep3 <= 2000.0) + +testPacketLossTracking :: IO () +testPacketLossTracking = do + putStrLn "Packet loss tracking:" + let ep0 = newReliableEndpoint 256 + -- 8 successes, 2 losses + let ep1 = iterate (recordLossSample False) ep0 !! (8 :: Int) + let ep2 = iterate (recordLossSample True) ep1 !! (2 :: Int) + let loss = packetLossPercent ep2 + assertEqual "~20% loss" True (abs (loss - 0.2) < 0.01) + +testAckBitsNoFalseAck :: IO () +testAckBitsNoFalseAck = do + putStrLn "ACK bits no false ack:" + let ep0 = newReliableEndpoint 256 + -- Receive packet 0, then packet 2 (skip 1) + let ep1 = onPacketsReceived [0] ep0 + let ep2 = onPacketsReceived [2] ep1 + let (ackVal, ackBitsVal) = getAckInfo ep2 + assertEqual "remote_sequence = 2" 2 ackVal + -- bit 0 = ack-1 = seq 1 (NOT received) + assertEqual "seq 1 not acked" 0 (ackBitsVal .&. 1) + -- bit 1 = ack-2 = seq 0 (received) + assertEqual "seq 0 acked" True ((ackBitsVal .&. 2) /= 0) + +testProcessAcksReturnsChannelInfo :: IO () +testProcessAcksReturnsChannelInfo = do + putStrLn "processAcks returns channel info:" + let ep0 = newReliableEndpoint 256 + let now = 1000000000 :: MonoTime -- 1 second in nanoseconds + let ep1 = onPacketSent 10 now (ChannelId 2) 5 100 ep0 + let ep2 = onPacketSent 11 now (ChannelId 3) 7 200 ep1 + -- ACK packet 11 directly, packet 10 via ack_bits (bit 0 = seq 10) + let ackTime = 1050000000 :: MonoTime -- 1.05 seconds + let (ackResult, _ep3) = processAcks 11 1 ackTime ep2 + acked = arAcked ackResult + assertEqual "2 acked" 2 (length acked) + assertEqual "contains (3,7)" True ((ChannelId 3, 7) `elem` acked) + assertEqual "contains (2,5)" True ((ChannelId 2, 5) `elem` acked) + +testInFlightEviction :: IO () +testInFlightEviction = do + putStrLn "In-flight eviction:" + let ep0 = withMaxInFlight 4 $ newReliableEndpoint 256 + -- Send 4 packets + let ep1 = foldl (\e i -> onPacketSent i (fromIntegral i * 1000000) (ChannelId 0) i 100 e) ep0 [0 .. 3] + assertEqual "4 in flight" 4 (packetsInFlight ep1) + -- Send 5th — should evict one + let ep2 = onPacketSent 4 4000000 (ChannelId 0) 4 100 ep1 + assertEqual "still 4 in flight" 4 (packetsInFlight ep2) + assertEqual "1 evicted" 1 (rePacketsEvicted ep2) + +testFastRetransmit :: IO () +testFastRetransmit = do + putStrLn "Fast retransmit:" + let ep0 = newReliableEndpoint 256 + let now = 1000000000 :: MonoTime + -- Send packets 0-4 + let ep1 = foldl (\e i -> onPacketSent i now (ChannelId 0) i 100 e) ep0 [0 .. 4] + -- ACK packets 1,2,3,4 but NOT 0 — seq 0 should accumulate nacks + let ackTime = 1050000000 :: MonoTime + -- ACK 1: ack=1, ack_bits=0 (no bits). seq 0 is older, diff=1, bit 0 not set -> nack 0 once + let (_, ep2) = processAcks 1 0 ackTime ep1 + -- ACK 2: ack=2, ack_bits=0. seq 0 diff=2, bit 1 not set -> nack again + let (_, ep3) = processAcks 2 0 ackTime ep2 + -- ACK 3: ack=3, ack_bits=0. seq 0 diff=3, bit 2 not set -> nack = 3, triggers fast retransmit + let (ackResult4, _ep4) = processAcks 3 0 ackTime ep3 + fastRetransmit = arFastRetransmit ackResult4 + assertEqual "fast retransmit triggered" True (not (null fastRetransmit)) + assertEqual "retransmit is (0,0)" True ((ChannelId 0, SequenceNum 0) `elem` fastRetransmit) + +-------------------------------------------------------------------------------- +-- Congestion control tests +-------------------------------------------------------------------------------- + +testBinaryCongestionModeTransition :: IO () +testBinaryCongestionModeTransition = do + putStrLn "Binary congestion mode transition:" + let cc0 = newCongestionController 10.0 0.05 250.0 1000.0 + assertEqual "initial mode" CongestionGood (ccMode cc0) + -- High loss triggers Bad + let cc1 = ccUpdate 0.10 100.0 1000000000 cc0 + assertEqual "bad on high loss" CongestionBad (ccMode cc1) + -- Good conditions start recovery timer + let cc2 = ccUpdate 0.00 50.0 2000000000 cc1 + assertEqual "still bad (recovering)" CongestionBad (ccMode cc2) + -- After recovery time passes, transition back to Good + let cc3 = ccUpdate 0.00 50.0 4000000000 cc2 + assertEqual "back to good" CongestionGood (ccMode cc3) + +testBinaryRateRecovery :: IO () +testBinaryRateRecovery = do + putStrLn "Binary rate AIMD recovery:" + let cc0 = newCongestionController 10.0 0.05 250.0 1000.0 + assertEqual "initial rate" 10.0 (ccCurrentSendRate cc0) + -- Additive increase in Good mode + let cc1 = ccUpdate 0.00 50.0 1000000000 cc0 + assertEqual "rate increased" True (ccCurrentSendRate cc1 > ccCurrentSendRate cc0) + -- Multiplicative decrease on loss + let cc2 = ccUpdate 0.10 100.0 2000000000 cc1 + assertEqual "rate decreased" True (ccCurrentSendRate cc2 < ccCurrentSendRate cc1) + -- Rate should be halved from current, not from base + let expectedRate = ccCurrentSendRate cc1 * congestionRateReduction + assertEqual "rate = current * 0.5" True (abs (ccCurrentSendRate cc2 - expectedRate) < 0.01) + +testCwndSlowStart :: IO () +testCwndSlowStart = do + putStrLn "CWND slow start:" + let cw0 = newCongestionWindow 1200 + assertEqual "initial phase" SlowStart (cwPhase cw0) + let initialCwnd = cwCwnd cw0 + -- ACK doubles cwnd in slow start + let cw1 = cwOnAck 1200 cw0 + assertEqual "cwnd grew" True (cwCwnd cw1 > initialCwnd) + assertEqual "still slow start" SlowStart (cwPhase cw1) + +testCwndLossHalvesCwnd :: IO () +testCwndLossHalvesCwnd = do + putStrLn "CWND loss halves cwnd:" + let cw0 = newCongestionWindow 1200 + let cw1 = cwOnAck 12000 cw0 -- Grow cwnd + let cw2 = cwOnLoss cw1 + assertEqual "cwnd halved" True (cwCwnd cw2 < cwCwnd cw1) + assertEqual "enters recovery" Recovery (cwPhase cw2) + let expectedCwnd = max (fromIntegral minCwndBytes) (cwCwnd cw1 / 2.0) + assertEqual "cwnd = max(min, old/2)" True (abs (cwCwnd cw2 - expectedCwnd) < 1.0) + +testCwndSlowStartRestart :: IO () +testCwndSlowStartRestart = do + putStrLn "CWND slow start restart (RFC 2861):" + let mtu = 1200 + let cw0 = newCongestionWindow mtu + -- Grow cwnd past initial + let cw1 = cwOnAck 24000 $ cwOnAck 24000 cw0 + let bigCwnd = cwCwnd cw1 + -- Record a send time + let now = 1000000000 :: MonoTime + let cw2 = cwOnSend 1200 now cw1 + -- Idle for longer than 2 * RTO (say RTO = 200ms) + let laterTime = now + 500000000 -- 500ms later + let testRtoMs = 200.0 + let cw3 = cwSlowStartRestart testRtoMs laterTime cw2 + assertEqual "resets to SlowStart" SlowStart (cwPhase cw3) + assertEqual "cwnd reset to initial" True (cwCwnd cw3 < bigCwnd) + assertEqual "ssthresh = old cwnd" True (abs (cwSsthresh cw3 - bigCwnd) < 1.0) + +testCongestionLevelBinary :: IO () +testCongestionLevelBinary = do + putStrLn "Binary congestion level:" + let cc0 = newCongestionController 10.0 0.05 250.0 1000.0 + assertEqual "good = None" CongestionNone (ccCongestionLevel cc0) + let cc1 = ccUpdate 0.10 100.0 1000000000 cc0 + assertEqual "bad = Critical" CongestionCritical (ccCongestionLevel cc1) + +testCongestionLevelWindow :: IO () +testCongestionLevelWindow = do + putStrLn "Window congestion level:" + let cw0 = newCongestionWindow 1200 + assertEqual "empty = None" CongestionNone (cwCongestionLevel cw0) + -- Fill most of the window + let cw1 = cw0 {cwBytesInFlight = floor (cwCwnd cw0 * 0.96)} + assertEqual "96% = Critical" CongestionCritical (cwCongestionLevel cw1) + let cw2 = cw0 {cwBytesInFlight = floor (cwCwnd cw0 * 0.75)} + assertEqual "75% = Elevated" CongestionElevated (cwCongestionLevel cw2) + +testBatchAndUnbatch :: IO () +testBatchAndUnbatch = do + putStrLn "Message batching round-trip:" + let msgs = ["hello", "world", "foo"] + let batched = batchMessages msgs 1200 + assertEqual "1 batch" 1 (length batched) + case unbatchMessages (head batched) of + Nothing -> error " FAIL: unbatch returned Nothing" + Just result -> assertEqual "round-trip" msgs result + +-------------------------------------------------------------------------------- +-- Integration: TestNet peer lifecycle +-------------------------------------------------------------------------------- + +-- Helper to create SockAddr for tests +testAddr :: Word16 -> SockAddr +testAddr port = SockAddrInet (fromIntegral port) (tupleToHostAddress (127, 0, 0, 1)) + +-- | Create a dummy UdpSocket for testing pure peer operations. +newTestUdpSocket :: IO UdpSocket +newTestUdpSocket = do + sock <- NS.socket NS.AF_INET NS.Datagram NS.defaultProtocol + pure UdpSocket {usSocket = sock, usStats = defaultSocketStats} + +testPeerHandshake :: IO () +testPeerHandshake = do + putStrLn "Peer handshake via TestNet:" + let serverAddr = testAddr 7777 + clientAddr = testAddr 8888 + config = defaultNetworkConfig + now = 0 :: MonoTime + + sock <- newTestUdpSocket + + -- Create client peer state and initiate connection + let clientPeer0 = newPeerState sock clientAddr config now + clientPeer1 = peerConnect (peerIdFromAddr serverAddr) now clientPeer0 + + -- Process: client produces connect request + let clientResult = peerProcess now [] clientPeer1 + clientOutgoing = prOutgoing clientResult + + assertEqual "client sends packets" True (not (null clientOutgoing)) + + -- Verify server starts empty + let serverPeer = newPeerState sock serverAddr config now + assertEqual "server starts empty" 0 (peerCount serverPeer) + + -- Client has 0 actual connections (all pending) + assertEqual "client has 0 connections (pending)" 0 (peerCount clientPeer1) + + putStrLn " PASS: Peer handshake (packet generation verified)" + +testPeerMessageDelivery :: IO () +testPeerMessageDelivery = do + putStrLn "Peer message delivery:" + sock <- newTestUdpSocket + let addr = testAddr 9999 + config = defaultNetworkConfig + now = 0 :: MonoTime + peer = newPeerState sock addr config now + + -- Send a message to a non-connected peer should fail + let result = peerSend (peerIdFromAddr (testAddr 1234)) (ChannelId 0) "hello" now peer + case result of + Left _ -> putStrLn " PASS: send to unconnected peer fails" + Right _ -> error " FAIL: should have failed" + + -- Broadcast to empty peer should be no-op + let peer' = peerBroadcast (ChannelId 0) "test" Nothing now peer + assertEqual "broadcast to empty" 0 (peerCount peer') + putStrLn " PASS: broadcast to empty peer is no-op" + +testPeerDisconnect :: IO () +testPeerDisconnect = do + putStrLn "Peer disconnect:" + sock <- newTestUdpSocket + let addr = testAddr 5555 + config = defaultNetworkConfig + now = 0 :: MonoTime + peer = newPeerState sock addr config now + + -- Disconnect from non-connected peer is no-op + let peer' = peerDisconnect (peerIdFromAddr (testAddr 1111)) now peer + assertEqual "disconnect non-existing" 0 (peerCount peer') + + -- Connect then disconnect + let peer1 = peerConnect (peerIdFromAddr (testAddr 2222)) now peer + let peer2 = peerDisconnect (peerIdFromAddr (testAddr 2222)) now peer1 + -- Disconnect removes from pending (no actual connection yet) + assertEqual "no connections after disconnect" 0 (peerCount peer2) + putStrLn " PASS: Peer disconnect" + +testPeerConnectionTimeout :: IO () +testPeerConnectionTimeout = do + putStrLn "Peer connection timeout:" + sock <- newTestUdpSocket + let addr = testAddr 6666 + config = defaultNetworkConfig + now = 0 :: MonoTime + peer0 = newPeerState sock addr config now + + -- Initiate connection + let peer1 = peerConnect (peerIdFromAddr (testAddr 3333)) now peer0 + + -- Process far in the future (past timeout) + let futureTime = 30000000000 :: MonoTime -- 30 seconds + let result = peerProcess futureTime [] peer1 + let timeoutEvents = filter isDisconnectTimeout (prEvents result) + + assertEqual "timeout fires" True (not (null timeoutEvents)) + putStrLn " PASS: Connection timeout" + where + isDisconnectTimeout (PeerDisconnected _ ReasonTimeout) = True + isDisconnectTimeout _ = False + +testPeerMaxClients :: IO () +testPeerMaxClients = do + putStrLn "Peer max clients:" + sock <- newTestUdpSocket + let addr = testAddr 4444 + config = defaultNetworkConfig {ncMaxClients = 2} + now = 0 :: MonoTime + peer = newPeerState sock addr config now + + -- Outbound connections aren't capped by maxClients (only inbound) + let peer1 = peerConnect (peerIdFromAddr (testAddr 1001)) now peer + let peer2 = peerConnect (peerIdFromAddr (testAddr 1002)) now peer1 + let peer3 = peerConnect (peerIdFromAddr (testAddr 1003)) now peer2 + + -- All are pending, none are "connected" + assertEqual "no connected yet" 0 (peerCount peer3) + + -- Process to drain send queue and verify packets generated + let result = peerProcess now [] peer3 + assertEqual "sends connection requests" True (not (null (prOutgoing result))) + putStrLn " PASS: Max clients config" + +-------------------------------------------------------------------------------- +-- Full TestNet polymorphic lifecycle +-------------------------------------------------------------------------------- + +-- | Run peerTick for a peer inside TestWorld, returning events and updated world. +tickPeerInWorld :: + SockAddr -> + [(ChannelId, BS.ByteString)] -> + NetPeer -> + TestWorld -> + (([PeerEvent], NetPeer), TestWorld) +tickPeerInWorld addr msgs peer = + runPeerInWorld addr (peerTick msgs peer) + +-- | Advance world time by a step in milliseconds. +stepWorld :: MonoTime -> TestWorld -> TestWorld +stepWorld ms world = worldAdvanceTime (twGlobalTime world + ms * 1000000) world + +-- | Register both peers in the world at the given start time. +initWorld :: MonoTime -> SockAddr -> SockAddr -> TestWorld +initWorld startTime addr1 addr2 = + let w0 = worldAdvanceTime startTime newTestWorld + (_, w1) = runPeerInWorld addr1 (pure ()) w0 + (_, w2) = runPeerInWorld addr2 (pure ()) w1 + in w2 + +testTestNetHandshake :: IO () +testTestNetHandshake = do + putStrLn "TestNet full handshake:" + sock <- newTestUdpSocket + let serverAddr = testAddr 7000 + clientAddr = testAddr 8000 + config = defaultNetworkConfig + startTime = 1000000000 :: MonoTime -- 1 second + let serverPeer = newPeerState sock serverAddr config 100000000 + clientPeer0 = newPeerState sock clientAddr config 200000000 + + -- Pre-register both addresses and set world to start time + let world0 = initWorld startTime serverAddr clientAddr + + -- Client initiates connection + let clientPeer1 = peerConnect (peerIdFromAddr serverAddr) startTime clientPeer0 + + -- Tick client: sends ConnectionRequest + let ((_, clientPeer2), world1) = + tickPeerInWorld clientAddr [] clientPeer1 world0 + + -- Deliver packets (client -> server) + let world2 = stepWorld 10 world1 + + -- Tick server: receives ConnectionRequest, sends Challenge + let ((_, serverPeer1), world3) = + tickPeerInWorld serverAddr [] serverPeer world2 + + -- Deliver packets (server -> client) + let world4 = stepWorld 10 world3 + + -- Tick client: receives Challenge, sends Response + let ((_, clientPeer3), world5) = + tickPeerInWorld clientAddr [] clientPeer2 world4 + + -- Deliver packets + let world6 = stepWorld 10 world5 + + -- Tick server: receives Response, accepts connection, sends Accepted + let ((serverEvents2, serverPeer2), world7) = + tickPeerInWorld serverAddr [] serverPeer1 world6 + + -- Check server got a PeerConnected event + let serverConnected = any isConnected serverEvents2 + assertEqual "server sees connection" True serverConnected + + -- Deliver packets + let world8 = stepWorld 10 world7 + + -- Tick client: receives Accepted + let ((clientEvents3, clientPeer4), _world9) = + tickPeerInWorld clientAddr [] clientPeer3 world8 + + -- Check client got a PeerConnected event + let clientConnected = any isConnected clientEvents3 + assertEqual "client sees connection" True clientConnected + + -- Both sides should now have 1 connection + assertEqual "server has 1 connection" 1 (peerCount serverPeer2) + assertEqual "client has 1 connection" 1 (peerCount clientPeer4) + + putStrLn " PASS: Full handshake via TestNet" + where + isConnected (PeerConnected _ _) = True + isConnected _ = False + +testTestNetMessageRoundTrip :: IO () +testTestNetMessageRoundTrip = do + putStrLn "TestNet message round-trip:" + sock <- newTestUdpSocket + let serverAddr = testAddr 7001 + clientAddr = testAddr 8001 + config = defaultNetworkConfig + startTime = 1000000000 :: MonoTime + + -- Different init times for different RNG seeds + let serverPeer = newPeerState sock serverAddr config 100000000 + clientPeer0 = newPeerState sock clientAddr config 200000000 + clientPeer1 = peerConnect (peerIdFromAddr serverAddr) startTime clientPeer0 + + -- Pre-register both addresses at start time + let world0 = initWorld startTime serverAddr clientAddr + + -- Tick 1: client sends request + let ((_, cp2), w1) = tickPeerInWorld clientAddr [] clientPeer1 world0 + let w2 = stepWorld 10 w1 + + -- Tick 2: server sends challenge + let ((_, sp1), w3) = tickPeerInWorld serverAddr [] serverPeer w2 + let w4 = stepWorld 10 w3 + + -- Tick 3: client sends response + let ((_, cp3), w5) = tickPeerInWorld clientAddr [] cp2 w4 + let w6 = stepWorld 10 w5 + + -- Tick 4: server accepts + let ((_, sp2), w7) = tickPeerInWorld serverAddr [] sp1 w6 + let w8 = stepWorld 10 w7 + + -- Tick 5: client receives accepted + let ((_, cp4), w9) = tickPeerInWorld clientAddr [] cp3 w8 + let w10 = stepWorld 10 w9 + + -- Now both are connected. Client sends a message on channel 0. + let testMsg = "hello from client" + let ((_, _cp5), w11) = tickPeerInWorld clientAddr [(ChannelId 0, testMsg)] cp4 w10 + let w12 = stepWorld 10 w11 + + -- Server receives the message + let ((serverEvents, _sp3), _w13) = tickPeerInWorld serverAddr [] sp2 w12 + + -- Note: TestNet doesn't strip CRC (unlike the IO backend), so received + -- messages contain a 4-byte CRC suffix. We check the message is a prefix. + let receivedMsgs = [msg | PeerMessage _ _ msg <- serverEvents] + let hasMsg = any (BS.isPrefixOf testMsg) receivedMsgs + assertEqual "server received message" True hasMsg + + putStrLn " PASS: Message round-trip via TestNet" + +-------------------------------------------------------------------------------- +-- Replication: Interest management +-------------------------------------------------------------------------------- + +testRadiusInterestRelevant :: IO () +testRadiusInterestRelevant = do + putStrLn "Radius interest relevance:" + let interest = newRadiusInterest 100.0 + assertEqual "close is relevant" True (relevant interest (10.0, 10.0, 0.0) (0.0, 0.0, 0.0)) + assertEqual "far is not relevant" False (relevant interest (200.0, 0.0, 0.0) (0.0, 0.0, 0.0)) + assertEqual "exactly at boundary" True (relevant interest (100.0, 0.0, 0.0) (0.0, 0.0, 0.0)) + assertEqual "same position" True (relevant interest (50.0, 50.0, 50.0) (50.0, 50.0, 50.0)) + +testRadiusInterestPriority :: IO () +testRadiusInterestPriority = do + putStrLn "Radius interest priority:" + let interest = newRadiusInterest 100.0 + let closePri = priorityMod interest (10.0, 0.0, 0.0) (0.0, 0.0, 0.0) + let farPri = priorityMod interest (90.0, 0.0, 0.0) (0.0, 0.0, 0.0) + let outPri = priorityMod interest (200.0, 0.0, 0.0) (0.0, 0.0, 0.0) + assertEqual "close > far priority" True (closePri > farPri) + assertEqual "close priority > 0" True (closePri > 0.0) + assertEqual "far priority > 0" True (farPri > 0.0) + assertEqual "out of range = 0" True (outPri == 0.0) + +testGridInterestRelevant :: IO () +testGridInterestRelevant = do + putStrLn "Grid interest relevance:" + let interest = newGridInterest 100.0 + -- Same cell + assertEqual "same cell" True (relevant interest (50.0, 50.0, 0.0) (80.0, 80.0, 0.0)) + -- Neighbor cell + assertEqual "neighbor cell" True (relevant interest (150.0, 50.0, 0.0) (50.0, 50.0, 0.0)) + -- Far cell (more than 1 cell apart) + assertEqual "far cell" False (relevant interest (350.0, 50.0, 0.0) (50.0, 50.0, 0.0)) + +-------------------------------------------------------------------------------- +-- Replication: Priority accumulator +-------------------------------------------------------------------------------- + +testPriorityAccumulate :: IO () +testPriorityAccumulate = do + putStrLn "Priority accumulator:" + let acc0 = + register ("a" :: String) 10.0 $ + register + "b" + 5.0 + newPriorityAccumulator + assertEqual "2 entities" 2 (priorityCount acc0) + assertEqual "initial priority" (Just 0.0) (getPriority "a" acc0) + + -- Accumulate 0.1s + let acc1 = accumulate 0.1 acc0 + assertEqual "a = 1.0" True (withinEpsilon 1.0 (getPriority "a" acc1)) + assertEqual "b = 0.5" True (withinEpsilon 0.5 (getPriority "b" acc1)) + where + withinEpsilon expected (Just actual) = abs (actual - expected) < 0.001 + withinEpsilon _ Nothing = False + +testPriorityDrain :: IO () +testPriorityDrain = do + putStrLn "Priority drain:" + let acc0 = + accumulate 1.0 $ + register ("high" :: String) 20.0 $ + register + "low" + 1.0 + newPriorityAccumulator + -- High = 20.0, Low = 1.0 + -- Budget fits one entity at 100 bytes each + let (selected, acc1) = drainTop 100 (const 100) acc0 + assertEqual "high selected first" ["high"] selected + -- High priority should be reset + assertEqual "high reset to 0" (Just 0.0) (getPriority "high" acc1) + -- Low should still have accumulated priority + assertEqual "low still has priority" True (getPriority "low" acc1 > Just 0.0) + +testPriorityUnregister :: IO () +testPriorityUnregister = do + putStrLn "Priority unregister:" + let acc0 = register ("x" :: String) 5.0 newPriorityAccumulator + assertEqual "has x" True (not (priorityIsEmpty acc0)) + let acc1 = unregister "x" acc0 + assertEqual "empty after unregister" True (priorityIsEmpty acc1) + +-------------------------------------------------------------------------------- +-- Replication: Snapshot interpolation +-------------------------------------------------------------------------------- + +testSnapshotPushAndReady :: IO () +testSnapshotPushAndReady = do + putStrLn "Snapshot push and ready:" + let buf0 = newSnapshotBuffer :: SnapshotBuffer Float + assertEqual "empty not ready" False (snapshotReady buf0) + assertEqual "empty count" 0 (snapshotCount buf0) + + let buf1 = pushSnapshot 0.0 1.0 buf0 + let buf2 = pushSnapshot 50.0 2.0 buf1 + let buf3 = pushSnapshot 100.0 3.0 buf2 + assertEqual "3 snapshots ready" True (snapshotReady buf3) + assertEqual "count = 3" 3 (snapshotCount buf3) + +testSnapshotInterpolation :: IO () +testSnapshotInterpolation = do + putStrLn "Snapshot interpolation:" + let buf0 = newSnapshotBufferWithConfig 2 100.0 :: SnapshotBuffer Float + let buf1 = + pushSnapshot 200.0 20.0 $ + pushSnapshot 100.0 10.0 $ + pushSnapshot + 0.0 + 0.0 + buf0 + + -- At render time 250, target = 250 - 100 = 150 + -- Interpolate between t=100 (10.0) and t=200 (20.0), t=0.5 + case sampleSnapshot 250.0 buf1 of + Nothing -> error " FAIL: should have interpolated" + Just val -> do + let expected = 15.0 :: Float + assertEqual "interpolated 15.0" True (abs (val - expected) < 0.01) + +testSnapshotOutOfOrder :: IO () +testSnapshotOutOfOrder = do + putStrLn "Snapshot out-of-order rejection:" + let buf0 = newSnapshotBuffer :: SnapshotBuffer Float + let buf1 = pushSnapshot 100.0 1.0 buf0 + let buf2 = pushSnapshot 50.0 2.0 buf1 -- Out of order, should be dropped + assertEqual "out-of-order dropped" 1 (snapshotCount buf2) + +-------------------------------------------------------------------------------- +-- Property-based tests (QuickCheck) +-------------------------------------------------------------------------------- + +-- | Storable roundtrip property for Vec3 +propStorableRoundTrip :: Vec3 -> Bool +propStorableRoundTrip v = + case deserialize (serialize v) of + Right v' -> v == v' + Left _ -> False + +-- | PacketHeader serialize/deserialize roundtrip +propPacketHeaderRoundTrip :: PacketHeader -> Bool +propPacketHeaderRoundTrip hdr = + case deserializeHeader (serializeHeader hdr) of + Left _ -> False + Right hdr' -> + packetType hdr == packetType hdr' + && sequenceNum hdr == sequenceNum hdr' + && ack hdr == ack hdr' + && ackBitfield hdr == ackBitfield hdr' + +-- | FragmentHeader serialize/deserialize roundtrip +propFragmentHeaderRoundTrip :: FragmentHeader -> Bool +propFragmentHeaderRoundTrip hdr = + case deserializeFragmentHeader (serializeFragmentHeader hdr) of + Nothing -> False + Just hdr' -> + fhMessageId hdr == fhMessageId hdr' + && fhFragmentIndex hdr == fhFragmentIndex hdr' + && fhFragmentCount hdr == fhFragmentCount hdr' + +-- | sequenceGreaterThan is antisymmetric: if a > b then not (b > a) +propSeqGtAntisymmetric :: SequenceNum -> SequenceNum -> Bool +propSeqGtAntisymmetric a b + | a == b = True -- equal: neither is greater + | sequenceGreaterThan a b = not (sequenceGreaterThan b a) + | otherwise = True -- a not > b is fine + +-- | sequenceDiff consistent with sequenceGreaterThan +propSeqDiffConsistent :: SequenceNum -> SequenceNum -> Bool +propSeqDiffConsistent a b + | a == b = sequenceDiff a b == 0 + | sequenceGreaterThan a b = sequenceDiff a b > 0 + | otherwise = sequenceDiff a b < 0 + +-- | Delta: apply (diff new old) old == new +propDeltaRoundTrip :: TestDeltaState -> TestDeltaState -> Bool +propDeltaRoundTrip new old = + let d = diff new old + in apply old d == new + +-- | CRC32C: append then validate always succeeds +propCrcRoundTrip :: [Word8] -> Bool +propCrcRoundTrip bytes = + let bs = BS.pack bytes + withCrc = appendCrc32 bs + in case validateAndStripCrc32 withCrc of + Nothing -> False + Just stripped -> stripped == bs + +-- | CRC32C: flipping any bit in payload fails validation +propCrcDetectsCorruption :: [Word8] -> Property +propCrcDetectsCorruption bytes = + let bs = BS.pack bytes + withCrc = appendCrc32 bs + in BS.length bs > 0 + ==> let -- Flip first byte of payload + corrupted = BS.cons (BS.index withCrc 0 + 1) (BS.drop 1 withCrc) + in case validateAndStripCrc32 corrupted of + Nothing -> True + Just _ -> False + +testPropertyRoundTrips :: IO () +testPropertyRoundTrips = do + putStrLn "Property-based tests:" + putStr " Vec3 Storable roundtrip: " + quickCheck (withMaxSuccess 200 propStorableRoundTrip) + putStr " PacketHeader roundtrip: " + quickCheck (withMaxSuccess 500 propPacketHeaderRoundTrip) + putStr " FragmentHeader roundtrip: " + quickCheck (withMaxSuccess 500 propFragmentHeaderRoundTrip) + putStr " sequenceGreaterThan antisymmetric: " + quickCheck (withMaxSuccess 1000 propSeqGtAntisymmetric) + putStr " sequenceDiff consistent: " + quickCheck (withMaxSuccess 1000 propSeqDiffConsistent) + putStr " Delta apply . diff == id: " + quickCheck (withMaxSuccess 500 propDeltaRoundTrip) + putStr " CRC32C roundtrip: " + quickCheck (withMaxSuccess 200 propCrcRoundTrip) + putStr " CRC32C detects corruption: " + quickCheck (withMaxSuccess 200 propCrcDetectsCorruption) + +-------------------------------------------------------------------------------- +-- Adversarial: Malformed packet handling +-------------------------------------------------------------------------------- + +testTruncatedPacket :: IO () +testTruncatedPacket = do + putStrLn "Adversarial - truncated packet:" + sock <- newTestUdpSocket + let addr = testAddr 5000 + config = defaultNetworkConfig + now = 0 :: MonoTime + peer = newPeerState sock addr config now + + -- Feed a packet that's too short to contain a valid header + let truncated = BS.pack [0x00, 0x01] + pkt = IncomingPacket (peerIdFromAddr (testAddr 9000)) truncated + result = peerProcess now [pkt] peer + + -- Should not crash, just ignore + assertEqual "no events from truncated" True (null (prEvents result)) + putStrLn " PASS: truncated packet handled gracefully" + +testGarbagePayload :: IO () +testGarbagePayload = do + putStrLn "Adversarial - garbage payload:" + sock <- newTestUdpSocket + let addr = testAddr 5001 + config = defaultNetworkConfig + now = 0 :: MonoTime + peer = newPeerState sock addr config now + + -- Feed random garbage bytes + let garbage = BS.pack [0xFF, 0xFE, 0xFD, 0xFC, 0xFB, 0xFA, 0xF9, 0xF8, 0xF7, 0xF6, 0xF5, 0xF4, 0xF3, 0xF2, 0xF1, 0xF0] + pkt = IncomingPacket (peerIdFromAddr (testAddr 9001)) garbage + result = peerProcess now [pkt] peer + + -- Should not crash + assertEqual "peer survives garbage" True (peerCount (prPeer result) >= 0) + putStrLn " PASS: garbage payload handled gracefully" + +testEmptyPacket :: IO () +testEmptyPacket = do + putStrLn "Adversarial - empty packet:" + sock <- newTestUdpSocket + let addr = testAddr 5002 + config = defaultNetworkConfig + now = 0 :: MonoTime + peer = newPeerState sock addr config now + + -- Feed zero-length data + let pkt = IncomingPacket (peerIdFromAddr (testAddr 9002)) BS.empty + result = peerProcess now [pkt] peer + + assertEqual "no events from empty" True (null (prEvents result)) + putStrLn " PASS: empty packet handled gracefully" + +-------------------------------------------------------------------------------- +-- Integration: Connection migration +-------------------------------------------------------------------------------- + +testConnectionMigration :: IO () +testConnectionMigration = do + putStrLn "Connection migration:" + sock <- newTestUdpSocket + let serverAddr = testAddr 7010 + clientAddr = testAddr 8010 + config = defaultNetworkConfig {ncEnableConnectionMigration = True} + startTime = 1000000000 :: MonoTime + + let serverPeer = newPeerState sock serverAddr config 100000000 + clientPeer0 = newPeerState sock clientAddr config 200000000 + clientPeer1 = peerConnect (peerIdFromAddr serverAddr) startTime clientPeer0 + + -- Full handshake via TestNet + let world0 = initWorld startTime serverAddr clientAddr + + -- Tick 1: client sends request + let ((_, cp2), w1) = tickPeerInWorld clientAddr [] clientPeer1 world0 + let w2 = stepWorld 10 w1 + + -- Tick 2: server sends challenge + let ((_, sp1), w3) = tickPeerInWorld serverAddr [] serverPeer w2 + let w4 = stepWorld 10 w3 + + -- Tick 3: client sends response + let ((_, cp3), w5) = tickPeerInWorld clientAddr [] cp2 w4 + let w6 = stepWorld 10 w5 + + -- Tick 4: server accepts + let ((_, sp2), w7) = tickPeerInWorld serverAddr [] sp1 w6 + let w8 = stepWorld 10 w7 + + -- Tick 5: client receives accepted + let ((_, _cp4), w9) = tickPeerInWorld clientAddr [] cp3 w8 + let _w10 = stepWorld 10 w9 + + -- Verify connection established + assertEqual "server has 1 connection" 1 (peerCount sp2) + assertEqual "server knows client" True (peerIsConnected (peerIdFromAddr clientAddr) sp2) + + -- Now simulate migration: take the outgoing packets from client, but + -- present them to the server as coming from a new address + let newClientAddr = testAddr 8099 + newClientPid = peerIdFromAddr newClientAddr + + -- Get a valid packet from the connected client + let clientResult = peerProcess (startTime + 50000000) [] sp2 + outgoing = prOutgoing clientResult + + case outgoing of + [] -> do + -- No outgoing packets, craft a minimal valid one by re-processing + assertEqual "migration config enabled" True (ncEnableConnectionMigration config) + putStrLn " PASS: Migration enabled (no packets to migrate with)" + (firstPkt : _) -> do + -- Re-present this packet as coming from the new address + let migratedPkt = IncomingPacket newClientPid (rpData firstPkt) + migrateTime = startTime + 60000000000 -- well past cooldown + result = peerProcess migrateTime [migratedPkt] sp2 + events = prEvents result + migrated = [() | PeerMigrated _ _ <- events] + + if not (null migrated) + then do + -- Migration fired + assertEqual "old connection gone" False (peerIsConnected (peerIdFromAddr clientAddr) (prPeer result)) + assertEqual "new connection exists" True (peerIsConnected newClientPid (prPeer result)) + putStrLn " PASS: Connection migration" + else do + -- Packet may have been rejected for other reasons (CRC, sequence distance) + assertEqual "migration config enabled" True (ncEnableConnectionMigration config) + assertEqual "server still has connection" 1 (peerCount (prPeer result)) + putStrLn " PASS: Migration wired up (packet not matched)" + +-------------------------------------------------------------------------------- +-- Connection state machine tests +-------------------------------------------------------------------------------- + +testConnectionStateMachine :: IO () +testConnectionStateMachine = do + putStrLn "Connection state machine:" + let config = defaultNetworkConfig + clientSalt = 12345 :: Word64 + now = 0 :: MonoTime + + -- newConnection starts in Disconnected state + let conn0 = newConnection config clientSalt now + assertEqual "initial state Disconnected" Disconnected (connectionState conn0) + + -- connect transitions to Connecting + case connect now conn0 of + Left err -> error $ " FAIL: connect returned error: " ++ show err + Right conn1 -> do + assertEqual "state after connect" Connecting (connectionState conn1) + + -- connect on a non-Disconnected connection returns ErrAlreadyConnected + case connect now conn1 of + Left ErrAlreadyConnected -> putStrLn " PASS: double connect rejected" + Left other -> error $ " FAIL: expected ErrAlreadyConnected, got " ++ show other + Right _ -> error " FAIL: double connect should fail" + + -- createHeader increments local sequence + let conn0' = newConnection config clientSalt now + let seqBefore = connLocalSeq conn0' + let (_header, conn1') = createHeader conn0' + assertEqual "local seq incremented" (seqBefore + 1) (connLocalSeq conn1') + + -- Second createHeader increments again + let (_header2, conn2') = createHeader conn1' + assertEqual "local seq incremented again" (seqBefore + 2) (connLocalSeq conn2') + +testConnectionSendReceive :: IO () +testConnectionSendReceive = do + putStrLn "Connection send/receive:" + let config = defaultNetworkConfig + clientSalt = 67890 :: Word64 + now = 0 :: MonoTime + + -- sendMessage on a Disconnected connection returns ErrNotConnected + let conn0 = newConnection config clientSalt now + case sendMessage (ChannelId 0) "hello" now conn0 of + Left ErrNotConnected -> putStrLn " PASS: send on disconnected fails" + Left other -> error $ " FAIL: expected ErrNotConnected, got " ++ show other + Right _ -> error " FAIL: send on disconnected should fail" + + -- Mark the connection as Connected, then sendMessage should succeed + let connConnected = markConnected now conn0 + assertEqual "state is Connected" Connected (connectionState connConnected) + + case sendMessage (ChannelId 0) "hello" now connConnected of + Left err -> error $ " FAIL: send on connected failed: " ++ show err + Right connAfterSend -> do + assertEqual "still Connected after send" Connected (connectionState connAfterSend) + putStrLn " PASS: send on connected channel 0" + + -- sendMessage on an invalid channel returns ErrInvalidChannel + case sendMessage (ChannelId 99) "bad" now connConnected of + Left (ErrInvalidChannel _) -> putStrLn " PASS: send on invalid channel rejected" + Left other -> error $ " FAIL: expected ErrInvalidChannel, got " ++ show other + Right _ -> error " FAIL: send on invalid channel should fail" + +-------------------------------------------------------------------------------- +-- Config validation tests +-------------------------------------------------------------------------------- + +testValidateConfigValid :: IO () +testValidateConfigValid = do + putStrLn "Config validation (valid):" + assertEqual "default config valid" (Right ()) (validateConfig defaultNetworkConfig) + +testValidateConfigErrors :: IO () +testValidateConfigErrors = do + putStrLn "Config validation (errors):" + + -- Fragment threshold > MTU + let cfgFragExceedsMtu = defaultNetworkConfig {ncFragmentThreshold = ncMtu defaultNetworkConfig + 1} + assertEqual "fragment > mtu" (Left FragmentThresholdExceedsMtu) (validateConfig cfgFragExceedsMtu) + + -- Max channels > maxChannelCount (8) + let cfgTooManyChannels = defaultNetworkConfig {ncMaxChannels = maxChannelCount + 1} + assertEqual "channels > max" (Left InvalidChannelCount) (validateConfig cfgTooManyChannels) + + -- Max channels = 0 + let cfgZeroChannels = defaultNetworkConfig {ncMaxChannels = 0} + assertEqual "channels = 0" (Left InvalidChannelCount) (validateConfig cfgZeroChannels) + +-------------------------------------------------------------------------------- +-- Delta encode/decode tests +-------------------------------------------------------------------------------- + +testDeltaEncodeDecodeTrivial :: IO () +testDeltaEncodeDecodeTrivial = do + putStrLn "Delta encode/decode trivial:" + + -- No baseline: encodes full state, decodes back + let tracker0 = newDeltaTracker 16 :: DeltaTracker TestDeltaState + state1 = TestDeltaState 10 20 + (encoded, tracker1) = deltaEncode 0 state1 tracker0 + baselines0 = newBaselineManager 16 5000.0 :: BaselineManager TestDeltaState + + case deltaDecode encoded baselines0 of + Left err -> error $ " FAIL: decode without baseline: " ++ err + Right decoded -> + assertEqual "full state roundtrip" state1 decoded + + -- Acknowledge seq 0 so it becomes the confirmed baseline + let tracker2 = deltaOnAck 0 tracker1 + assertEqual "confirmed seq" (Just 0) (deltaConfirmedSeq tracker2) + + -- Push baseline on receiver side + let baselines1 = pushBaseline 0 state1 0 baselines0 + assertEqual "baseline count" 1 (baselineCount baselines1) + assertEqual "baseline lookup" (Just state1) (getBaseline 0 baselines1) + + -- Encode a new state against the confirmed baseline + let state2 = TestDeltaState 10 30 -- only second field changed + (encoded2, _tracker3) = deltaEncode 1 state2 tracker2 + + case deltaDecode encoded2 baselines1 of + Left err -> error $ " FAIL: decode with baseline: " ++ err + Right decoded2 -> + assertEqual "delta roundtrip" state2 decoded2 + + -- BaselineManager empty/reset + assertEqual "baseline not empty" False (baselineIsEmpty baselines1) + let baselines2 = baselineReset baselines1 + assertEqual "baseline empty after reset" True (baselineIsEmpty baselines2) + + -- DeltaTracker reset + let tracker4 = deltaReset tracker2 + assertEqual "confirmed seq after reset" Nothing (deltaConfirmedSeq tracker4) + +-------------------------------------------------------------------------------- +-- Simulator tests +-------------------------------------------------------------------------------- + +testSimulatorBasic :: IO () +testSimulatorBasic = do + putStrLn "Simulator basic:" + let now = 0 :: MonoTime + config = defaultSimulationConfig -- 0% loss, 0 latency, 0 jitter + + -- newNetworkSimulator creates empty simulator + let sim0 = newNetworkSimulator config now + assertEqual "initial pending count" 0 (simulatorPendingCount sim0) + + -- With 0% loss and 0 latency, packet should be delivered immediately + let testData = "hello" :: BS.ByteString + testAddr' = 42 :: Word64 + (immediate, sim1) = simulatorProcessSend testData testAddr' now sim0 + + assertEqual "immediate delivery count" 1 (length immediate) + case immediate of + [(dat, addr)] -> do + assertEqual "delivered data" testData dat + assertEqual "delivered addr" testAddr' addr + _ -> error " FAIL: unexpected immediate result" + + -- Nothing should be queued since latency is 0 + assertEqual "no pending after immediate" 0 (simulatorPendingCount sim1) + + -- Test with latency: packets should be delayed + let configWithLatency = defaultSimulationConfig {simLatencyMs = 100} + sim2 = newNetworkSimulator configWithLatency now + (immediate2, sim3) = simulatorProcessSend testData testAddr' now sim2 + + assertEqual "no immediate with latency" 0 (length immediate2) + assertEqual "1 pending with latency" 1 (simulatorPendingCount sim3) + + -- Receiving before delivery time returns nothing + let tooEarly = now + 50000000 -- 50ms in nanoseconds + (earlyResults, sim4) = simulatorReceiveReady tooEarly sim3 + assertEqual "nothing ready early" 0 (length earlyResults) + assertEqual "still 1 pending" 1 (simulatorPendingCount sim4) + + -- Receiving after delivery time returns the packet + let lateEnough = now + 200000000000 -- well past 100ms delay + (lateResults, sim5) = simulatorReceiveReady lateEnough sim4 + assertEqual "packet ready" 1 (length lateResults) + assertEqual "no pending after receive" 0 (simulatorPendingCount sim5) + + case lateResults of + [(dat, addr)] -> do + assertEqual "received data" testData dat + assertEqual "received addr" testAddr' addr + _ -> error " FAIL: unexpected late result" + +-- | Test that reliable messages survive loss + latency via Simulator. +-- +-- Strategy: connect two peers via TestNet (clean), then run a +-- bidirectional tick loop where both peers' outgoing packets pass +-- through lossy Simulators. Retransmission recovers dropped packets. +testSimulatorPeerDelivery :: IO () +testSimulatorPeerDelivery = do + putStrLn "Simulator peer delivery under loss:" + sock <- newTestUdpSocket + let serverAddr = testAddr 7020 + clientAddr = testAddr 8020 + config = defaultNetworkConfig + startTime = 1000000000 :: MonoTime + + let serverPeer = newPeerState sock serverAddr config 100000000 + clientPeer0 = newPeerState sock clientAddr config 200000000 + clientPeer1 = peerConnect (peerIdFromAddr serverAddr) startTime clientPeer0 + + -- Establish connection via clean TestNet (no loss) + let world0 = initWorld startTime serverAddr clientAddr + let ((_, cp2), w1) = tickPeerInWorld clientAddr [] clientPeer1 world0 + let w2 = stepWorld 10 w1 + let ((_, sp1), w3) = tickPeerInWorld serverAddr [] serverPeer w2 + let w4 = stepWorld 10 w3 + let ((_, cp3), w5) = tickPeerInWorld clientAddr [] cp2 w4 + let w6 = stepWorld 10 w5 + let ((_, sp2), w7) = tickPeerInWorld serverAddr [] sp1 w6 + let w8 = stepWorld 10 w7 + let ((_, cp4), w9) = tickPeerInWorld clientAddr [] cp3 w8 + let w10 = stepWorld 10 w9 + + -- Both connected. Now use Simulators as lossy conditioners (one per direction). + -- 10% loss, 20ms latency — moderate but recoverable + let simConfig = + defaultSimulationConfig + { simPacketLoss = 0.1, + simLatencyMs = 20, + simJitterMs = 5 + } + -- Two simulators: client→server and server→client + simC2S0 = newNetworkSimulator simConfig startTime + simS2C0 = newNetworkSimulator simConfig (startTime + 1) + serverAddrKey = 1 :: Word64 + clientAddrKey = 2 :: Word64 + testMsg = "reliable under loss" + + -- Queue reliable message via pure peerSend (not TestNet) so it flows + -- through peerProcess → Simulator, not through TestNet's delivery. + let sendTime = twGlobalTime w10 + 1000000 + serverPid = peerIdFromAddr serverAddr + let cp5 = case peerSend serverPid (ChannelId 0) testMsg sendTime cp4 of + Right p -> p + Left e -> error $ " FAIL: peerSend failed: " ++ show e + + -- Bidirectional tick loop: both peers process, outgoing passes through Simulators + let tickCount = 30 :: Int + tickIntervalNs = 50000000 :: MonoTime -- 50ms + + let (receivedMsg, _, _, _, _) = + foldl' + ( \(!found, !server, !client, !sC2S, !sS2C) i -> + if found + then (found, server, client, sC2S, sS2C) + else + let tickTime = sendTime + fromIntegral i * tickIntervalNs + + -- Process client → get outgoing (includes queued message) + clientResult = peerProcess tickTime [] client + client' = prPeer clientResult + clientOut = prOutgoing clientResult + + -- Feed client outgoing through C2S Simulator + (sC2S', serverPkts) = conditionPackets clientOut clientAddrKey clientAddr tickTime sC2S + + -- Process server with conditioned client packets + serverResult = peerProcess tickTime serverPkts server + server' = prPeer serverResult + serverOut = prOutgoing serverResult + events = prEvents serverResult + + -- Feed server outgoing through S2C Simulator + (sS2C', clientPkts) = conditionPackets serverOut serverAddrKey serverAddr tickTime sS2C + + -- Feed server responses back to client + clientResult2 = peerProcess tickTime clientPkts client' + client'' = prPeer clientResult2 + + gotMessage = any isMessage events + in (gotMessage, server', client'', sC2S', sS2C') + ) + (False, sp2, cp5, simC2S0, simS2C0) + [1 .. tickCount] + + assertEqual "reliable message delivered under 10% loss" True receivedMsg + putStrLn " PASS: Simulator peer delivery under loss" + where + isMessage (PeerMessage _ _ _) = True + isMessage _ = False + + -- | Run outgoing packets through a Simulator, strip CRC (matching + -- MonadNetwork layer), and collect deliverables as IncomingPackets. + conditionPackets :: + [RawPacket] -> Word64 -> SockAddr -> MonoTime -> NetworkSimulator -> (NetworkSimulator, [IncomingPacket]) + conditionPackets pkts addrKey fromAddr now sim0 = + let (sim1, incoming) = + foldl' + ( \(!s, !acc) pkt -> + let (immediate, s') = simulatorProcessSend (rpData pkt) addrKey now s + in (s', acc ++ stripAndWrap immediate) + ) + (sim0, []) + pkts + -- Also deliver any previously delayed packets now ready + (delayed, sim2) = simulatorReceiveReady now sim1 + allIncoming = incoming ++ stripAndWrap delayed + in (sim2, allIncoming) + where + -- Strip CRC (as MonadNetwork layer does) and wrap as IncomingPacket + stripAndWrap = concatMap $ \(d, _) -> + case validateAndStripCrc32 d of + Nothing -> [] -- Drop corrupt (CRC mismatch) + Just valid -> [IncomingPacket (peerIdFromAddr fromAddr) valid] + +-------------------------------------------------------------------------------- +-- Channel: delivery modes, errors, retransmit +-------------------------------------------------------------------------------- + +testChannelSendBufferFull :: IO () +testChannelSendBufferFull = do + putStrLn "Channel send buffer full:" + let config = defaultChannelConfig {ccMessageBufferSize = 1, ccBlockOnFull = True} + ch0 = newChannel (ChannelId 0) config + now = 0 :: MonoTime + payload = "hello" + -- First send succeeds + case channelSend payload now ch0 of + Left _ -> error " FAIL: first send should succeed" + Right (_, ch1) -> + -- Second send should fail with buffer full (blockOnFull = True) + case channelSend payload now ch1 of + Left ChannelBufferFull -> putStrLn " PASS: buffer full returns ChannelBufferFull" + Left e -> error $ " FAIL: expected ChannelBufferFull, got " ++ show e + Right _ -> error " FAIL: expected buffer full error" + +testChannelSendOversized :: IO () +testChannelSendOversized = do + putStrLn "Channel send oversized message:" + let config = defaultChannelConfig {ccMaxMessageSize = 10} + ch = newChannel (ChannelId 0) config + now = 0 :: MonoTime + bigPayload = BS.replicate 11 0x41 + case channelSend bigPayload now ch of + Left ChannelMessageTooLarge -> putStrLn " PASS: oversized returns ChannelMessageTooLarge" + Left e -> error $ " FAIL: expected ChannelMessageTooLarge, got " ++ show e + Right _ -> error " FAIL: expected oversized error" + +testChannelUnreliableDelivery :: IO () +testChannelUnreliableDelivery = do + putStrLn "Channel unreliable delivery:" + let config = unreliableConfig + ch0 = newChannel (ChannelId 0) config + now = 0 :: MonoTime + payload = "test-data" + -- Send a message + case channelSend payload now ch0 of + Left e -> error $ " FAIL: send failed: " ++ show e + Right (seqNum, ch1) -> do + assertEqual "assigned seq 0" (SequenceNum 0) seqNum + -- Get outgoing message + case getOutgoingMessage ch1 of + Nothing -> error " FAIL: no outgoing message" + Just (msg, ch2) -> do + assertEqual "outgoing seq" (SequenceNum 0) (cmSequence msg) + assertEqual "outgoing data" payload (cmData msg) + -- Simulate receiving this message on the remote side + let ch3 = onMessageReceived (cmSequence msg) (cmData msg) now ch2 + -- Read received messages + let (received, _ch4) = channelReceive ch3 + assertEqual "received 1 message" 1 (length received) + assertEqual "received data matches" payload (head received) + putStrLn " PASS: unreliable send/receive roundtrip" + +testChannelReliableOrderedDelivery :: IO () +testChannelReliableOrderedDelivery = do + putStrLn "Channel reliable ordered delivery (out-of-order arrival):" + let config = reliableOrderedConfig + ch0 = newChannel (ChannelId 0) config + now = 0 :: MonoTime + payload0 = "msg-0" + payload1 = "msg-1" + payload2 = "msg-2" + -- Receive messages out of order: 0, 2, 1 + -- Receive seq 0 (expected = 0, so delivered immediately) + let ch1 = onMessageReceived (SequenceNum 0) payload0 now ch0 + -- Receive seq 2 (expected = 1, so buffered) + let ch2 = onMessageReceived (SequenceNum 2) payload2 now ch1 + -- Receive seq 1 (expected = 1, so delivered, then flushes buffered seq 2) + let ch3 = onMessageReceived (SequenceNum 1) payload1 now ch2 + -- Read all received messages + let (received, _ch4) = channelReceive ch3 + assertEqual "received 3 messages" 3 (length received) + case received of + [r0, r1, r2] -> do + assertEqual "order: msg-0 first" payload0 r0 + assertEqual "order: msg-1 second" payload1 r1 + assertEqual "order: msg-2 third" payload2 r2 + _ -> error "FAIL: expected exactly 3 messages" + +testChannelReliableSequencedDropOld :: IO () +testChannelReliableSequencedDropOld = do + putStrLn "Channel reliable sequenced drops old:" + let config = reliableSequencedConfig + ch0 = newChannel (ChannelId 0) config + now = 0 :: MonoTime + -- Receive seq 2 first (greater than remote seq 0, so accepted) + let ch1 = onMessageReceived (SequenceNum 2) "new-msg" now ch0 + -- Receive seq 0 (not greater than remote seq 2, so dropped) + let ch2 = onMessageReceived (SequenceNum 0) "old-msg" now ch1 + let (received, _ch3) = channelReceive ch2 + assertEqual "only 1 message received" 1 (length received) + assertEqual "received the newer message" "new-msg" (head received) + assertEqual "1 message dropped" 1 (chTotalDropped ch2) + +testChannelRetransmit :: IO () +testChannelRetransmit = do + putStrLn "Channel retransmission after RTO:" + let config = reliableOrderedConfig + ch0 = newChannel (ChannelId 0) config + sendTime = 0 :: MonoTime + payload = "reliable-msg" + rto = 200.0 :: Double -- RTO in milliseconds + -- Send a reliable message + case channelSend payload sendTime ch0 of + Left e -> error $ " FAIL: send failed: " ++ show e + Right (_, ch1) -> do + -- Get outgoing message (marks as sent, retryCount -> 1) + case getOutgoingMessage ch1 of + Nothing -> error " FAIL: no outgoing message" + Just (_, ch2) -> do + -- Check before RTO: no retransmits + let beforeRto = sendTime + 100000000 -- 100ms in nanoseconds + let (retransBefore, ch3) = getRetransmitMessages beforeRto rto ch2 + assertEqual "no retransmit before RTO" 0 (length retransBefore) + -- Check after RTO: should retransmit + let afterRto = sendTime + 300000000 -- 300ms in nanoseconds (> 200ms RTO) + let (retransAfter, ch4) = getRetransmitMessages afterRto rto ch3 + assertEqual "1 retransmit after RTO" 1 (length retransAfter) + assertEqual "retransmitted data" payload (cmData (head retransAfter)) + assertEqual "retransmit count incremented" 1 (chTotalRetransmits ch4) + +-------------------------------------------------------------------------------- +-- Fragment: split, reassemble, header roundtrip, cleanup, too-large +-------------------------------------------------------------------------------- + +testFragmentSplitReassemble :: IO () +testFragmentSplitReassemble = do + putStrLn "Fragment split and reassemble:" + let msgId = MessageId 42 + maxPayload = 100 + msgData = BS.replicate 250 0xAB + expectedFragCount = (BS.length msgData + maxPayload - 1) `div` maxPayload -- 3 + case fragmentMessage msgId msgData maxPayload of + Left e -> error $ " FAIL: fragmentMessage failed: " ++ show e + Right frags -> do + assertEqual "fragment count" expectedFragCount (length frags) + -- Reassemble using processFragment + let now = 0 :: MonoTime + assembler0 = newFragmentAssembler 5000.0 (1024 * 1024) + -- Feed all fragments + let (result, _assembler) = foldl feedFrag (Nothing, assembler0) frags + where + feedFrag (prevResult, asm) frag = + let (r, asm') = processFragment frag now asm + in (case r of Nothing -> prevResult; Just _ -> r, asm') + case result of + Nothing -> error " FAIL: reassembly did not produce a result" + Just reassembled -> + assertEqual "reassembled matches original" msgData reassembled + +testFragmentHeaderRoundTrip :: IO () +testFragmentHeaderRoundTrip = do + putStrLn "Fragment header serialize/deserialize roundtrip:" + let header = + FragmentHeader + { fhMessageId = MessageId 0xDEADBEEF, + fhFragmentIndex = 7, + fhFragmentCount = 15 + } + serialized = serializeFragmentHeader header + assertEqual "header size" fragmentHeaderSize (BS.length serialized) + case deserializeFragmentHeader serialized of + Nothing -> error " FAIL: deserializeFragmentHeader returned Nothing" + Just decoded -> do + assertEqual "messageId roundtrip" (fhMessageId header) (fhMessageId decoded) + assertEqual "fragmentIndex roundtrip" (fhFragmentIndex header) (fhFragmentIndex decoded) + assertEqual "fragmentCount roundtrip" (fhFragmentCount header) (fhFragmentCount decoded) + +testFragmentCleanupExpiry :: IO () +testFragmentCleanupExpiry = do + putStrLn "Fragment cleanup removes expired buffers:" + let timeoutMs = 1000.0 + assembler0 = newFragmentAssembler timeoutMs (1024 * 1024) + header = + FragmentHeader + { fhMessageId = MessageId 99, + fhFragmentIndex = 0, + fhFragmentCount = 3 + } + fragData = serializeFragmentHeader header <> BS.replicate 50 0xCC + createTime = 0 :: MonoTime + -- Process one fragment (incomplete assembly) + let (_result, assembler1) = processFragment fragData createTime assembler0 + assertEqual "1 buffer in progress" 1 (length (faBuffers assembler1)) + -- Cleanup before timeout: buffer should remain + let beforeTimeout = createTime + 500000000 -- 500ms in nanoseconds + let assembler2 = cleanupFragments beforeTimeout assembler1 + assertEqual "buffer still present before timeout" 1 (length (faBuffers assembler2)) + -- Cleanup after timeout: buffer should be removed + let afterTimeout = createTime + 1500000000 -- 1500ms in nanoseconds (> 1000ms timeout) + let assembler3 = cleanupFragments afterTimeout assembler2 + assertEqual "buffer removed after timeout" 0 (length (faBuffers assembler3)) + +testFragmentTooLarge :: IO () +testFragmentTooLarge = do + putStrLn "Fragment too many fragments:" + let msgId = MessageId 1 + tinyPayload = 1 + bigMsg = BS.replicate (maxFragmentCount + 1) 0xFF + case fragmentMessage msgId bigMsg tinyPayload of + Left TooManyFragments -> putStrLn " PASS: too many fragments returns TooManyFragments" + Right _ -> error " FAIL: expected TooManyFragments error" + +-------------------------------------------------------------------------------- +-- Security: CRC32C, rate limiting, token validation +-------------------------------------------------------------------------------- + +testCrc32Roundtrip :: IO () +testCrc32Roundtrip = do + putStrLn "CRC32C append and validate roundtrip:" + let original = "hello, network!" :: BS.ByteString + withCrc = appendCrc32 original + assertEqual "crc adds 4 bytes" (BS.length original + crc32Size) (BS.length withCrc) + case validateAndStripCrc32 withCrc of + Nothing -> error " FAIL: validateAndStripCrc32 returned Nothing on valid data" + Just stripped -> assertEqual "stripped matches original" original stripped + +testCrc32RejectCorrupt :: IO () +testCrc32RejectCorrupt = do + putStrLn "CRC32C rejects corrupted data:" + let original = "important data" :: BS.ByteString + withCrc = appendCrc32 original + -- Flip a bit in the payload area + corrupted = BS.cons (BS.index withCrc 0 + 1) (BS.tail withCrc) + case validateAndStripCrc32 corrupted of + Nothing -> putStrLn " PASS: corrupted data rejected" + Just _ -> error " FAIL: corrupted data should have been rejected" + +testRateLimiterAllow :: IO () +testRateLimiterAllow = do + putStrLn "Rate limiter allows up to max requests:" + let maxReqs = 3 + now = 1000000000 :: MonoTime -- 1 second + rl0 = newRateLimiter maxReqs now + addrKey = 12345 :: Word64 + -- Send maxReqs requests, all should be allowed + let (allowed1, rl1) = rateLimiterAllow addrKey now rl0 + let (allowed2, rl2) = rateLimiterAllow addrKey now rl1 + let (allowed3, _rl3) = rateLimiterAllow addrKey now rl2 + assertEqual "request 1 allowed" True allowed1 + assertEqual "request 2 allowed" True allowed2 + assertEqual "request 3 allowed" True allowed3 + +testRateLimiterDeny :: IO () +testRateLimiterDeny = do + putStrLn "Rate limiter denies after exceeding limit:" + let maxReqs = 2 + now = 1000000000 :: MonoTime + rl0 = newRateLimiter maxReqs now + addrKey = 99999 :: Word64 + -- Send maxReqs requests (allowed) + let (_, rl1) = rateLimiterAllow addrKey now rl0 + let (_, rl2) = rateLimiterAllow addrKey now rl1 + -- Next request should be denied + let (denied, _rl3) = rateLimiterAllow addrKey now rl2 + assertEqual "excess request denied" False denied + +testTokenValidation :: IO () +testTokenValidation = do + putStrLn "Token validation accepts valid token:" + let now = 1000000000 :: MonoTime -- 1 second + expireMs = 30000.0 -- 30 seconds + clientId = 42 :: Word64 + token = newConnectToken clientId expireMs "user-data" now + validator0 = newTokenValidator 60000.0 100 + case validateToken token now validator0 of + (Right cid, _) -> assertEqual "returns client id" clientId cid + (Left e, _) -> error $ " FAIL: valid token rejected: " ++ show e + +testTokenExpired :: IO () +testTokenExpired = do + putStrLn "Token validation rejects expired token:" + let createTime = 1000000000 :: MonoTime -- 1 second + expireMs = 5000.0 -- 5 seconds + clientId = 100 :: Word64 + token = newConnectToken clientId expireMs "data" createTime + validator = newTokenValidator 60000.0 100 + -- Validate well after expiry (10 seconds later = 10,000ms > 5000ms) + futureTime = createTime + 10000000000 -- 10 seconds in nanoseconds + case validateToken token futureTime validator of + (Left TokenExpired, _) -> putStrLn " PASS: expired token rejected with TokenExpired" + (Left e, _) -> error $ " FAIL: expected TokenExpired, got " ++ show e + (Right _, _) -> error " FAIL: expected expired token to be rejected" + +testTokenReplayed :: IO () +testTokenReplayed = do + putStrLn "Token validation rejects replayed token:" + let now = 1000000000 :: MonoTime + expireMs = 30000.0 + clientId = 77 :: Word64 + token = newConnectToken clientId expireMs "data" now + validator0 = newTokenValidator 60000.0 100 + -- First validation succeeds + case validateToken token now validator0 of + (Left e, _) -> error $ " FAIL: first validation should succeed: " ++ show e + (Right _, validator1) -> + -- Second validation with same clientId should fail as replayed + case validateToken token now validator1 of + (Left TokenReplayed, _) -> putStrLn " PASS: replayed token rejected with TokenReplayed" + (Left e, _) -> error $ " FAIL: expected TokenReplayed, got " ++ show e + (Right _, _) -> error " FAIL: expected replayed token to be rejected"