path: root/OnionRouter.hs

{-# LANGUAGE NondecreasingIndentation #-}
module OnionRouter where

import Control.Concurrent.Lifted.Instrument
import Crypto.Tox
import Network.Kademlia
import Network.Kademlia.Routing
import Network.QueryResponse
import Network.Tox.NodeId
import Network.Tox.Onion.Transport

import Control.Arrow
import Control.Concurrent.STM
import Control.Monad
import Crypto.PubKey.Curve25519
import Crypto.Random
import Data.Bits
import qualified Data.ByteString     as B
import qualified Data.HashMap.Strict as HashMap
         ;import Data.HashMap.Strict (HashMap)
import qualified Data.IntMap         as IntMap
         ;import Data.IntMap         (IntMap)
import Data.Maybe
import qualified Data.Serialize      as S
import Data.Typeable
import Data.Word
import qualified Data.Word64Map      as W64
         ;import Data.Word64Map      (Word64Map, fitsInInt)
import Network.Socket
import System.Endian
import System.IO

newtype RouteId = RouteId Int
 deriving Show

-- Toxcore saves a maximum of 12 paths: 6 paths are reserved for announcing
-- ourselves and 6 others are used to search for friends.
--
-- Note: This is pointless because a man-in-the-middle attack currently makes
-- it trivial to glean friend relationships: the storing node can swap the
-- published to-route key with his own giving him access to one layer of
-- encryption and thus the real public key of the sender. TODO:
-- Counter-measures.
--
-- Unlike toxcore, we don't currently reserve paths for only-searching or
-- only-announcing.  Instead, we maintain 12 multi-purpose routes.
data OnionRouter = OnionRouter
    { pendingQueries  :: TVar (Word64Map NodeId)
    , routeMap        :: TVar (IntMap RouteRecord)
    , trampolineNodes :: TVar (IntMap NodeInfo)
    , trampolineIds   :: TVar (HashMap NodeId Int)
    , trampolineCount :: TVar Int
    , onionDRG        :: TVar ChaChaDRG
    , routeThread     :: ThreadId
    , pendingRoutes   :: IntMap (TVar Bool)
    , routeLog        :: TChan String
    }

data RouteRecord = RouteRecord
    { storedRoute :: OnionRoute
    , responseCount :: !Int
    , timeoutCount :: !Int
    }

-- Onion paths have different timeouts depending on whether the path is
-- confirmed or unconfirmed. Unconfirmed paths (paths that core has never
-- received any responses from) have a timeout of 4 seconds with 2 tries before
-- they are deemed non working. This is because, due to network conditions,
-- there may be a large number of newly created paths that do not work and so
-- trying them a lot would make finding a working path take much longer. The
-- timeout for a confirmed path (from which a response was received) is 10
-- seconds with 4 tries without a response. A confirmed path has a maximum
-- lifetime of 1200 seconds to make possible deanonimization attacks more
-- difficult.
timeoutForRoute :: RouteRecord -> Int
timeoutForRoute RouteRecord{ responseCount = 0 } = 4000000
timeoutForRoute RouteRecord{ responseCount = _ } = 10000000

freshRoute :: OnionRoute -> RouteRecord
freshRoute r = RouteRecord
    { storedRoute = r
    , responseCount = 0
    , timeoutCount = 0
    }

gotResponse :: RouteRecord -> RouteRecord
gotResponse rr = rr
    { responseCount = succ $ responseCount rr
    , timeoutCount = 0
    }

gotTimeout :: RouteRecord -> RouteRecord
gotTimeout rr = rr
    { timeoutCount = succ $ timeoutCount rr
    }

data RouteEvent = BuildRoute RouteId

newOnionRouter :: IO OnionRouter
newOnionRouter = do
    drg0 <- drgNew
    or <- atomically $ do
        chan <- newTChan
        drg <- newTVar drg0
        forM_ [0..11] $ \n -> writeTChan chan $ BuildRoute (RouteId n)
        pq <- newTVar W64.empty
        rm <- newTVar IntMap.empty
        tn <- newTVar IntMap.empty
        ti <- newTVar HashMap.empty
        tc <- newTVar 0
        vs <- sequence $ replicate 12 (newTVar True)
        rlog <- newTChan
        return OnionRouter
            { pendingRoutes = IntMap.fromList $ zip [0..11] vs
            , onionDRG = drg
            , pendingQueries = pq
            , routeMap = rm
            , trampolineNodes = tn
            , trampolineIds = ti
            , trampolineCount = tc
            , routeLog = rlog
            , routeThread = error "Failed to invoke forkRouteBuilder"
            }
    return or

forkRouteBuilder :: OnionRouter -> (NodeId -> NodeInfo -> IO [NodeInfo]) -> IO OnionRouter
forkRouteBuilder or getnodes = do
    tid <- forkIO $ do
        me <- myThreadId
        labelThread me "OnionRouter"
        forever $ do
            let checkRebuild rid want_build stm = flip orElse stm $ do
                    readTVar want_build >>= check
                    writeTVar want_build False
                    return $ BuildRoute $ RouteId rid
            io <- atomically $
                   (readTChan (routeLog or) >>= return . hPutStrLn stderr)
                `orElse`
                   (IntMap.foldrWithKey checkRebuild retry (pendingRoutes or)
                        >>= return . handleEvent getnodes or { routeThread = me })
            io
    return or { routeThread = tid }

generateNodeId :: MonadRandom m => m NodeId
generateNodeId = either (error "unable to make random nodeid")
                        id
                 . S.decode <$> getRandomBytes 32

distinct3by :: Eq t => (a -> t) -> a -> a -> a -> Bool
distinct3by f a b c = f a /= f b && f b /= f c && f c /= f a

-- The two integer functions below take an [inclusive,inclusive] range.
randomR :: (DRG g, Integral a) => (a, a) -> g -> (a, g)
randomR (l,h) = randomIvalInteger (toInteger l, toInteger h)

next :: DRG g => g -> (Int,g)
next g = withDRG g $ do bs <- getRandomBytes $ if fitsInInt (Proxy :: Proxy Word64)
                                                then 8
                                                else 4
                        either (return . error) return $ S.decode bs

randomIvalInteger :: (DRG g, Num a) => (Integer, Integer) -> g -> (a, g)
randomIvalInteger (l,h) rng
 | l > h     = randomIvalInteger (h,l) rng
 | otherwise = case (f 1 0 rng) of (v, rng') -> (fromInteger (l + v `mod` k), rng')
     where
       (genlo, genhi) = (minBound :: Int, maxBound :: Int) -- genRange :: RandomGen g => g -> (Int, Int)
       b = fromIntegral genhi - fromIntegral genlo + 1

       -- Probabilities of the most likely and least likely result
       -- will differ at most by a factor of (1 +- 1/q).  Assuming the RandomGen
       -- is uniform, of course

       -- On average, log q / log b more random values will be generated
       -- than the minimum
       q = 1000
       k = h - l + 1
       magtgt = k * q

       -- generate random values until we exceed the target magnitude
       f mag v g | mag >= magtgt = (v, g)
                 | otherwise = v' `seq`f (mag*b) v' g' where
                        (x,g') = next g -- next :: RandomGen g => g -> (Int, g)
                        v' = (v * b + (fromIntegral x - fromIntegral genlo))

selectTrampolines :: OnionRouter -> STM [NodeInfo]
selectTrampolines or = do
    cnt <- readTVar (trampolineCount or)
    drg0 <- readTVar (onionDRG or)
    ts <- readTVar (trampolineNodes or)
    let (a, drg1) = randomR (0,cnt - 1) drg0
        (b0, drg2) = randomR (0,cnt - 2) drg1
        (c0, drg ) = randomR (0,cnt - 3) drg2
        b | b0 < a    = b0
          | otherwise = b0 + 1
        c1 | c0 < a     = c0
           | otherwise  = c0 + 1
        c | c1 < b    = c1
          | otherwise = c1 + 1
        ns = mapMaybe (\n -> IntMap.lookup n ts) [a,b,c]
    ns' <- case ns of
            [an,bn,cn] -> do check $ distinct3by nodeClass an bn cn
                             return ns
            _          -> retry
    writeTVar (onionDRG or) drg
    return ns'

handleEvent :: (NodeId -> NodeInfo -> IO [NodeInfo]) -> OnionRouter -> RouteEvent -> IO ()
handleEvent getnodes or e@(BuildRoute (RouteId rid)) = do
    hPutStrLn stderr $ "ONION Rebuilding RouteId " ++ show rid
    mb <- join . atomically $ do
        ts <- selectTrampolines or
        drg <- readTVar (onionDRG or)
        [av,bv,cv] <- sequence $ replicate 3 (newTVar Nothing)
        let (getr, drg') = withDRG drg $ do
                asec <- generateSecretKey -- Three aliases
                bsec <- generateSecretKey
                csec <- generateSecretKey
                aq <- generateNodeId -- Three queries
                bq <- generateNodeId
                cq <- generateNodeId
                sel <- B.head <$> getRandomBytes 1 -- Three two-bit result selectors (6 bits)
                let asel = sel .&. 0x3
                    bsel = shiftR sel 2 .&. 0x3
                    csel = shiftR sel 4 .&. 0x3
                    sendq s q ni = listToMaybe . drop (fromIntegral s) <$> getnodes q ni
                    sendqs = do
                        forkIO $ sendq asel aq (ts !! 0) >>= atomically . writeTVar av . Just
                        forkIO $ sendq bsel bq (ts !! 1) >>= atomically . writeTVar bv . Just
                        forkIO $ sendq csel cq (ts !! 2) >>= atomically . writeTVar cv . Just
                        atomically $ do -- Wait for all 3 results.
                            rs <- catMaybes <$> sequence [readTVar av,readTVar bv,readTVar cv]
                            case rs of [_,_,_] -> do
                                                    return $ catMaybes rs
                                                    -- self <- IntMap.lookup (-1) <$> readTVar (trampolineNodes or)
                                                    -- return $ maybe (catMaybes rs) (\x -> [x,x,x]) self
                                       _       -> retry
                return $ do
                    nodes <- sendqs
                    hPutStr stderr $ unlines
                        [ "ONION trampolines:   " ++ show ts
                        , "ONION query results: " ++ show nodes ]
                    case nodes of
                        [a,b,c] | distinct3by nodeClass a b c -> do
                            atomically $ writeTChan (routeLog or) $ unwords [ "ONION using route:", show $ nodeAddr a, show $ nodeAddr b, show $ nodeAddr c]
                            return $ Just OnionRoute
                                { routeAliasA = asec
                                , routeAliasB = bsec
                                , routeAliasC = csec
                                , routeNodeA = a
                                , routeNodeB = b
                                , routeNodeC = c
                                }
                        [a,b,c] -> do
                            atomically $ writeTChan (routeLog or) $ unwords [ "ONION Discarding insecure route:", show $ nodeAddr a, show $ nodeAddr b, show $ nodeAddr c]
                            return Nothing
                        _ -> return Nothing
        writeTVar (onionDRG or) drg'
        return $ getr
    atomically $ maybe (writeTVar (pendingRoutes or IntMap.! rid) True)
                       (\r -> modifyTVar' (routeMap or)
                                          (IntMap.insert rid $ freshRoute r))
                       mb
    case mb of
        Just _  -> hPutStrLn stderr $ "ONION Finished RouteId " ++ show rid
        Nothing -> hPutStrLn stderr $ "ONION Failed RouteId " ++ show rid

routeId :: Nonce8 -> RouteId
routeId (Nonce8 w8) = RouteId $ mod (fromIntegral w8) 12

lookupSender :: OnionRouter -> SockAddr -> Nonce8 -> IO (Maybe (NodeInfo, RouteId))
lookupSender or saddr (Nonce8 w8) = do
    result <- atomically $ do
        ks <- readTVar (pendingQueries or)
        let r = W64.lookup w8 ks
        writeTChan (routeLog or) $ "ONION lookupSender " ++ unwords [show w8, "->", show r]
        return r
    return $ do
        nid <- result
        ni <- either (const Nothing) Just $ nodeInfo nid saddr
        Just (ni, routeId (Nonce8 w8))

lookupRoute :: OnionRouter -> NodeInfo -> RouteId -> IO (Maybe OnionRoute)
lookupRoute or ni (RouteId rid) = do
    mb <- atomically $ IntMap.lookup rid <$> readTVar (routeMap or)
    return $ storedRoute <$> mb

lookupTimeout :: OnionRouter -> Nonce8 -> OnionDestination r -> STM (OnionDestination RouteId, Int)
lookupTimeout or n8 (OnionDestination ni Nothing) = do
    let RouteId rid = routeId n8
    mrr <- IntMap.lookup rid <$> readTVar (routeMap or)
    readTVar (routeMap or) >>= \rm -> writeTChan (routeLog or) $ "ONION lookupTimeout " ++ unwords [show rid,show (IntMap.keys rm)]
    case mrr of
        Just rr -> return ( OnionDestination ni (Just $ routeId n8), timeoutForRoute rr)
        Nothing -> return ( OnionDestination ni Nothing            , 0 )

hookQueries :: OnionRouter -> (tid -> Nonce8)
                           -> TransactionMethods d tid (OnionDestination r) x
                           -> TransactionMethods d tid (OnionDestination r) x
hookQueries or t8 tmethods = TransactionMethods
    { dispatchRegister = \mvar od d -> do --  :: MVar x -> d -> STM (tid, d)
        (tid,d') <- dispatchRegister tmethods mvar od d
        let Nonce8 w8 = t8 tid
            ni = onionNodeInfo od
        modifyTVar' (pendingQueries or) (W64.insert w8 (nodeId ni))
        writeTChan (routeLog or) $ "ONION query add " ++ unwords [ show w8, ":=", show ni ]
        return (tid,d')
    , dispatchResponse = \tid x d -> do -- :: tid -> x -> d -> STM (d, IO ())
        let (Nonce8 w8, RouteId rid) = (id &&& routeId) $ t8 tid
        modifyTVar' (pendingQueries or) (W64.delete w8)
        modifyTVar' (routeMap or) (IntMap.adjust gotResponse rid)
        writeTChan (routeLog or) $ "ONION query del " ++ show w8
        dispatchResponse tmethods tid x d
    , dispatchCancel = \tid d -> do -- :: tid -> d -> STM d
        let (Nonce8 w8, RouteId rid) = (id &&& routeId) $ t8 tid
        modifyTVar' (pendingQueries or) (W64.delete w8)
        writeTChan (routeLog or) $ "ONION query can " ++ show w8
        modifyTVar' (routeMap or) (IntMap.adjust gotTimeout rid)
        let expireRoute = writeTVar (pendingRoutes or IntMap.! rid) True
        rr <- IntMap.lookup rid <$> readTVar (routeMap or)
        case rr of
            Just RouteRecord{ responseCount=0
                            , timeoutCount = c } | c >= 2 -> expireRoute
            Just RouteRecord{ timeoutCount = c } | c >= 4 -> expireRoute
            _ -> return ()
        dispatchCancel tmethods tid d
    }


addtramp :: NodeInfo -> Maybe (HashMap NodeId NodeInfo) -> Maybe (HashMap NodeId NodeInfo)
addtramp ni Nothing  = Just $ HashMap.singleton (nodeId ni) ni
addtramp ni (Just m) = Just $ HashMap.insert (nodeId ni) ni m

deltramp :: NodeInfo -> Maybe (HashMap NodeId v) -> Maybe (HashMap NodeId v)
deltramp ni Nothing  = Nothing
deltramp ni (Just m) = case HashMap.delete (nodeId ni) m of
                        m' | HashMap.null m' -> Nothing
                        m'                   -> Just m'

hookBucketList :: KademliaSpace NodeId NodeInfo -> TVar (BucketList NodeInfo) ->  OnionRouter -> RoutingTransition NodeInfo -> STM ()
hookBucketList kademlia bkts0 or (RoutingTransition ni Accepted) = do
    s <- do
        drg0 <- readTVar (onionDRG or)
        bkts <- readTVar bkts0
        let antibias = 2 ^ bucketNumber kademlia (nodeId ni) bkts
            (s,drg)  = randomR (0,antibias - 1) drg0
        writeTVar (onionDRG or) drg
        -- let self = (thisNode bkts) { nodeIP = read "127.0.0.1" }
        -- modifyTVar' (trampolineNodes or) (IntMap.insert (-1) self)
        return s
    -- debias via stochastic filter
    when (s == 0) $ do
    ns <- readTVar (trampolineIds or)
    case HashMap.lookup (nodeId ni) ns of
        Just _ -> return ()
        Nothing -> do
            cnt <- readTVar (trampolineCount or)
            writeTChan (routeLog or) $ "ONION trampoline Accepted " ++ unwords [show cnt, show ni]
            modifyTVar' (trampolineIds or) (HashMap.insert (nodeId ni) cnt)
            modifyTVar' (trampolineNodes or) (IntMap.insert cnt ni)
            writeTVar (trampolineCount or) (succ cnt)
hookBucketList _ _ or (RoutingTransition ni Stranger) = do
    ns <- readTVar (trampolineIds or)
    case HashMap.lookup (nodeId ni) ns of
        Just n -> do writeTVar (trampolineIds or) (HashMap.delete (nodeId ni) ns)
                     cnt <- pred <$> readTVar (trampolineCount or)
                     writeTVar (trampolineCount or) cnt
                     if n == cnt
                        then modifyTVar' (trampolineNodes or) (IntMap.delete n)
                        else do lastnode <- (IntMap.! cnt) <$> readTVar (trampolineNodes or)
                                modifyTVar' (trampolineNodes or) (IntMap.insert n lastnode . IntMap.delete cnt)
                     writeTChan (routeLog or) $ "ONION trampoline Stranger " ++ unwords [show n,show ni]
        Nothing -> return ()
hookBucketList _ _ _ _ = return () -- ignore Applicant event.

newtype IPClass = IPClass Word32
 deriving Eq

ipkey :: IPClass -> Int
ipkey (IPClass k) = fromIntegral k

nodeClass :: NodeInfo -> IPClass
nodeClass = ipClass. nodeAddr

ipClass :: SockAddr -> IPClass
ipClass (SockAddrInet _ addr)                = IPClass (fromBE32 addr .&. 0xFF000000)
ipClass (SockAddrInet6 _ _ (hiword,_,_,_) _) = IPClass hiword
ipClass _                                    = IPClass 0 -- unreachable.