-- | DHT node state, RPC handler, and record/provider stores.
--
-- The DHTNode is the top-level coordinator for Kademlia DHT operations.
-- It owns the routing table, record store, provider store, and handles
-- both inbound (as handler) and outbound (dhtSendRequest) RPC.
--
-- The outbound sender is wired to the Switch by 'newDHTNode'; it remains
-- a record field so tests can inject mocks without a real network.
module LibP2P.DHT
  ( -- * Types
    DHTNode (..)
  , DHTMode (..)
  , PeerSession
  , ProviderEntry (..)
  , Validator (..)
    -- * Validators
  , defaultValidator
  , namespacedValidator
  , pkValidator
    -- * Construction
  , newDHTNode
  , newPeerSession
  , stopDHTNode
  , defaultQueryTimeoutMicros
    -- * Handler registration
  , registerDHTHandler
    -- * Inbound RPC handler
  , handleDHTRequest
    -- * Routing table maintenance
  , addPeerToTable
    -- * Store operations
  , storeRecord
  , lookupRecord
  , addProvider
  , getProviders
    -- * Wire helpers
  , decodePeerAddrs
    -- * Constants
  , dhtProtocolId
  , providerRecordTTL
  ) where

import Control.Concurrent.Async (Async, cancel)
import Control.Concurrent.MVar (MVar, newMVar, putMVar, takeMVar, tryTakeMVar)
import Control.Concurrent.STM
import Control.Exception (SomeException, catch, mask, onException, try)
import Data.ByteString (ByteString)
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import Data.Text (Text)
import qualified Data.Text as T
import Data.Time (NominalDiffTime, UTCTime, diffUTCTime, getCurrentTime)
import Data.Time.Format.ISO8601 (iso8601Show)
import LibP2P.Crypto.PeerId (PeerId (..), peerIdBytes)
import LibP2P.DHT.Distance (keyToDHTKey, peerIdToKey)
import LibP2P.DHT.Message
import LibP2P.DHT.RoutingTable
  ( RoutingTable
  , allPeers
  , closestPeers
  , insertPeer
  , newRoutingTable
  , removePeer
  )
import LibP2P.DHT.Types
import LibP2P.DHT.Validator
  ( Validator (..)
  , defaultValidator
  , namespacedValidator
  , pkValidator
  )
import LibP2P.Multiaddr (Multiaddr, fromBytes, toBytes)
import LibP2P.MultistreamSelect.Negotiation
  ( NegotiationResult (..)
  , StreamIO (..)
  , closeQuietly
  , negotiateInitiator
  )
import LibP2P.Switch (setStreamHandler)
import LibP2P.Switch.ConnPool (lookupConn)
import LibP2P.Switch.Types (Connection (..), MuxerSession (..), Switch (..))

-- | DHT protocol identifier for multistream-select.
dhtProtocolId :: Text
dhtProtocolId :: Text
dhtProtocolId = Text
"/ipfs/kad/1.0.0"

-- | Server or client mode.
data DHTMode = DHTServer | DHTClient
  deriving (Int -> DHTMode -> ShowS
[DHTMode] -> ShowS
DHTMode -> [Char]
(Int -> DHTMode -> ShowS)
-> (DHTMode -> [Char]) -> ([DHTMode] -> ShowS) -> Show DHTMode
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> DHTMode -> ShowS
showsPrec :: Int -> DHTMode -> ShowS
$cshow :: DHTMode -> [Char]
show :: DHTMode -> [Char]
$cshowList :: [DHTMode] -> ShowS
showList :: [DHTMode] -> ShowS
Show, DHTMode -> DHTMode -> Bool
(DHTMode -> DHTMode -> Bool)
-> (DHTMode -> DHTMode -> Bool) -> Eq DHTMode
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: DHTMode -> DHTMode -> Bool
== :: DHTMode -> DHTMode -> Bool
$c/= :: DHTMode -> DHTMode -> Bool
/= :: DHTMode -> DHTMode -> Bool
Eq)

-- | A provider record for content routing.
data ProviderEntry = ProviderEntry
  { ProviderEntry -> PeerId
peProvider  :: !PeerId
  , ProviderEntry -> [Multiaddr]
peAddrs     :: ![Multiaddr]
  , ProviderEntry -> UTCTime
peTimestamp :: !UTCTime
  } deriving (Int -> ProviderEntry -> ShowS
[ProviderEntry] -> ShowS
ProviderEntry -> [Char]
(Int -> ProviderEntry -> ShowS)
-> (ProviderEntry -> [Char])
-> ([ProviderEntry] -> ShowS)
-> Show ProviderEntry
forall a.
(Int -> a -> ShowS) -> (a -> [Char]) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> ProviderEntry -> ShowS
showsPrec :: Int -> ProviderEntry -> ShowS
$cshow :: ProviderEntry -> [Char]
show :: ProviderEntry -> [Char]
$cshowList :: [ProviderEntry] -> ShowS
showList :: [ProviderEntry] -> ShowS
Show, ProviderEntry -> ProviderEntry -> Bool
(ProviderEntry -> ProviderEntry -> Bool)
-> (ProviderEntry -> ProviderEntry -> Bool) -> Eq ProviderEntry
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: ProviderEntry -> ProviderEntry -> Bool
== :: ProviderEntry -> ProviderEntry -> Bool
$c/= :: ProviderEntry -> ProviderEntry -> Bool
/= :: ProviderEntry -> ProviderEntry -> Bool
Eq)

-- | Provider record expiration interval, per specs/kad-dht (48 hours).
-- Expired entries are pruned on read in 'getProviders'.
providerRecordTTL :: NominalDiffTime
providerRecordTTL :: NominalDiffTime
providerRecordTTL = NominalDiffTime
48 NominalDiffTime -> NominalDiffTime -> NominalDiffTime
forall a. Num a => a -> a -> a
* NominalDiffTime
3600

-- | A cached outbound @/ipfs/kad/1.0.0@ stream together with its
-- exchange lock, held as a single 'MVar' that is both.
--
-- Kademlia RPC messages carry no request identifier, so two exchanges
-- interleaved on one stream cannot be reassociated afterwards: a caller
-- reads whichever response arrives next, not necessarily its own. The
-- complete write + read exchange therefore has to be serialized per
-- peer. go-libp2p pairs its per-peer cached stream with exactly this
-- kind of exchange-wide lock (@peerMessageSender.lk@).
--
-- Making the 'MVar' hold the stream slot rather than guard a separate
-- one means replacing a dead stream is, by construction, something only
-- the caller currently holding the exchange can do.
--
-- 'psInvalid' is set when the peer's last connection closes so a caller
-- that still holds the slot cannot put a live stream back into a map
-- entry that has already been removed (go-libp2p's @invalidate()@).
data PeerSession = PeerSession
  { PeerSession -> MVar (Maybe StreamIO)
psSlot    :: !(MVar (Maybe StreamIO))
  , PeerSession -> TVar Bool
psInvalid :: !(TVar Bool)
  }

-- | Top-level DHT node state.
data DHTNode = DHTNode
  { DHTNode -> Switch
dhtSwitch        :: !Switch
  , DHTNode -> TVar RoutingTable
dhtRoutingTable  :: !(TVar RoutingTable)
  , DHTNode -> TVar (Map ByteString DHTRecord)
dhtRecordStore   :: !(TVar (Map ByteString DHTRecord))
  , DHTNode -> TVar (Map ByteString [ProviderEntry])
dhtProviderStore :: !(TVar (Map ByteString [ProviderEntry]))
  , DHTNode -> DHTKey
dhtLocalKey      :: !DHTKey
  , DHTNode -> PeerId
dhtLocalPeerId   :: !PeerId
  , DHTNode -> DHTMode
dhtMode          :: !DHTMode
  , DHTNode -> Validator
dhtValidator     :: !Validator
    -- ^ Record validator applied to PUT_VALUE records before storage
    -- (and available to GET_VALUE conflict resolution). Defaults to
    -- 'defaultValidator' (the @/pk/@ namespace).
  , DHTNode -> TVar (Map PeerId PeerSession)
dhtStreams       :: !(TVar (Map PeerId PeerSession))
    -- ^ Cached outbound @/ipfs/kad/1.0.0@ sessions, one per peer
    -- (go-libp2p reuses a single long-lived stream per peer)
  , DHTNode -> PeerId -> DHTMessage -> IO (Either [Char] DHTMessage)
dhtSendRequest   :: !(PeerId -> DHTMessage -> IO (Either String DHTMessage))
    -- ^ Outbound RPC sender. Wired to the Switch by 'newDHTNode';
    -- kept as a field so tests can inject mocks.
  , DHTNode -> IORef (Maybe (Connection -> IO ()))
dhtDisconnectHook :: !(IORef (Maybe (Connection -> IO ())))
    -- ^ Disconnect notifier; 'stopDHTNode' clears it so a stopped node
    -- does not keep a callback alive on the Switch.
  , DHTNode -> Int
dhtQueryTimeout :: !Int
    -- ^ Deadline for one iterative lookup / bootstrap run, in microseconds.
    -- Default 'defaultQueryTimeoutMicros' (10s, specs/kad-dht QueryTimeout).
  , DHTNode -> TVar (Maybe (Async ()))
dhtBootstrapWorker :: !(TVar (Maybe (Async ())))
    -- ^ Periodic bootstrap loop; cancelled by 'stopDHTNode'.
  }

-- | Default query/bootstrap timeout: 10 seconds (specs/kad-dht).
defaultQueryTimeoutMicros :: Int
defaultQueryTimeoutMicros :: Int
defaultQueryTimeoutMicros = Int
10000000

-- | Create an empty or pre-loaded peer session (tests inject the latter).
newPeerSession :: Maybe StreamIO -> IO PeerSession
newPeerSession :: Maybe StreamIO -> IO PeerSession
newPeerSession Maybe StreamIO
slot = MVar (Maybe StreamIO) -> TVar Bool -> PeerSession
PeerSession (MVar (Maybe StreamIO) -> TVar Bool -> PeerSession)
-> IO (MVar (Maybe StreamIO)) -> IO (TVar Bool -> PeerSession)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe StreamIO -> IO (MVar (Maybe StreamIO))
forall a. a -> IO (MVar a)
newMVar Maybe StreamIO
slot IO (TVar Bool -> PeerSession) -> IO (TVar Bool) -> IO PeerSession
forall a b. IO (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> Bool -> IO (TVar Bool)
forall a. a -> IO (TVar a)
newTVarIO Bool
False

-- | Create a new DHT node with the outbound sender wired to the Switch.
--
-- Registers a disconnect notifier so a cached session is dropped when
-- the peer's last connection closes (#279).
newDHTNode :: Switch -> DHTMode -> IO DHTNode
newDHTNode :: Switch -> DHTMode -> IO DHTNode
newDHTNode Switch
sw DHTMode
mode = do
  let localPid :: PeerId
localPid = Switch -> PeerId
swLocalPeerId Switch
sw
  rt <- RoutingTable -> IO (TVar RoutingTable)
forall a. a -> IO (TVar a)
newTVarIO (PeerId -> RoutingTable
newRoutingTable PeerId
localPid)
  records <- newTVarIO Map.empty
  providers <- newTVarIO Map.empty
  streams <- newTVarIO Map.empty
  hook <- newIORef Nothing
  worker <- newTVarIO Nothing
  let node = DHTNode
        { dhtSwitch :: Switch
dhtSwitch           = Switch
sw
        , dhtRoutingTable :: TVar RoutingTable
dhtRoutingTable     = TVar RoutingTable
rt
        , dhtRecordStore :: TVar (Map ByteString DHTRecord)
dhtRecordStore      = TVar (Map ByteString DHTRecord)
records
        , dhtProviderStore :: TVar (Map ByteString [ProviderEntry])
dhtProviderStore    = TVar (Map ByteString [ProviderEntry])
providers
        , dhtLocalKey :: DHTKey
dhtLocalKey         = PeerId -> DHTKey
peerIdToKey PeerId
localPid
        , dhtLocalPeerId :: PeerId
dhtLocalPeerId      = PeerId
localPid
        , dhtMode :: DHTMode
dhtMode             = DHTMode
mode
        , dhtValidator :: Validator
dhtValidator        = Validator
defaultValidator
        , dhtStreams :: TVar (Map PeerId PeerSession)
dhtStreams          = TVar (Map PeerId PeerSession)
streams
        , dhtSendRequest :: PeerId -> DHTMessage -> IO (Either [Char] DHTMessage)
dhtSendRequest      = Switch
-> TVar (Map PeerId PeerSession)
-> PeerId
-> DHTMessage
-> IO (Either [Char] DHTMessage)
sendRequestViaSwitch Switch
sw TVar (Map PeerId PeerSession)
streams
        , dhtDisconnectHook :: IORef (Maybe (Connection -> IO ()))
dhtDisconnectHook   = IORef (Maybe (Connection -> IO ()))
hook
        , dhtQueryTimeout :: Int
dhtQueryTimeout     = Int
defaultQueryTimeoutMicros
        , dhtBootstrapWorker :: TVar (Maybe (Async ()))
dhtBootstrapWorker  = TVar (Maybe (Async ()))
worker
        }
  writeIORef hook (Just (dropCachedSession sw streams))
  atomically $ modifyTVar' (swDisconnectNotifiers sw) (runDisconnectHook hook :)
  pure node

-- | Stop the DHT node: drop cached sessions and deregister the disconnect
-- notifier so a stopped node cannot keep a callback alive on the Switch.
stopDHTNode :: DHTNode -> IO ()
stopDHTNode :: DHTNode -> IO ()
stopDHTNode DHTNode
node = do
  IORef (Maybe (Connection -> IO ()))
-> Maybe (Connection -> IO ()) -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef (DHTNode -> IORef (Maybe (Connection -> IO ()))
dhtDisconnectHook DHTNode
node) Maybe (Connection -> IO ())
forall a. Maybe a
Nothing
  mWorker <- STM (Maybe (Async ())) -> IO (Maybe (Async ()))
forall a. STM a -> IO a
atomically (STM (Maybe (Async ())) -> IO (Maybe (Async ())))
-> STM (Maybe (Async ())) -> IO (Maybe (Async ()))
forall a b. (a -> b) -> a -> b
$ do
    w <- TVar (Maybe (Async ())) -> STM (Maybe (Async ()))
forall a. TVar a -> STM a
readTVar (DHTNode -> TVar (Maybe (Async ()))
dhtBootstrapWorker DHTNode
node)
    writeTVar (dhtBootstrapWorker node) Nothing
    pure w
  mapM_ cancel mWorker
  sessions <- atomically $ do
    m <- readTVar (dhtStreams node)
    writeTVar (dhtStreams node) Map.empty
    pure (Map.elems m)
  mapM_ invalidateHeldSession sessions

-- | Register the DHT handler on the Switch.
--
-- Per specs/kad-dht (client and server mode), nodes operating in client
-- mode do not offer the Kademlia protocol identifier for incoming
-- streams, so this is a no-op for 'DHTClient' nodes: they keep issuing
-- outbound queries via 'dhtSendRequest' but never serve inbound RPC.
registerDHTHandler :: DHTNode -> IO ()
registerDHTHandler :: DHTNode -> IO ()
registerDHTHandler DHTNode
node = case DHTNode -> DHTMode
dhtMode DHTNode
node of
  DHTMode
DHTClient -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
  DHTMode
DHTServer ->
    Switch -> Text -> StreamHandler -> IO ()
setStreamHandler (DHTNode -> Switch
dhtSwitch DHTNode
node) Text
dhtProtocolId
      (\Connection
conn StreamIO
stream -> DHTNode -> StreamIO -> PeerId -> IO ()
handleDHTRequest DHTNode
node StreamIO
stream (Connection -> PeerId
connPeerId Connection
conn))

-- | Handle an inbound DHT stream.
--
-- Per specs/kad-dht, implementations must handle additional RPC request
-- messages on the same incoming stream: go-libp2p keeps one long-lived
-- stream per peer and pipelines requests over it. Loop until the stream
-- errors, is reset, or reaches EOF.
handleDHTRequest :: DHTNode -> StreamIO -> PeerId -> IO ()
handleDHTRequest :: DHTNode -> StreamIO -> PeerId -> IO ()
handleDHTRequest DHTNode
node StreamIO
stream PeerId
remotePeerId = IO ()
loop
  where
    loop :: IO ()
loop = do
      result <- IO (Either [Char] ())
-> IO (Either SomeException (Either [Char] ()))
forall e a. Exception e => IO a -> IO (Either e a)
try (IO (Either [Char] ())
 -> IO (Either SomeException (Either [Char] ())))
-> IO (Either [Char] ())
-> IO (Either SomeException (Either [Char] ()))
forall a b. (a -> b) -> a -> b
$ do
        readResult <- StreamIO -> Int -> IO (Either [Char] DHTMessage)
readFramedMessage StreamIO
stream Int
maxDHTMessageSize
        case readResult of
          Left [Char]
err -> Either [Char] () -> IO (Either [Char] ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Char] -> Either [Char] ()
forall a b. a -> Either a b
Left [Char]
err)
          Right DHTMessage
msg -> do
            response <- DHTNode -> DHTMessage -> PeerId -> IO DHTMessage
processRequest DHTNode
node DHTMessage
msg PeerId
remotePeerId
            writeFramedMessage stream response
            -- Routing-table growth: a peer speaking the DHT protocol to
            -- us is a live contact; insert (or refresh) it. Note this
            -- cannot distinguish client-mode senders (the spec would
            -- exclude them) without identify-provided protocol lists.
            now <- getCurrentTime
            _ <- addPeerToTable node BucketEntry
              { entryPeerId   = remotePeerId
              , entryKey      = peerIdToKey remotePeerId
              , entryAddrs    = []
              , entryLastSeen = now
              , entryConnType = Connected
              }
            pure (Right ())
      case result of
        Left (SomeException
_ :: SomeException) -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()  -- Stream closed or reset
        Right (Left [Char]
_err) -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()          -- Framing/decode error: stop serving
        Right (Right ()) -> IO ()
loop

-- | Process a single DHT request and produce a response.
processRequest :: DHTNode -> DHTMessage -> PeerId -> IO DHTMessage
processRequest :: DHTNode -> DHTMessage -> PeerId -> IO DHTMessage
processRequest DHTNode
node DHTMessage
msg PeerId
remotePeerId =
  case DHTMessage -> MessageType
msgType DHTMessage
msg of
    MessageType
FindNode -> DHTNode -> DHTMessage -> IO DHTMessage
handleFindNode DHTNode
node DHTMessage
msg
    MessageType
GetValue -> DHTNode -> DHTMessage -> IO DHTMessage
handleGetValue DHTNode
node DHTMessage
msg
    MessageType
PutValue -> DHTNode -> DHTMessage -> IO DHTMessage
handlePutValue DHTNode
node DHTMessage
msg
    MessageType
AddProvider -> DHTNode -> DHTMessage -> PeerId -> IO DHTMessage
handleAddProvider DHTNode
node DHTMessage
msg PeerId
remotePeerId
    MessageType
GetProviders -> DHTNode -> DHTMessage -> IO DHTMessage
handleGetProviders DHTNode
node DHTMessage
msg

-- | FIND_NODE: return k closest peers to the requested key.
handleFindNode :: DHTNode -> DHTMessage -> IO DHTMessage
handleFindNode :: DHTNode -> DHTMessage -> IO DHTMessage
handleFindNode DHTNode
node DHTMessage
msg = do
  rt <- TVar RoutingTable -> IO RoutingTable
forall a. TVar a -> IO a
readTVarIO (DHTNode -> TVar RoutingTable
dhtRoutingTable DHTNode
node)
  -- The wire key is raw (a binary peer ID); the spec distance metric is
  -- XOR over SHA-256 digests, so hash before comparing.
  let targetKey = ByteString -> DHTKey
keyToDHTKey (DHTMessage -> ByteString
msgKey DHTMessage
msg)
      closest = DHTKey -> Int -> RoutingTable -> [BucketEntry]
closestPeers DHTKey
targetKey Int
kValue RoutingTable
rt
      peers = (BucketEntry -> DHTPeer) -> [BucketEntry] -> [DHTPeer]
forall a b. (a -> b) -> [a] -> [b]
map BucketEntry -> DHTPeer
entryToDHTPeer [BucketEntry]
closest
  pure emptyDHTMessage
    { msgType = FindNode
    , msgCloserPeers = peers
    }

-- | GET_VALUE: return stored record + k closest peers.
handleGetValue :: DHTNode -> DHTMessage -> IO DHTMessage
handleGetValue :: DHTNode -> DHTMessage -> IO DHTMessage
handleGetValue DHTNode
node DHTMessage
msg = do
  rt <- TVar RoutingTable -> IO RoutingTable
forall a. TVar a -> IO a
readTVarIO (DHTNode -> TVar RoutingTable
dhtRoutingTable DHTNode
node)
  records <- readTVarIO (dhtRecordStore node)
  let key = DHTMessage -> ByteString
msgKey DHTMessage
msg
      -- Store lookup uses the raw key; distance uses its SHA-256.
      targetKey = ByteString -> DHTKey
keyToDHTKey ByteString
key
      closest = DHTKey -> Int -> RoutingTable -> [BucketEntry]
closestPeers DHTKey
targetKey Int
kValue RoutingTable
rt
      peers = (BucketEntry -> DHTPeer) -> [BucketEntry] -> [DHTPeer]
forall a b. (a -> b) -> [a] -> [b]
map BucketEntry -> DHTPeer
entryToDHTPeer [BucketEntry]
closest
      rec = ByteString -> Map ByteString DHTRecord -> Maybe DHTRecord
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup ByteString
key Map ByteString DHTRecord
records
  pure emptyDHTMessage
    { msgType = GetValue
    , msgRecord = rec
    , msgCloserPeers = peers
    }

-- | PUT_VALUE: validate, store with a receiver-set timestamp, and echo.
--
-- Per specs/kad-dht (Entry validation), incoming records are validated
-- before being stored: the record key must match the message key and
-- the configured 'dhtValidator' must accept the key/value binding
-- (e.g. @/pk/@ records must carry the public key hashing to the key's
-- multihash). Rejected records are neither stored nor echoed back.
--
-- The stored record's @timeReceived@ is set by the receiver (Record
-- field 5: "Time the record was received, set by receiver"), never
-- taken from the sender's claim.
handlePutValue :: DHTNode -> DHTMessage -> IO DHTMessage
handlePutValue :: DHTNode -> DHTMessage -> IO DHTMessage
handlePutValue DHTNode
node DHTMessage
msg = do
  case DHTMessage -> Maybe DHTRecord
msgRecord DHTMessage
msg of
    Maybe DHTRecord
Nothing -> DHTMessage -> IO DHTMessage
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DHTMessage
rejected
    Just DHTRecord
rec
      | DHTRecord -> ByteString
recKey DHTRecord
rec ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
/= DHTMessage -> ByteString
msgKey DHTMessage
msg -> DHTMessage -> IO DHTMessage
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DHTMessage
rejected
      | Bool
otherwise ->
          case Validator -> ByteString -> ByteString -> Either [Char] ()
valValidate (DHTNode -> Validator
dhtValidator DHTNode
node) (DHTRecord -> ByteString
recKey DHTRecord
rec) (DHTRecord -> ByteString
recValue DHTRecord
rec) of
            Left [Char]
_err -> DHTMessage -> IO DHTMessage
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure DHTMessage
rejected
            Right () -> do
              now <- IO UTCTime
getCurrentTime
              storeRecord node rec { recTimeReceived = T.pack (iso8601Show now) }
              pure emptyDHTMessage
                { msgType = PutValue
                , msgKey = msgKey msg
                , msgRecord = Just rec
                }
  where
    rejected :: DHTMessage
rejected = DHTMessage
emptyDHTMessage { msgType = PutValue }

-- | ADD_PROVIDER: verify sender and store provider record.
handleAddProvider :: DHTNode -> DHTMessage -> PeerId -> IO DHTMessage
handleAddProvider :: DHTNode -> DHTMessage -> PeerId -> IO DHTMessage
handleAddProvider DHTNode
node DHTMessage
msg PeerId
remotePeerId = do
  now <- IO UTCTime
getCurrentTime
  -- Verify that provider peers match sender's Peer ID
  let validProviders = (DHTPeer -> Bool) -> [DHTPeer] -> [DHTPeer]
forall a. (a -> Bool) -> [a] -> [a]
filter (\DHTPeer
p -> DHTPeer -> ByteString
dhtPeerId DHTPeer
p ByteString -> ByteString -> Bool
forall a. Eq a => a -> a -> Bool
== PeerId -> ByteString
peerIdBytes PeerId
remotePeerId) (DHTMessage -> [DHTPeer]
msgProviderPeers DHTMessage
msg)
  -- Store each valid provider keyed by msgKey
  mapM_ (\DHTPeer
p -> DHTNode -> ByteString -> ProviderEntry -> IO ()
addProvider DHTNode
node (DHTMessage -> ByteString
msgKey DHTMessage
msg) (DHTPeer -> UTCTime -> ProviderEntry
dhtPeerToProvider DHTPeer
p UTCTime
now)) validProviders
  pure emptyDHTMessage { msgType = AddProvider }

-- | GET_PROVIDERS: return stored providers + k closest peers.
handleGetProviders :: DHTNode -> DHTMessage -> IO DHTMessage
handleGetProviders :: DHTNode -> DHTMessage -> IO DHTMessage
handleGetProviders DHTNode
node DHTMessage
msg = do
  rt <- TVar RoutingTable -> IO RoutingTable
forall a. TVar a -> IO a
readTVarIO (DHTNode -> TVar RoutingTable
dhtRoutingTable DHTNode
node)
  let key = DHTMessage -> ByteString
msgKey DHTMessage
msg
      -- Store lookup uses the raw key; distance uses its SHA-256.
      targetKey = ByteString -> DHTKey
keyToDHTKey ByteString
key
      closest = DHTKey -> Int -> RoutingTable -> [BucketEntry]
closestPeers DHTKey
targetKey Int
kValue RoutingTable
rt
      closerPeers = (BucketEntry -> DHTPeer) -> [BucketEntry] -> [DHTPeer]
forall a b. (a -> b) -> [a] -> [b]
map BucketEntry -> DHTPeer
entryToDHTPeer [BucketEntry]
closest
  -- getProviders prunes entries past the 48h provider TTL on read.
  providers <- getProviders node key
  let providerPeers = (ProviderEntry -> DHTPeer) -> [ProviderEntry] -> [DHTPeer]
forall a b. (a -> b) -> [a] -> [b]
map ProviderEntry -> DHTPeer
providerToDHTPeer [ProviderEntry]
providers
  pure emptyDHTMessage
    { msgType = GetProviders
    , msgCloserPeers = closerPeers
    , msgProviderPeers = providerPeers
    }

-- Routing table maintenance

-- | Insert a peer into the routing table, applying the Kademlia
-- full-bucket eviction policy.
--
-- When the target bucket is full, the least-recently-seen peer is
-- probed with a FIND_NODE request (the DHT liveness check; our
-- 'MessageType' has no PING, and go-libp2p likewise treats any
-- successful RPC as proof of liveness):
--
-- * if the LRS peer answers, it is kept (refreshed to
--   most-recently-seen) and the new peer is dropped ('BucketFull');
-- * if it does not answer, it is evicted and the new peer takes its
--   place ('Inserted').
addPeerToTable :: DHTNode -> BucketEntry -> IO InsertResult
addPeerToTable :: DHTNode -> BucketEntry -> IO InsertResult
addPeerToTable DHTNode
node BucketEntry
entry = do
  result <- STM InsertResult -> IO InsertResult
forall a. STM a -> IO a
atomically (STM InsertResult -> IO InsertResult)
-> STM InsertResult -> IO InsertResult
forall a b. (a -> b) -> a -> b
$ do
    rt <- TVar RoutingTable -> STM RoutingTable
forall a. TVar a -> STM a
readTVar (DHTNode -> TVar RoutingTable
dhtRoutingTable DHTNode
node)
    let (rt', res) = insertPeer entry rt
    writeTVar (dhtRoutingTable node) rt'
    pure res
  case result of
    BucketFull PeerId
lrs -> PeerId -> IO InsertResult
evictOrKeep PeerId
lrs
    InsertResult
other -> InsertResult -> IO InsertResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure InsertResult
other
  where
    evictOrKeep :: PeerId -> IO InsertResult
evictOrKeep PeerId
lrs = do
      let ping :: DHTMessage
ping = DHTMessage
emptyDHTMessage { msgType = FindNode, msgKey = peerIdBytes lrs }
      response <- DHTNode -> PeerId -> DHTMessage -> IO (Either [Char] DHTMessage)
dhtSendRequest DHTNode
node PeerId
lrs DHTMessage
ping
        IO (Either [Char] DHTMessage)
-> (SomeException -> IO (Either [Char] DHTMessage))
-> IO (Either [Char] DHTMessage)
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
`catch` \(SomeException
e :: SomeException) -> Either [Char] DHTMessage -> IO (Either [Char] DHTMessage)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Char] -> Either [Char] DHTMessage
forall a b. a -> Either a b
Left (SomeException -> [Char]
forall a. Show a => a -> [Char]
show SomeException
e))
      case response of
        Right DHTMessage
_ -> do
          -- LRS peer is alive: move it to most-recently-seen and drop
          -- the newcomer.
          now <- IO UTCTime
getCurrentTime
          atomically $ modifyTVar' (dhtRoutingTable node) (refreshPeer lrs now)
          pure (BucketFull lrs)
        Left [Char]
_ -> do
          -- LRS peer is dead: evict it and insert the newcomer.
          STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar RoutingTable -> (RoutingTable -> RoutingTable) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' (DHTNode -> TVar RoutingTable
dhtRoutingTable DHTNode
node) ((RoutingTable -> RoutingTable) -> STM ())
-> (RoutingTable -> RoutingTable) -> STM ()
forall a b. (a -> b) -> a -> b
$ \RoutingTable
rt ->
            (RoutingTable, InsertResult) -> RoutingTable
forall a b. (a, b) -> a
fst (BucketEntry -> RoutingTable -> (RoutingTable, InsertResult)
insertPeer BucketEntry
entry (PeerId -> RoutingTable -> RoutingTable
removePeer PeerId
lrs RoutingTable
rt))
          InsertResult -> IO InsertResult
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure InsertResult
Inserted

-- | Move an existing peer to the most-recently-seen position of its
-- bucket with a fresh last-seen timestamp. No-op if the peer is gone.
refreshPeer :: PeerId -> UTCTime -> RoutingTable -> RoutingTable
refreshPeer :: PeerId -> UTCTime -> RoutingTable -> RoutingTable
refreshPeer PeerId
pid UTCTime
now RoutingTable
rt =
  case (BucketEntry -> Bool) -> [BucketEntry] -> [BucketEntry]
forall a. (a -> Bool) -> [a] -> [a]
filter ((PeerId -> PeerId -> Bool
forall a. Eq a => a -> a -> Bool
== PeerId
pid) (PeerId -> Bool) -> (BucketEntry -> PeerId) -> BucketEntry -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. BucketEntry -> PeerId
entryPeerId) (RoutingTable -> [BucketEntry]
allPeers RoutingTable
rt) of
    (BucketEntry
e : [BucketEntry]
_) -> (RoutingTable, InsertResult) -> RoutingTable
forall a b. (a, b) -> a
fst (BucketEntry -> RoutingTable -> (RoutingTable, InsertResult)
insertPeer BucketEntry
e { entryLastSeen = now } RoutingTable
rt)
    [] -> RoutingTable
rt

-- Outbound RPC

-- | Send a DHT request to a peer over the Switch.
--
-- Reuses a cached @/ipfs/kad/1.0.0@ stream per peer when one exists
-- (go-libp2p also keeps one long-lived stream per peer); otherwise opens
-- a new muxer stream on an existing connection and negotiates the
-- protocol. A failed exchange on a cached stream evicts it and retries
-- once on a fresh stream.
--
-- Requests to one peer are serialized, not pipelined: the caller holds
-- the peer's 'PeerSession' for the whole write + read exchange. Requests
-- to different peers stay concurrent.
sendRequestViaSwitch
  :: Switch
  -> TVar (Map PeerId PeerSession)
  -> PeerId
  -> DHTMessage
  -> IO (Either String DHTMessage)
sendRequestViaSwitch :: Switch
-> TVar (Map PeerId PeerSession)
-> PeerId
-> DHTMessage
-> IO (Either [Char] DHTMessage)
sendRequestViaSwitch Switch
sw TVar (Map PeerId PeerSession)
sessionsVar PeerId
pid DHTMessage
request = do
  session <- TVar (Map PeerId PeerSession) -> PeerId -> IO PeerSession
peerSession TVar (Map PeerId PeerSession)
sessionsVar PeerId
pid
  -- Taking the session out for the duration of the exchange is what keeps
  -- a second caller from reading this caller's response. If the exchange
  -- is interrupted (a query deadline, say) the stream may be left holding
  -- a partial request or an unread reply, so it is closed and the slot
  -- put back empty rather than handed to the next caller.
  mask $ \forall a. IO a -> IO a
restore -> do
    cached <- MVar (Maybe StreamIO) -> IO (Maybe StreamIO)
forall a. MVar a -> IO a
takeMVar (PeerSession -> MVar (Maybe StreamIO)
psSlot PeerSession
session)
    let abandon = do
          (StreamIO -> IO ()) -> Maybe StreamIO -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ StreamIO -> IO ()
closeQuietly Maybe StreamIO
cached
          MVar (Maybe StreamIO) -> Maybe StreamIO -> IO ()
forall a. MVar a -> a -> IO ()
putMVar (PeerSession -> MVar (Maybe StreamIO)
psSlot PeerSession
session) Maybe StreamIO
forall a. Maybe a
Nothing
    (slot, result) <- restore (exchange cached) `onException` abandon
    commitSession session slot
    pure result
  where
    exchange :: Maybe StreamIO -> IO (Maybe StreamIO, Either [Char] DHTMessage)
exchange Maybe StreamIO
Nothing = IO (Maybe StreamIO, Either [Char] DHTMessage)
openAndExchange
    exchange (Just StreamIO
stream) = do
      result <- StreamIO -> DHTMessage -> IO (Either [Char] DHTMessage)
exchangeFramed StreamIO
stream DHTMessage
request
      case result of
        Right DHTMessage
resp -> (Maybe StreamIO, Either [Char] DHTMessage)
-> IO (Maybe StreamIO, Either [Char] DHTMessage)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StreamIO -> Maybe StreamIO
forall a. a -> Maybe a
Just StreamIO
stream, DHTMessage -> Either [Char] DHTMessage
forall a b. b -> Either a b
Right DHTMessage
resp)
        Left [Char]
_ -> do
          -- Cached stream is dead: close it and retry once on a fresh one.
          StreamIO -> IO ()
closeQuietly StreamIO
stream
          IO (Maybe StreamIO, Either [Char] DHTMessage)
openAndExchange

    openAndExchange :: IO (Maybe StreamIO, Either [Char] DHTMessage)
openAndExchange = do
      opened <- Switch -> PeerId -> IO (Either [Char] StreamIO)
openDHTStream Switch
sw PeerId
pid
      case opened of
        Left [Char]
err -> (Maybe StreamIO, Either [Char] DHTMessage)
-> IO (Maybe StreamIO, Either [Char] DHTMessage)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe StreamIO
forall a. Maybe a
Nothing, [Char] -> Either [Char] DHTMessage
forall a b. a -> Either a b
Left [Char]
err)
        Right StreamIO
stream ->
          (IO (Maybe StreamIO, Either [Char] DHTMessage)
-> IO () -> IO (Maybe StreamIO, Either [Char] DHTMessage)
forall a b. IO a -> IO b -> IO a
`onException` StreamIO -> IO ()
closeQuietly StreamIO
stream) (IO (Maybe StreamIO, Either [Char] DHTMessage)
 -> IO (Maybe StreamIO, Either [Char] DHTMessage))
-> IO (Maybe StreamIO, Either [Char] DHTMessage)
-> IO (Maybe StreamIO, Either [Char] DHTMessage)
forall a b. (a -> b) -> a -> b
$ do
            result <- StreamIO -> DHTMessage -> IO (Either [Char] DHTMessage)
exchangeFramed StreamIO
stream DHTMessage
request
            case result of
              Left [Char]
err -> do
                StreamIO -> IO ()
closeQuietly StreamIO
stream
                (Maybe StreamIO, Either [Char] DHTMessage)
-> IO (Maybe StreamIO, Either [Char] DHTMessage)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Maybe StreamIO
forall a. Maybe a
Nothing, [Char] -> Either [Char] DHTMessage
forall a b. a -> Either a b
Left [Char]
err)
              Either [Char] DHTMessage
ok -> (Maybe StreamIO, Either [Char] DHTMessage)
-> IO (Maybe StreamIO, Either [Char] DHTMessage)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (StreamIO -> Maybe StreamIO
forall a. a -> Maybe a
Just StreamIO
stream, Either [Char] DHTMessage
ok)

-- | Look up the peer's session, creating an empty one on first contact.
--
-- Two callers racing to create the same session agree on one: the loser
-- discards the 'MVar' it just allocated, so the exchange lock is never
-- split in two.
peerSession :: TVar (Map PeerId PeerSession) -> PeerId -> IO PeerSession
peerSession :: TVar (Map PeerId PeerSession) -> PeerId -> IO PeerSession
peerSession TVar (Map PeerId PeerSession)
sessionsVar PeerId
pid = do
  existing <- PeerId -> Map PeerId PeerSession -> Maybe PeerSession
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup PeerId
pid (Map PeerId PeerSession -> Maybe PeerSession)
-> IO (Map PeerId PeerSession) -> IO (Maybe PeerSession)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (Map PeerId PeerSession) -> IO (Map PeerId PeerSession)
forall a. TVar a -> IO a
readTVarIO TVar (Map PeerId PeerSession)
sessionsVar
  case existing of
    Just PeerSession
session -> PeerSession -> IO PeerSession
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure PeerSession
session
    Maybe PeerSession
Nothing -> do
      fresh <- Maybe StreamIO -> IO PeerSession
newPeerSession Maybe StreamIO
forall a. Maybe a
Nothing
      atomically $ do
        sessions <- readTVar sessionsVar
        case Map.lookup pid sessions of
          Just PeerSession
winner -> PeerSession -> STM PeerSession
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure PeerSession
winner
          Maybe PeerSession
Nothing -> do
            TVar (Map PeerId PeerSession) -> Map PeerId PeerSession -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Map PeerId PeerSession)
sessionsVar (PeerId
-> PeerSession -> Map PeerId PeerSession -> Map PeerId PeerSession
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert PeerId
pid PeerSession
fresh Map PeerId PeerSession
sessions)
            PeerSession -> STM PeerSession
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure PeerSession
fresh

-- | Put a stream back only if the session is still valid. An invalidated
-- session has been removed from the map; caching into it would leak a
-- live stream that no later caller can close.
commitSession :: PeerSession -> Maybe StreamIO -> IO ()
commitSession :: PeerSession -> Maybe StreamIO -> IO ()
commitSession PeerSession
session Maybe StreamIO
slot = do
  invalid <- TVar Bool -> IO Bool
forall a. TVar a -> IO a
readTVarIO (PeerSession -> TVar Bool
psInvalid PeerSession
session)
  if invalid
    then do
      mapM_ closeQuietly slot
      putMVar (psSlot session) Nothing
    else putMVar (psSlot session) slot

-- | Run the DHT disconnect hook if it has not been deregistered.
runDisconnectHook :: IORef (Maybe (Connection -> IO ())) -> Connection -> IO ()
runDisconnectHook :: IORef (Maybe (Connection -> IO ())) -> Connection -> IO ()
runDisconnectHook IORef (Maybe (Connection -> IO ()))
hook Connection
conn = do
  mfn <- IORef (Maybe (Connection -> IO ()))
-> IO (Maybe (Connection -> IO ()))
forall a. IORef a -> IO a
readIORef IORef (Maybe (Connection -> IO ()))
hook
  mapM_ ($ conn) mfn

-- | Drop the cached session when this was the peer's last connection.
dropCachedSession
  :: Switch -> TVar (Map PeerId PeerSession) -> Connection -> IO ()
dropCachedSession :: Switch -> TVar (Map PeerId PeerSession) -> Connection -> IO ()
dropCachedSession Switch
sw TVar (Map PeerId PeerSession)
sessionsVar Connection
conn = do
  remaining <- STM (Maybe Connection) -> IO (Maybe Connection)
forall a. STM a -> IO a
atomically (STM (Maybe Connection) -> IO (Maybe Connection))
-> STM (Maybe Connection) -> IO (Maybe Connection)
forall a b. (a -> b) -> a -> b
$ TVar (Map PeerId [Connection]) -> PeerId -> STM (Maybe Connection)
lookupConn (Switch -> TVar (Map PeerId [Connection])
swConnPool Switch
sw) (Connection -> PeerId
connPeerId Connection
conn)
  case remaining of
    Just Connection
_  -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    Maybe Connection
Nothing -> TVar (Map PeerId PeerSession) -> PeerId -> IO ()
invalidatePeerSession TVar (Map PeerId PeerSession)
sessionsVar (Connection -> PeerId
connPeerId Connection
conn)

-- | Remove the peer's session from the map, mark it invalid, and close
-- the stream if no exchange currently holds it.
invalidatePeerSession :: TVar (Map PeerId PeerSession) -> PeerId -> IO ()
invalidatePeerSession :: TVar (Map PeerId PeerSession) -> PeerId -> IO ()
invalidatePeerSession TVar (Map PeerId PeerSession)
sessionsVar PeerId
pid = do
  mSession <- STM (Maybe PeerSession) -> IO (Maybe PeerSession)
forall a. STM a -> IO a
atomically (STM (Maybe PeerSession) -> IO (Maybe PeerSession))
-> STM (Maybe PeerSession) -> IO (Maybe PeerSession)
forall a b. (a -> b) -> a -> b
$ do
    sessions <- TVar (Map PeerId PeerSession) -> STM (Map PeerId PeerSession)
forall a. TVar a -> STM a
readTVar TVar (Map PeerId PeerSession)
sessionsVar
    case Map.lookup pid sessions of
      Maybe PeerSession
Nothing -> Maybe PeerSession -> STM (Maybe PeerSession)
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe PeerSession
forall a. Maybe a
Nothing
      Just PeerSession
session -> do
        TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (PeerSession -> TVar Bool
psInvalid PeerSession
session) Bool
True
        TVar (Map PeerId PeerSession) -> Map PeerId PeerSession -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar TVar (Map PeerId PeerSession)
sessionsVar (PeerId -> Map PeerId PeerSession -> Map PeerId PeerSession
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete PeerId
pid Map PeerId PeerSession
sessions)
        Maybe PeerSession -> STM (Maybe PeerSession)
forall a. a -> STM a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (PeerSession -> Maybe PeerSession
forall a. a -> Maybe a
Just PeerSession
session)
  mapM_ invalidateHeldSession mSession

-- | Close a held session's stream unless an in-flight exchange owns the slot.
invalidateHeldSession :: PeerSession -> IO ()
invalidateHeldSession :: PeerSession -> IO ()
invalidateHeldSession PeerSession
session = do
  STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (PeerSession -> TVar Bool
psInvalid PeerSession
session) Bool
True
  mSlot <- MVar (Maybe StreamIO) -> IO (Maybe (Maybe StreamIO))
forall a. MVar a -> IO (Maybe a)
tryTakeMVar (PeerSession -> MVar (Maybe StreamIO)
psSlot PeerSession
session)
  case mSlot of
    Maybe (Maybe StreamIO)
Nothing -> () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    Just Maybe StreamIO
slot -> do
      (StreamIO -> IO ()) -> Maybe StreamIO -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ StreamIO -> IO ()
closeQuietly Maybe StreamIO
slot
      MVar (Maybe StreamIO) -> Maybe StreamIO -> IO ()
forall a. MVar a -> a -> IO ()
putMVar (PeerSession -> MVar (Maybe StreamIO)
psSlot PeerSession
session) Maybe StreamIO
forall a. Maybe a
Nothing

-- | Write a framed request and read the framed response, capturing IO errors.
exchangeFramed :: StreamIO -> DHTMessage -> IO (Either String DHTMessage)
exchangeFramed :: StreamIO -> DHTMessage -> IO (Either [Char] DHTMessage)
exchangeFramed StreamIO
stream DHTMessage
request = do
  result <- IO (Either [Char] DHTMessage)
-> IO (Either SomeException (Either [Char] DHTMessage))
forall e a. Exception e => IO a -> IO (Either e a)
try (IO (Either [Char] DHTMessage)
 -> IO (Either SomeException (Either [Char] DHTMessage)))
-> IO (Either [Char] DHTMessage)
-> IO (Either SomeException (Either [Char] DHTMessage))
forall a b. (a -> b) -> a -> b
$ do
    StreamIO -> DHTMessage -> IO ()
writeFramedMessage StreamIO
stream DHTMessage
request
    StreamIO -> Int -> IO (Either [Char] DHTMessage)
readFramedMessage StreamIO
stream Int
maxDHTMessageSize
  pure $ case result of
    Left (SomeException
e :: SomeException) -> [Char] -> Either [Char] DHTMessage
forall a b. a -> Either a b
Left ([Char]
"DHT stream I/O failed: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ SomeException -> [Char]
forall a. Show a => a -> [Char]
show SomeException
e)
    Right Either [Char] DHTMessage
r -> Either [Char] DHTMessage
r

-- | Open a new muxer stream to the peer and negotiate @/ipfs/kad/1.0.0@.
openDHTStream :: Switch -> PeerId -> IO (Either String StreamIO)
openDHTStream :: Switch -> PeerId -> IO (Either [Char] StreamIO)
openDHTStream Switch
sw PeerId
pid = do
  mConn <- STM (Maybe Connection) -> IO (Maybe Connection)
forall a. STM a -> IO a
atomically (STM (Maybe Connection) -> IO (Maybe Connection))
-> STM (Maybe Connection) -> IO (Maybe Connection)
forall a b. (a -> b) -> a -> b
$ TVar (Map PeerId [Connection]) -> PeerId -> STM (Maybe Connection)
lookupConn (Switch -> TVar (Map PeerId [Connection])
swConnPool Switch
sw) PeerId
pid
  case mConn of
    Maybe Connection
Nothing -> Either [Char] StreamIO -> IO (Either [Char] StreamIO)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ([Char] -> Either [Char] StreamIO
forall a b. a -> Either a b
Left [Char]
"no open connection to peer")
    Just Connection
conn -> do
      result <- IO (StreamIO, NegotiationResult)
-> IO (Either SomeException (StreamIO, NegotiationResult))
forall e a. Exception e => IO a -> IO (Either e a)
try (IO (StreamIO, NegotiationResult)
 -> IO (Either SomeException (StreamIO, NegotiationResult)))
-> IO (StreamIO, NegotiationResult)
-> IO (Either SomeException (StreamIO, NegotiationResult))
forall a b. (a -> b) -> a -> b
$ do
        stream <- MuxerSession -> IO StreamIO
muxOpenStream (Connection -> MuxerSession
connSession Connection
conn)
        negotiated <- negotiateInitiator stream [dhtProtocolId]
        pure (stream, negotiated)
      pure $ case result of
        Left (SomeException
e :: SomeException) -> [Char] -> Either [Char] StreamIO
forall a b. a -> Either a b
Left ([Char]
"failed to open DHT stream: " [Char] -> ShowS
forall a. [a] -> [a] -> [a]
++ SomeException -> [Char]
forall a. Show a => a -> [Char]
show SomeException
e)
        Right (StreamIO
stream, Accepted Text
_) -> StreamIO -> Either [Char] StreamIO
forall a b. b -> Either a b
Right StreamIO
stream
        Right (StreamIO
_, NegotiationResult
NoProtocol) -> [Char] -> Either [Char] StreamIO
forall a b. a -> Either a b
Left [Char]
"peer does not support /ipfs/kad/1.0.0"

-- Store operations

-- | Store a record in the local datastore.
storeRecord :: DHTNode -> DHTRecord -> IO ()
storeRecord :: DHTNode -> DHTRecord -> IO ()
storeRecord DHTNode
node DHTRecord
rec = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$
  TVar (Map ByteString DHTRecord)
-> (Map ByteString DHTRecord -> Map ByteString DHTRecord) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' (DHTNode -> TVar (Map ByteString DHTRecord)
dhtRecordStore DHTNode
node) (ByteString
-> DHTRecord
-> Map ByteString DHTRecord
-> Map ByteString DHTRecord
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert (DHTRecord -> ByteString
recKey DHTRecord
rec) DHTRecord
rec)

-- | Look up a record by key.
lookupRecord :: DHTNode -> ByteString -> IO (Maybe DHTRecord)
lookupRecord :: DHTNode -> ByteString -> IO (Maybe DHTRecord)
lookupRecord DHTNode
node ByteString
key = ByteString -> Map ByteString DHTRecord -> Maybe DHTRecord
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup ByteString
key (Map ByteString DHTRecord -> Maybe DHTRecord)
-> IO (Map ByteString DHTRecord) -> IO (Maybe DHTRecord)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (Map ByteString DHTRecord) -> IO (Map ByteString DHTRecord)
forall a. TVar a -> IO a
readTVarIO (DHTNode -> TVar (Map ByteString DHTRecord)
dhtRecordStore DHTNode
node)

-- | Add a provider entry for a content key.
--
-- A provider republishing on schedule replaces its previous entry
-- (deduplicated by peer ID) instead of appending a duplicate, so the
-- entry's timestamp is refreshed and GET_PROVIDERS responses stay
-- bounded.
addProvider :: DHTNode -> ByteString -> ProviderEntry -> IO ()
addProvider :: DHTNode -> ByteString -> ProviderEntry -> IO ()
addProvider DHTNode
node ByteString
key ProviderEntry
entry = STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$
  TVar (Map ByteString [ProviderEntry])
-> (Map ByteString [ProviderEntry]
    -> Map ByteString [ProviderEntry])
-> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' (DHTNode -> TVar (Map ByteString [ProviderEntry])
dhtProviderStore DHTNode
node) ((Map ByteString [ProviderEntry] -> Map ByteString [ProviderEntry])
 -> STM ())
-> (Map ByteString [ProviderEntry]
    -> Map ByteString [ProviderEntry])
-> STM ()
forall a b. (a -> b) -> a -> b
$ \Map ByteString [ProviderEntry]
m ->
    ([ProviderEntry] -> [ProviderEntry] -> [ProviderEntry])
-> ByteString
-> [ProviderEntry]
-> Map ByteString [ProviderEntry]
-> Map ByteString [ProviderEntry]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith [ProviderEntry] -> [ProviderEntry] -> [ProviderEntry]
merge ByteString
key [ProviderEntry
entry] Map ByteString [ProviderEntry]
m
  where
    merge :: [ProviderEntry] -> [ProviderEntry] -> [ProviderEntry]
merge [ProviderEntry]
new [ProviderEntry]
old = [ProviderEntry]
new [ProviderEntry] -> [ProviderEntry] -> [ProviderEntry]
forall a. [a] -> [a] -> [a]
++ (ProviderEntry -> Bool) -> [ProviderEntry] -> [ProviderEntry]
forall a. (a -> Bool) -> [a] -> [a]
filter (\ProviderEntry
e -> ProviderEntry -> PeerId
peProvider ProviderEntry
e PeerId -> PeerId -> Bool
forall a. Eq a => a -> a -> Bool
/= ProviderEntry -> PeerId
peProvider ProviderEntry
entry) [ProviderEntry]
old

-- | Get providers for a content key, pruning entries older than
-- 'providerRecordTTL' (48h expiration interval per specs/kad-dht).
getProviders :: DHTNode -> ByteString -> IO [ProviderEntry]
getProviders :: DHTNode -> ByteString -> IO [ProviderEntry]
getProviders DHTNode
node ByteString
key = do
  now <- IO UTCTime
getCurrentTime
  atomically $ do
    m <- readTVar (dhtProviderStore node)
    let live = (ProviderEntry -> Bool) -> [ProviderEntry] -> [ProviderEntry]
forall a. (a -> Bool) -> [a] -> [a]
filter (\ProviderEntry
e -> UTCTime -> UTCTime -> NominalDiffTime
diffUTCTime UTCTime
now (ProviderEntry -> UTCTime
peTimestamp ProviderEntry
e) NominalDiffTime -> NominalDiffTime -> Bool
forall a. Ord a => a -> a -> Bool
< NominalDiffTime
providerRecordTTL)
                      ([ProviderEntry]
-> ByteString -> Map ByteString [ProviderEntry] -> [ProviderEntry]
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault [] ByteString
key Map ByteString [ProviderEntry]
m)
        m' = if [ProviderEntry] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [ProviderEntry]
live then ByteString
-> Map ByteString [ProviderEntry] -> Map ByteString [ProviderEntry]
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete ByteString
key Map ByteString [ProviderEntry]
m else ByteString
-> [ProviderEntry]
-> Map ByteString [ProviderEntry]
-> Map ByteString [ProviderEntry]
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert ByteString
key [ProviderEntry]
live Map ByteString [ProviderEntry]
m
    writeTVar (dhtProviderStore node) m'
    pure live

-- Helpers

-- | Convert a BucketEntry to a DHTPeer protobuf message.
-- Per specs/kad-dht, Peer records carry the peer's known multiaddrs so
-- the requester can dial them (go-libp2p filters address-less peers).
entryToDHTPeer :: BucketEntry -> DHTPeer
entryToDHTPeer :: BucketEntry -> DHTPeer
entryToDHTPeer BucketEntry
entry = DHTPeer
  { dhtPeerId :: ByteString
dhtPeerId = PeerId -> ByteString
peerIdBytes (BucketEntry -> PeerId
entryPeerId BucketEntry
entry)
  , dhtPeerAddrs :: [ByteString]
dhtPeerAddrs = (Multiaddr -> ByteString) -> [Multiaddr] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map Multiaddr -> ByteString
toBytes (BucketEntry -> [Multiaddr]
entryAddrs BucketEntry
entry)
  , dhtPeerConnType :: ConnectionType
dhtPeerConnType = BucketEntry -> ConnectionType
entryConnType BucketEntry
entry
  }

-- | Convert a DHTPeer from ADD_PROVIDER into a ProviderEntry.
dhtPeerToProvider :: DHTPeer -> UTCTime -> ProviderEntry
dhtPeerToProvider :: DHTPeer -> UTCTime -> ProviderEntry
dhtPeerToProvider DHTPeer
peer UTCTime
now = ProviderEntry
  { peProvider :: PeerId
peProvider  = ByteString -> PeerId
PeerId (DHTPeer -> ByteString
dhtPeerId DHTPeer
peer)
  , peAddrs :: [Multiaddr]
peAddrs     = [ByteString] -> [Multiaddr]
decodePeerAddrs (DHTPeer -> [ByteString]
dhtPeerAddrs DHTPeer
peer)
  , peTimestamp :: UTCTime
peTimestamp = UTCTime
now
  }

-- | Convert a ProviderEntry to a DHTPeer protobuf message.
providerToDHTPeer :: ProviderEntry -> DHTPeer
providerToDHTPeer :: ProviderEntry -> DHTPeer
providerToDHTPeer ProviderEntry
pe = DHTPeer
  { dhtPeerId :: ByteString
dhtPeerId = PeerId -> ByteString
peerIdBytes (ProviderEntry -> PeerId
peProvider ProviderEntry
pe)
  , dhtPeerAddrs :: [ByteString]
dhtPeerAddrs = (Multiaddr -> ByteString) -> [Multiaddr] -> [ByteString]
forall a b. (a -> b) -> [a] -> [b]
map Multiaddr -> ByteString
toBytes (ProviderEntry -> [Multiaddr]
peAddrs ProviderEntry
pe)
  , dhtPeerConnType :: ConnectionType
dhtPeerConnType = ConnectionType
Connected
  }

-- | Decode raw wire multiaddrs from a Peer record, dropping any that fail
-- to parse: a malformed address from a remote peer must not poison the
-- rest of the record.
decodePeerAddrs :: [ByteString] -> [Multiaddr]
decodePeerAddrs :: [ByteString] -> [Multiaddr]
decodePeerAddrs [ByteString]
raw = [Multiaddr
addr | Right Multiaddr
addr <- (ByteString -> Either [Char] Multiaddr)
-> [ByteString] -> [Either [Char] Multiaddr]
forall a b. (a -> b) -> [a] -> [b]
map ByteString -> Either [Char] Multiaddr
fromBytes [ByteString]
raw]