path: root/Mainline.hs

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

import Control.Applicative
import Control.Arrow
import Control.Concurrent.STM
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 qualified Data.Serialize           as S
import Data.Set                           (Set)
import Data.Torrent
import Data.Typeable
import Data.Word
import Network.Address                    (Address, fromSockAddr, setPort,
                                           sockAddrPort, 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)

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 ()


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 "fromBEncode"

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 :: NodeId
    , routing4 :: !( TVar (Maybe RoutingInfo) )
    , routing6 :: !( TVar (Maybe RoutingInfo) )
    }

newClient ::
      SockAddr -> IO (Client String Method TransactionId NodeInfo (Message BValue))
newClient addr = do
    udp <- udpTransport addr
    nid <- NodeId <$> getRandomBytes 20
    self <- atomically $ newTVar
             $ NodeInfo nid (fromMaybe (toEnum 0) $ fromSockAddr addr)
                            (fromMaybe 0 $ sockAddrPort addr)
    routing <- atomically $ Routing nid <$> newTVar Nothing <*> newTVar Nothing
    swarms <- newSwarmsDatabase
    let net = onInbound (updateRouting routing)
              $ layerTransport parsePacket encodePacket
              $ udp
        dispatch tbl = DispatchMethods
            { classifyInbound = classify
            , lookupHandler = handlers
            , tableMethods = tbl
            }

        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)
    map_var <- atomically $ newTVar (0, mempty)
    return Client
            { clientNet           = net
            , clientDispatcher    = dispatch mapT
            , clientErrorReporter = ignoreErrors -- TODO
            , clientPending       = map_var
            , clientAddress       = atomically (readTVar self)
            , clientResponseId    = return
            }

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

updateRouting :: Routing -> NodeInfo -> Message BValue -> IO ()
updateRouting routing naddr _ = do
    -- TODO Update kademlia table.
    -- TODO Update external ip address and update BEP-42 node id.
    return ()

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 = do
        let myid   = tentativeId routing  :: NodeId
            k      = R.defaultK           :: Int
            nobkts = R.defaultBucketCount :: Int
        nfo <- atomically $ readTVar var
        let tbl = maybe (R.nullTable myid nobkts) R.myBuckets nfo
        return $ R.kclosest nodeId k node tbl

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