-- | 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 (..)
  , ProviderEntry (..)
  , Validator (..)
    -- * Validators
  , defaultValidator
  , namespacedValidator
  , pkValidator
    -- * Construction
  , newDHTNode
    -- * 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.STM
import Control.Exception (SomeException, catch, try)
import Data.ByteString (ByteString)
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 (..)
  , 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

-- | 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 StreamIO)
dhtStreams       :: !(TVar (Map PeerId StreamIO))
    -- ^ Cached outbound @/ipfs/kad/1.0.0@ streams, 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.
  }

-- | Create a new DHT node with the outbound sender wired to the Switch.
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
  pure DHTNode
    { dhtSwitch        = sw
    , dhtRoutingTable  = rt
    , dhtRecordStore   = records
    , dhtProviderStore = providers
    , dhtLocalKey      = peerIdToKey localPid
    , dhtLocalPeerId   = localPid
    , dhtMode          = mode
    , dhtValidator     = defaultValidator
    , dhtStreams       = streams
    , dhtSendRequest   = sendRequestViaSwitch sw streams
    }

-- | 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 pipelines all requests to a peer over one stream); 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.
sendRequestViaSwitch
  :: Switch
  -> TVar (Map PeerId StreamIO)
  -> PeerId
  -> DHTMessage
  -> IO (Either String DHTMessage)
sendRequestViaSwitch :: Switch
-> TVar (Map PeerId StreamIO)
-> PeerId
-> DHTMessage
-> IO (Either [Char] DHTMessage)
sendRequestViaSwitch Switch
sw TVar (Map PeerId StreamIO)
streamsVar PeerId
pid DHTMessage
request = do
  mCached <- PeerId -> Map PeerId StreamIO -> Maybe StreamIO
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup PeerId
pid (Map PeerId StreamIO -> Maybe StreamIO)
-> IO (Map PeerId StreamIO) -> IO (Maybe StreamIO)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (Map PeerId StreamIO) -> IO (Map PeerId StreamIO)
forall a. TVar a -> IO a
readTVarIO TVar (Map PeerId StreamIO)
streamsVar
  case mCached of
    Maybe StreamIO
Nothing -> IO (Either [Char] DHTMessage)
openAndExchange
    Just StreamIO
stream -> do
      result <- StreamIO -> DHTMessage -> IO (Either [Char] DHTMessage)
exchangeFramed StreamIO
stream DHTMessage
request
      case result of
        Right DHTMessage
resp -> Either [Char] DHTMessage -> IO (Either [Char] DHTMessage)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (DHTMessage -> Either [Char] DHTMessage
forall a b. b -> Either a b
Right DHTMessage
resp)
        Left [Char]
_ -> do
          -- Cached stream is dead: evict it and retry on a fresh one.
          STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar (Map PeerId StreamIO)
-> (Map PeerId StreamIO -> Map PeerId StreamIO) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (Map PeerId StreamIO)
streamsVar (PeerId -> Map PeerId StreamIO -> Map PeerId StreamIO
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete PeerId
pid)
          IO (Either [Char] DHTMessage)
openAndExchange
  where
    openAndExchange :: IO (Either [Char] DHTMessage)
openAndExchange = do
      opened <- Switch -> PeerId -> IO (Either [Char] StreamIO)
openDHTStream Switch
sw PeerId
pid
      case opened of
        Left [Char]
err -> 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 [Char]
err)
        Right StreamIO
stream -> do
          STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar (Map PeerId StreamIO)
-> (Map PeerId StreamIO -> Map PeerId StreamIO) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (Map PeerId StreamIO)
streamsVar (PeerId -> StreamIO -> Map PeerId StreamIO -> Map PeerId StreamIO
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert PeerId
pid StreamIO
stream)
          result <- StreamIO -> DHTMessage -> IO (Either [Char] DHTMessage)
exchangeFramed StreamIO
stream DHTMessage
request
          case result of
            Left [Char]
err -> do
              STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar (Map PeerId StreamIO)
-> (Map PeerId StreamIO -> Map PeerId StreamIO) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
modifyTVar' TVar (Map PeerId StreamIO)
streamsVar (PeerId -> Map PeerId StreamIO -> Map PeerId StreamIO
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete PeerId
pid)
              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 [Char]
err)
            Either [Char] DHTMessage
ok -> Either [Char] DHTMessage -> IO (Either [Char] DHTMessage)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Either [Char] DHTMessage
ok

-- | 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]