path: root/Mainline.hs

{-# LANGUAGE CPP                        #-}
{-# LANGUAGE DeriveDataTypeable         #-}
{-# LANGUAGE DeriveFoldable             #-}
{-# LANGUAGE DeriveFunctor              #-}
{-# LANGUAGE DeriveTraversable          #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE LambdaCase                 #-}
{-# LANGUAGE PatternSynonyms            #-}
{-# LANGUAGE StandaloneDeriving         #-}
{-# LANGUAGE TupleSections              #-}
module Mainline where

import Control.Applicative
import Control.Arrow
import Control.Concurrent.STM
import Control.Monad
import Crypto.Random
import Data.BEncode                       as BE
import qualified Data.BEncode.BDict       as BE
         ;import Data.BEncode.BDict       (BKey)
import Data.Bits
import Data.Bits.ByteString
import Data.Bool
import qualified Data.ByteArray           as BA
         ;import Data.ByteArray           (ByteArrayAccess)
import qualified Data.ByteString          as B
         ;import Data.ByteString          (ByteString)
import Data.ByteString.Lazy               (toStrict)
import Data.Coerce
import Data.Data
import Data.Default
import Data.Hashable
import Data.IP
import Data.List
import Data.Maybe
import Data.Monoid
import Data.Ord
import qualified Data.Serialize           as S
import Data.Set                           (Set)
import Data.Torrent
import Data.Typeable
import Data.Word
import Kademlia
import Network.Address                    (Address, fromSockAddr, setPort,
                                           sockAddrPort, testIdBit, toSockAddr)
import Network.BitTorrent.DHT.ContactInfo as Peers
import Network.BitTorrent.DHT.Token       as Token
import qualified Network.DHT.Routing      as R
         ;import Network.DHT.Routing      (Info, Timestamp, getTimestamp)
import Network.QueryResponse
import Network.Socket

newtype NodeId = NodeId ByteString
 deriving (Eq,Ord,Show,ByteArrayAccess, BEncode, Bits, Hashable)

instance FiniteBits NodeId where
    finiteBitSize _ = 160

data NodeInfo = NodeInfo
  { nodeId   :: NodeId
  , nodeIP :: IP
  , nodePort :: PortNumber
  }

instance Hashable NodeInfo where
  hashWithSalt s ni = hashWithSalt s (nodeIP ni , nodePort ni)
  {-# INLINE hashWithSalt #-}



{-

-- | KRPC 'compact list' compatible encoding: contact information for
-- nodes is encoded as a 26-byte string. Also known as "Compact node
-- info" the 20-byte Node ID in network byte order has the compact
-- IP-address/port info concatenated to the end.
  get = NodeInfo <$> (NodeId <$> S.getBytes 20 ) <*> S.get <*> S.get
-}

getNodeInfo4 :: S.Get NodeInfo
getNodeInfo4 = NodeInfo <$> (NodeId <$> S.getBytes 20)
                        <*> (IPv4 <$> S.get)
                        <*> S.get

putNodeInfo4 :: NodeInfo -> S.Put
putNodeInfo4 (NodeInfo (NodeId nid) ip port)
    | IPv4 ip4 <- ip                          = put4 ip4
    | IPv6 ip6 <- ip , Just ip4 <- un4map ip6 = put4 ip4
    | otherwise                               = return ()
 where
    put4 ip4 = S.putByteString nid >> S.put ip4 >> S.put port

getNodeInfo6 :: S.Get NodeInfo
getNodeInfo6 = NodeInfo <$> (NodeId <$> S.getBytes 20)
                        <*> (IPv6 <$> S.get)
                        <*> S.get

putNodeInfo6 :: NodeInfo -> S.Put
putNodeInfo6 (NodeInfo (NodeId nid) (IPv6 ip) port)
               = S.putByteString nid >> S.put ip >> S.put port
putNodeInfo6 _ = return ()


-- TODO: We should use a SocketAddrInet6 address for a dual-stack listen
-- socket.  Therefore, the behavior of this method should depend on the bind
-- address for outbound packets.
nodeAddr :: NodeInfo -> SockAddr
nodeAddr (NodeInfo _ ip port) = setPort port $ toSockAddr ip

nodeInfo :: NodeId -> SockAddr -> Either String NodeInfo
nodeInfo nid saddr
    | Just ip <- fromSockAddr saddr
    , Just port <- sockAddrPort saddr = Right $ NodeInfo nid ip port
    | otherwise                       = Left "Address family not supported."

-- | Types of RPC errors.
data ErrorCode
    -- | Some error doesn't fit in any other category.
  = GenericError

    -- | Occurs when server fail to process procedure call.
  | ServerError

    -- | Malformed packet, invalid arguments or bad token.
  | ProtocolError

    -- | Occurs when client trying to call method server don't know.
  | MethodUnknown
    deriving (Show, Read, Eq, Ord, Bounded, Typeable, Data)

-- | According to the table:
-- <http://bittorrent.org/beps/bep_0005.html#errors>
instance Enum ErrorCode where
  fromEnum GenericError  = 201
  fromEnum ServerError   = 202
  fromEnum ProtocolError = 203
  fromEnum MethodUnknown = 204
  {-# INLINE fromEnum #-}
  toEnum 201 = GenericError
  toEnum 202 = ServerError
  toEnum 203 = ProtocolError
  toEnum 204 = MethodUnknown
  toEnum _   = GenericError
  {-# INLINE toEnum #-}

instance BEncode ErrorCode where
  toBEncode = toBEncode . fromEnum
  {-# INLINE toBEncode #-}
  fromBEncode b = toEnum <$> fromBEncode b
  {-# INLINE fromBEncode #-}

data Error = Error
    { errorCode    :: !ErrorCode  -- ^ The type of error.
    , errorMessage :: !ByteString -- ^ Human-readable text message.
    } deriving ( Show, Eq, Ord, Typeable, Data, Read )

newtype TransactionId = TransactionId ByteString
 deriving (Eq, Ord, Show, BEncode)

newtype Method = Method ByteString
 deriving (Eq, Ord, Show, BEncode)

data Message a = Q { msgOrigin   :: NodeId
                   , msgID       :: TransactionId
                   , qryPayload  :: a
                   , qryMethod   :: Method
                   , qryReadOnly :: Bool }

               | R { msgOrigin      :: NodeId
                   , msgID          :: TransactionId
                   , rspPayload     :: Either Error a
                   , rspReflectedIP :: Maybe SockAddr }

instance BE.BEncode (Message BValue) where
    toBEncode = encodeMessage
    fromBEncode = error "TODO: fromBEncode (Mainline Message)"

encodeMessage (Q origin tid a meth ro)
    = case a of
        BDict args -> encodeQuery tid meth (BDict $ genericArgs origin ro `BE.union` args)
        _          -> encodeQuery tid meth a -- XXX: Not really a valid query.
encodeMessage (R origin tid v ip)
    = case v of
        Right vals -> encodeResponse tid vals (BString . encodeAddr <$> ip)
        Left  err  -> encodeError tid err

encodeAddr :: SockAddr -> ByteString
encodeAddr (SockAddrInet port addr)
    = S.runPut (S.putWord32host addr >> S.put (fromIntegral port :: Word16))
encodeAddr (SockAddrInet6 port _ addr _)
    = S.runPut (S.put addr >> S.put (fromIntegral port :: Word16))
encodeAddr _ = B.empty

genericArgs nodeid ro =
       "id" .=! nodeid
    .: "ro" .=? bool Nothing (Just (1 :: Int)) ro
    .: endDict

encodeError     tid (Error ecode emsg)       = encodeAny tid "e" (ecode,emsg) id

encodeResponse  tid rvals              rip   = encodeAny tid "r" rvals        ("ip" .=? rip   .:)

encodeQuery     tid qmeth              qargs = encodeAny tid "q" qmeth        ("a"  .=! qargs .:)

encodeAny tid key val aux = toDict $
    aux $  key .=! val
        .: "t" .=! tid
        .: "y" .=! key
        .: endDict

parsePacket :: ByteString -> SockAddr -> Either String (Message BValue, NodeInfo)
parsePacket bs addr = do pkt <- BE.decode bs
                         ni <- nodeInfo (msgOrigin pkt) addr
                         return (pkt, ni)

encodePacket :: Message BValue -> NodeInfo -> (ByteString, SockAddr)
encodePacket msg ni = ( toStrict $ BE.encode msg
                      , nodeAddr ni )

classify :: Message BValue -> MessageClass String Method TransactionId
classify (Q { msgID = tid, qryMethod = meth }) = IsQuery meth tid
classify (R { msgID = tid                   }) = IsResponse tid

encodePayload :: BEncode a => TransactionId -> NodeInfo -> NodeInfo -> a -> Message BValue
encodePayload tid self dest b = R (nodeId self) tid (Right $ BE.toBEncode b) (Just $ nodeAddr dest)

errorPayload :: TransactionId -> NodeInfo -> NodeInfo -> Error -> Message a
errorPayload tid self dest e = R (nodeId self) tid (Left e) (Just $ nodeAddr dest)

decodePayload :: BEncode a => Message BValue -> Either String a
decodePayload msg = BE.fromBEncode $ qryPayload msg

type Handler = MethodHandler String TransactionId NodeInfo (Message BValue)

handler :: ( BEncode a
           , BEncode b
           ) =>
           (NodeInfo -> a -> IO b) -> Maybe Handler
handler f = Just $ MethodHandler decodePayload encodePayload f


handlerE :: ( BEncode a
            , BEncode b
            ) =>
            (NodeInfo -> a -> IO (Either Error b)) -> Maybe Handler
handlerE f = Just $ MethodHandler decodePayload enc f
 where
    enc tid self dest (Left e)  = errorPayload  tid self dest e
    enc tid self dest (Right b) = encodePayload tid self dest b

type AnnounceSet = Set (InfoHash, PortNumber)

data SwarmsDatabase = SwarmsDatabase
    { contactInfo   :: !( TVar PeerStore     ) -- ^ Published by other nodes.
    , sessionTokens :: !( TVar SessionTokens ) -- ^ Query session IDs.
    , announceInfo  :: !( TVar AnnounceSet   ) -- ^ To publish by this node.
    }

newSwarmsDatabase :: IO SwarmsDatabase
newSwarmsDatabase = do
    toks <- nullSessionTokens
    atomically
        $ SwarmsDatabase <$> newTVar def
                    <*> newTVar toks
                    <*> newTVar def

type RoutingInfo = Info NodeInfo NodeId

data Routing = Routing
    { tentativeId :: NodeInfo
    , routing4 :: !( TVar (R.BucketList NodeInfo) )
    , committee4 :: TriadCommittee NodeId SockAddr
    , routing6 :: !( TVar (R.BucketList NodeInfo) )
    , committee6 :: TriadCommittee NodeId SockAddr
    }

type MainlineClient = Client String Method TransactionId NodeInfo (Message BValue)

newClient :: SockAddr -> IO MainlineClient
newClient addr = do
    udp <- udpTransport addr
    nid <- NodeId <$> getRandomBytes 20
    let tenative_info = NodeInfo
                { nodeId   = nid
                , nodeIP   = fromMaybe (toEnum 0) $ fromSockAddr addr
                , nodePort = fromMaybe 0 $ sockAddrPort addr
                }
    routing <- atomically $ do
        let nobkts = R.defaultBucketCount :: Int
        tbl4 <- newTVar $ R.nullTable (comparing nodeId) (\s -> hashWithSalt s . nodeId) tenative_info nobkts
        tbl6 <- newTVar $ R.nullTable (comparing nodeId) (\s -> hashWithSalt s . nodeId) tenative_info nobkts
        committee4 <- newTriadCommittee (const $ return ()) -- TODO: update tbl4
        committee6 <- newTriadCommittee (const $ return ()) -- TODO: update tbl6
        return $ Routing tenative_info tbl4 committee4 tbl6 committee6
    swarms <- newSwarmsDatabase
    map_var <- atomically $ newTVar (0, mempty)
    let net = onInbound (updateRouting outgoingClient routing)
              $ layerTransport parsePacket encodePacket
              $ udp

        -- Paranoid: It's safe to define /net/ and /client/ to be mutually
        -- recursive since 'updateRouting' does not invoke 'awaitMessage' which
        -- which was modified by 'onInbound'.  However, I'm going to avoid the
        -- mutual reference just to be safe.
        outgoingClient = client { clientNet = net { awaitMessage = return Nothing } }

        dispatch = DispatchMethods
            { classifyInbound = classify -- :: x -> MessageClass err meth tid
            , lookupHandler   = handlers -- :: meth -> Maybe (MethodHandler err tid addr x)
            , tableMethods    = mapT     -- :: TransactionMethods tbl tid x
            }

        handlers :: Method -> Maybe Handler
        handlers ( Method "ping"          ) = handler pingH
        handlers ( Method "find_node"     ) = handler $ findNodeH routing
        handlers ( Method "get_peers"     ) = handler $ getPeersH routing swarms
        handlers ( Method "announce_peer" ) = handlerE $ announceH swarms
        handlers ( Method meth            ) = Just $ defaultHandler meth

        mapT = transactionMethods mapMethods gen

        gen :: Word16 -> (TransactionId, Word16)
        gen cnt = (TransactionId $ S.encode cnt, cnt+1)

        client = Client
            { clientNet           = net
            , clientDispatcher    = dispatch
            , clientErrorReporter = ignoreErrors -- TODO
            , clientPending       = map_var
            , clientAddress       = \maddr -> atomically $ do
                let var = case flip prefer4or6 Nothing <$> maddr of
                            Just Want_IP6 -> routing6 routing
                            _             -> routing4 routing
                R.selfNode <$> readTVar var
            , clientResponseId    = return
            }

    return client

defaultHandler :: ByteString -> Handler
defaultHandler meth = MethodHandler decodePayload errorPayload returnError
 where
    returnError :: NodeInfo -> BValue -> IO Error
    returnError _ _ = return $ Error MethodUnknown ("Unknown method " <> meth)

mainlineKademlia :: MainlineClient -> TriadCommittee NodeId SockAddr -> TVar (R.BucketList NodeInfo) ->  Kademlia NodeId NodeInfo
mainlineKademlia client committee var
    = Kademlia quietInsertions
               mainlineSpace
               (vanillaIO var $ ping client)
                   { tblTransition = transitionCommittee committee }


mainlineSpace :: R.KademliaSpace NodeId NodeInfo
mainlineSpace = R.KademliaSpace
    { R.kademliaLocation = nodeId
    , R.kademliaTestBit  = testIdBit
    , R.kademliaXor      = xor
    }

transitionCommittee :: TriadCommittee NodeId SockAddr -> RoutingTransition NodeInfo -> STM (IO ())
transitionCommittee committee (RoutingTransition ni Stranger) = do
    delVote committee (nodeId ni)
    return $ return ()

updateRouting :: MainlineClient -> Routing -> NodeInfo -> Message BValue -> IO ()
updateRouting client routing naddr msg = do
    case prefer4or6 naddr Nothing of
        Want_IP4 -> go (routing4 routing) (committee4 routing)
        Want_IP6 -> go (routing6 routing) (committee6 routing)
 where
    go tbl committee = do
        case msg of
            R { rspReflectedIP = Just sockaddr }
                -> atomically $ addVote committee (nodeId naddr) sockaddr
            _   -> return ()
        insertNode (mainlineKademlia client committee tbl) naddr

data Ping = Ping deriving Show

-- Pong is the same as Ping.
type Pong = Ping
pattern Pong = Ping

instance BEncode Ping where
  toBEncode Ping = toDict endDict
  fromBEncode _  = pure Ping

data WantIP = Want_IP4 | Want_IP6 | Want_Both
 deriving (Eq, Enum, Ord, Show)

wantList :: WantIP -> [ByteString]
wantList Want_IP4  = ["ip4"]
wantList Want_IP6  = ["ip6"]
wantList Want_Both = ["ip4","ip6"]

instance BEncode WantIP where
    toBEncode w = toBEncode $ wantList w
    fromBEncode bval = do
        wants <- fromBEncode bval
        let _ = wants :: [ByteString]
        case (elem "ip4" wants, elem "ip6" wants) of
             (True,True)  -> Right Want_Both
             (True,False) -> Right Want_IP4
             (False,True) -> Right Want_IP6
             _            -> Left "Unrecognized IP type."

data FindNode = FindNode NodeId (Maybe WantIP)

instance BEncode FindNode where
  toBEncode (FindNode nid iptyp) = toDict   $ target_key .=! nid
                                              .: want_key .=? iptyp
                                              .: endDict
  fromBEncode                    = fromDict $ FindNode  <$>! target_key
                                                        <*>? want_key

data NodeFound = NodeFound
    { nodes4 :: [NodeInfo]
    , nodes6 :: [NodeInfo]
    }

instance BEncode NodeFound where
  toBEncode (NodeFound ns ns6) = toDict $
       nodes_key .=?
        (if Prelude.null ns then Nothing
                            else Just (S.runPut (mapM_ putNodeInfo4 ns)))
    .: nodes6_key .=?
        (if Prelude.null ns6 then Nothing
                             else Just (S.runPut (mapM_ putNodeInfo6 ns6)))
    .: endDict

  fromBEncode bval = NodeFound <$> ns4 <*> ns6
   where
    ns4 = fromDict (binary getNodeInfo4 nodes_key) bval
    ns6 = fromDict (binary getNodeInfo6 nodes6_key) bval

binary :: S.Get a -> BKey -> BE.Get [a]
binary get k = field (req k) >>= either (fail . format) return .
    S.runGet (many get)
  where
    format str = "fail to deserialize " ++ show k ++ " field: " ++ str

pingH :: NodeInfo -> Ping -> IO Pong
pingH _ Ping = return Pong

-- | True if the argument is an IPv4-mapped address with prefix ::FFFF:0:0/96
-- as defined in RFC 4291.
is4mapped :: IPv6 -> Bool
is4mapped ip
    | [0,0,0,0,0,0xffff,_,_] <- fromIPv6 ip
                = True
    | otherwise = False

un4map :: IPv6 -> Maybe IPv4
un4map ip
    | [0,0,0,0,0,0xffff,x,y] <- fromIPv6 ip
                = Just $ toIPv4
                       $ map (.&. 0xFF)
                             [x `shiftR` 8, x, y `shiftR` 8, y ]
    | otherwise = Nothing

prefer4or6 :: NodeInfo -> Maybe WantIP -> WantIP
prefer4or6 addr iptyp = fromMaybe (ipFamily $ nodeIP addr) iptyp

ipFamily :: IP -> WantIP
ipFamily ip = case ip of
    IPv4 _               -> Want_IP4
    IPv6 a | is4mapped a -> Want_IP4
           | otherwise   -> Want_IP6

findNodeH :: Routing -> NodeInfo -> FindNode -> IO NodeFound
findNodeH routing addr (FindNode node iptyp) = do
    let preferred = prefer4or6 addr iptyp
    ks  <- bool (return []) (go $ routing4 routing) (preferred /= Want_IP6)
    ks6 <- bool (return []) (go $ routing6 routing) (preferred /= Want_IP4)
    return $ NodeFound ks ks6
 where
    go var = R.kclosest nodeId k node <$> atomically (readTVar var)
    k = R.defaultK


data GetPeers = GetPeers InfoHash (Maybe WantIP)

instance BEncode GetPeers where
  toBEncode (GetPeers ih iptyp)
              = toDict   $ info_hash_key .=! ih
                        .: want_key .=? iptyp
                        .: endDict
  fromBEncode = fromDict $ GetPeers <$>! info_hash_key <*>? want_key


data GotPeers = GotPeers
  { -- | If the queried node has no peers for the infohash, returned
    -- the K nodes in the queried nodes routing table closest to the
    -- infohash supplied in the query.
    peers        :: [PeerAddr]

  , nodes        :: NodeFound

    -- | The token value is a required argument for a future
    -- announce_peer query.
  , grantedToken :: Token
  } -- deriving (Show, Eq, Typeable)

nodeIsIPv6 :: NodeInfo -> Bool
nodeIsIPv6 (NodeInfo _ (IPv6 _) _) = True
nodeIsIPv6 _                       = False

instance BEncode GotPeers where
  toBEncode GotPeers { nodes = NodeFound ns4 ns6, ..} = toDict $
       nodes_key  .=? (if null ns4 then Nothing
                                   else Just $ S.runPut (mapM_ putNodeInfo4 ns4))
    .: nodes6_key .=? (if null ns6 then Nothing
                                   else Just $ S.runPut (mapM_ putNodeInfo4 ns6))
    .: token_key .=! grantedToken
    .: peers_key .=! map S.encode peers
    .: endDict

  fromBEncode = fromDict $ do
    ns4 <- fromMaybe [] <$> optional (binary getNodeInfo4 nodes_key)        -- "nodes"
    ns6 <- fromMaybe [] <$> optional (binary getNodeInfo6 nodes6_key)       -- "nodes6"
    -- TODO: BEP 42...
    --
    -- Once enforced, responses to get_peers requests whose node ID does not
    -- match its external IP should be considered to not contain a token and
    -- thus not be eligible as storage target.  Implementations should take
    -- care that they find the closest set of nodes which return a token and
    -- whose IDs matches their IPs before sending a store request to those
    -- nodes.
    --
    -- Sounds like something to take care of at peer-search time, so I'll
    -- ignore it for now.
    tok <- field    (req    token_key)                                      -- "token"
    ps <- fromMaybe [] <$> optional (field (req peers_key) >>= decodePeers) -- "values"
    pure $ GotPeers ps (NodeFound ns4 ns6) tok
     where
       decodePeers = either fail pure . mapM S.decode

getPeersH :: Routing -> SwarmsDatabase -> NodeInfo -> GetPeers -> IO GotPeers
getPeersH routing (SwarmsDatabase peers toks _) naddr (GetPeers ih iptyp) = do
    ps <- do
        tm <- getTimestamp
        atomically $ do
          (ps,store') <- Peers.freshPeers ih tm <$> readTVar peers
          writeTVar peers store'
          return ps
    -- Filter peer results to only a single address family, IPv4 or IPv6, as
    -- per BEP 32.
    let notboth  = iptyp >>= \case Want_Both -> Nothing
                                   specific  -> Just specific
        selected = prefer4or6 naddr notboth
        ps'      = filter ( (== selected) . ipFamily . peerHost ) ps
    tok <- grantToken toks naddr
    ns <- findNodeH routing naddr (FindNode (coerce ih) iptyp)
    return $ GotPeers ps' ns tok

-- | Announce that the peer, controlling the querying node, is
-- downloading a torrent on a port.
data Announce = Announce
  { -- | If set, the 'port' field should be ignored and the source
    -- port of the UDP packet should be used as the peer's port
    -- instead. This is useful for peers behind a NAT that may not
    -- know their external port, and supporting uTP, they accept
    -- incoming connections on the same port as the DHT port.
    impliedPort :: Bool

    -- | infohash of the torrent;
  , topic    :: InfoHash

    -- | some clients announce the friendly name of the torrent here.
  , announcedName :: Maybe ByteString

    -- | the port /this/ peer is listening;
  , port     :: PortNumber

    -- TODO: optional boolean "seed" key

    -- | received in response to a previous get_peers query.
  , sessionToken :: Token

  } deriving (Show, Eq, Typeable)

peer_ip_key   = "ip"
peer_id_key   = "peer id"
peer_port_key = "port"
msg_type_key   = "msg_type"
piece_key      = "piece"
total_size_key = "total_size"
node_id_key :: BKey
node_id_key = "id"
read_only_key :: BKey
read_only_key = "ro"
want_key :: BKey
want_key = "want"
target_key :: BKey
target_key = "target"
nodes_key :: BKey
nodes_key = "nodes"
nodes6_key :: BKey
nodes6_key = "nodes6"
info_hash_key :: BKey
info_hash_key = "info_hash"
peers_key :: BKey
peers_key = "values"
token_key :: BKey
token_key = "token"
name_key :: BKey
name_key = "name"
port_key :: BKey
port_key = "port"
implied_port_key :: BKey
implied_port_key = "implied_port"

instance BEncode Announce where
  toBEncode Announce {..} = toDict $
       implied_port_key .=? flagField impliedPort
    .: info_hash_key    .=! topic
    .: name_key         .=? announcedName
    .: port_key         .=! port
    .: token_key        .=! sessionToken
    .: endDict
    where
      flagField flag = if flag then Just (1 :: Int) else Nothing

  fromBEncode = fromDict $ do
    Announce <$> (boolField <$> optional (field (req implied_port_key)))
             <*>! info_hash_key
             <*>? name_key
             <*>! port_key
             <*>! token_key
    where
      boolField = maybe False (/= (0 :: Int))



-- | The queried node must verify that the token was previously sent
-- to the same IP address as the querying node. Then the queried node
-- should store the IP address of the querying node and the supplied
-- port number under the infohash in its store of peer contact
-- information.
data Announced = Announced
  deriving (Show, Eq, Typeable)

instance BEncode Announced where
  toBEncode   _ = toBEncode Ping
  fromBEncode _ = pure  Announced

announceH :: SwarmsDatabase -> NodeInfo -> Announce -> IO (Either Error Announced)
announceH (SwarmsDatabase peers toks _) naddr announcement = do
  checkToken toks naddr (sessionToken announcement)
    >>= bool (Left  <$> return (Error ProtocolError "invalid parameter: token"))
             (Right <$> go)
 where
    go = atomically $ do
        modifyTVar' peers
            $ insertPeer (topic announcement) (announcedName announcement)
            $ PeerAddr
                { peerId = Nothing
                -- Avoid storing IPv4-mapped addresses.
                , peerHost = case nodeIP naddr of
                    IPv6 ip6 | Just ip4 <- un4map ip6 -> IPv4 ip4
                    a                                 -> a
                , peerPort = if impliedPort announcement
                                then nodePort naddr
                                else port announcement
                }
        return Announced

ping :: MainlineClient -> NodeInfo -> IO Bool
ping client addr = fromMaybe False <$> sendQuery client serializer Ping addr
 where
    serializer = MethodSerializer
        { methodTimeout  = 5
        , method         = Method "ping"
        , wrapQuery      = encodePayload
        , unwrapResponse = const True
        }

data TriadSlot = SlotA | SlotB | SlotC
 deriving (Eq,Ord,Enum,Show,Read)

data TriadCommittee voter a = TriadCommittee
    { triadDecider :: TVar TriadSlot
    , triadA :: TVar (Maybe (voter,a))
    , triadB :: TVar (Maybe (voter,a))
    , triadC :: TVar (Maybe (voter,a))
    , triadNewDecision :: a -> STM ()
    }

triadSlot :: TriadSlot -> TriadCommittee voter a -> TVar (Maybe (voter,a))
triadSlot SlotA = triadA
triadSlot SlotB = triadB
triadSlot SlotC = triadC

triadDecision :: a -> TriadCommittee voter a -> STM a
triadDecision fallback triad = do
    slot <- readTVar (triadDecider triad)
    maybe fallback snd <$> readTVar (triadSlot slot triad)


newTriadCommittee :: (a -> STM ()) -> STM (TriadCommittee voter a)
newTriadCommittee onChange =
    TriadCommittee <$> newTVar SlotA
                   <*> newTVar Nothing
                   <*> newTVar Nothing
                   <*> newTVar Nothing
                   <*> pure onChange

triadCountVotes :: Eq a => TriadCommittee voter a -> STM ()
triadCountVotes triad = do
    prior <- do
        slot <- readTVar (triadDecider triad)
        fmap snd <$> readTVar (triadSlot slot triad)
    a <- fmap ((SlotA,) . snd) <$> readTVar (triadA triad)
    b <- fmap ((SlotB,) . snd) <$> readTVar (triadB triad)
    c <- fmap ((SlotC,) . snd) <$> readTVar (triadC triad)
    let (slot,vote) = case catMaybes [a,b,c] of
            [  (x,xvote)
             , (y,yvote)
             , (z,zvote) ]  -> if xvote == yvote then (x,Just xvote)
                                                 else (z,Just zvote)
            []              -> (SlotA,Nothing)
            ((slot,vote):_) -> (slot, Just vote)
    writeTVar (triadDecider triad) slot
    case vote of
        Just v | vote /= prior -> triadNewDecision triad v
        _                      -> return ()


addVote :: (Eq voter, Eq a) => TriadCommittee voter a -> voter -> a -> STM ()
addVote triad voter vote = do
    a <- (SlotA,) . fmap fst <$> readTVar (triadA triad)
    b <- (SlotB,) . fmap fst <$> readTVar (triadB triad)
    c <- (SlotC,) . fmap fst <$> readTVar (triadC triad)
    let avail (_,Nothing) = True
        avail (_,Just x ) = (x == voter)
        slots = filter avail [a,b,c]
    forM_ (take 1 slots) $ \(slot,_) -> do
        writeTVar (triadSlot slot triad)
                  (Just (voter,vote))
        triadCountVotes triad

delVote :: (Eq voter, Eq a) => TriadCommittee voter a -> voter -> STM ()
delVote triad voter = do
    a <- (SlotA,) . fmap fst <$> readTVar (triadA triad)
    b <- (SlotB,) . fmap fst <$> readTVar (triadB triad)
    c <- (SlotC,) . fmap fst <$> readTVar (triadC triad)
    let match (_,Just x ) = (x == voter)
        slots = filter match [a,b,c]
    forM_ (take 1 slots) $ \(slot,_) -> do
        writeTVar (triadSlot slot triad) Nothing
        triadCountVotes triad