path: root/OnionRouter.hs

{-# LANGUAGE NondecreasingIndentation #-}
{-# LANGUAGE LambdaCase #-}
module OnionRouter where

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

import Control.Concurrent.STM
import Control.Concurrent.STM.TArray
import Control.Monad
import Crypto.Random
import Data.Array.MArray
import Data.Bits
import Data.Bool
import Data.List
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.Time.Clock.POSIX
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.Timeout

-- 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
    { -- | For every query, we remember the destination and source keys
      -- so we can decrypt the response.
      pendingQueries  :: TVar (Word64Map PendingQuery)
      -- | The current 12 routes that may be assigned to outgoing packets.
    , routeMap        :: TArray Int (Maybe RouteRecord)
      -- | A set of nodes used to query for random route nodes. These aren't
      -- used directly in onion routes, they are queried for route nodes that
      -- are nearby randomly selected ids.
      --
      -- These nodes are chosen from the kademlia buckets and when one of them
      -- is evicted from a bucket, it is no longer used as a trampoline node.
      --
      -- Also, currently our own address is (unnecessarily) stored here at
      -- index (-1).
    , trampolineNodes :: TVar (IntMap NodeInfo)
      -- | This map associates 'NodeId' values with the corresponding
      -- 'trampolineNodes' index.
    , trampolineIds   :: TVar (HashMap NodeId Int)
      -- | Indicates the current size of 'trampolineNodes'.
    , trampolineCount :: TVar Int
      -- | The pseudo-random generator used to select onion routes.
    , onionDRG        :: TVar ChaChaDRG
      -- | Building onion routes happens in a dedicated thread. See 'forkRouteBuilder'.
    , routeThread     :: ThreadId
      -- | Each of the 12 routes has a version number here that is set larger
      -- than the 'routeVersion' set in 'routeMap' when the route should be
      -- discarded and replaced with a fresh one.
    , pendingRoutes   :: TArray Int Int
      -- | Debug prints are written to this channel which is then flushed to
      -- 'routeLogger'.
    , routeLog        :: TChan String
      -- | User supplied log function.
    , routeLogger     :: String -> IO ()
    }

data PendingQuery = PendingQuery
    { pendingVersion     :: !Int -- ^ Remembered version number so timeouts can signal a rebuild.
    , pendingDestination :: OnionDestination RouteId
    }
 deriving Show

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

-- 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 :: POSIXTime -> OnionRoute -> Maybe RouteRecord -> Maybe RouteRecord
freshRoute birthday r mrec = Just $ RouteRecord
    { storedRoute    = r
    , responseCount  = 0
    , timeoutCount   = 0
    , routeVersion   = maybe 0 succ $ routeVersion <$> mrec
    , routeBirthTime = birthday
    }

modifyArray :: TArray Int r -> (r -> r) -> Int -> STM ()
modifyArray a f i = do
    mx <- readArray a i
    writeArray a i $ f mx
{-# INLINE modifyArray #-}

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

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

newtype RouteEvent = BuildRoute RouteId

newOnionRouter :: (String -> IO ()) -> IO OnionRouter
newOnionRouter perror = 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 <- newArray (0,11) Nothing
        tn <- newTVar IntMap.empty
        ti <- newTVar HashMap.empty
        tc <- newTVar 0
        pr <- newArray (0,11) 0
        rlog <- newTChan
        return OnionRouter
            { pendingRoutes = pr
            , onionDRG = drg
            , pendingQueries = pq
            , routeMap = rm
            , trampolineNodes = tn
            , trampolineIds = ti
            , trampolineCount = tc
            , routeLog = rlog
            , routeThread = error "Failed to invoke forkRouteBuilder"
            , routeLogger = perror
            }
    return or

forkRouteBuilder :: OnionRouter -> (NodeId -> NodeInfo -> IO (Maybe [NodeInfo])) -> IO OnionRouter
forkRouteBuilder or getnodes = do
    tid <- forkIO $ do
        me <- myThreadId
        labelThread me "OnionRouter"
        forever $ do
            let checkRebuild :: Int -> Int -> STM RouteEvent
                checkRebuild rid wanted_ver = do
                    current_ver <- fmap routeVersion <$> readArray (routeMap or) rid
                    writeTChan (routeLog or) $ "ONION checkRebuild "++show(current_ver,wanted_ver)
                    check $ maybe True (< wanted_ver) current_ver
                    return $ BuildRoute $ RouteId rid
            io <- atomically $ {-# SCC "forkRouteBuilder.log" #-}
                   (readTChan (routeLog or) >>= return . routeLogger or)
                `orElse` {-# SCC "forkRouteBuilder.checkRebuild" #-}
                   (let stms = map (\rid -> checkRebuild rid =<< readArray (pendingRoutes or) rid)
                                   [0..11]
                        in do event <- foldr1 orElse stms
                              return $ handleEvent getnodes or { routeThread = me } event)
            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))

-- Repeatedly attempt to select 3 nodes as a secure onion route letting 1 second
-- elapse between retries.
--
-- Only the DRG random seed is updated.  Hopefully another thread will change the
-- trampolineNodes set so that selection can succeed.
selectTrampolines :: OnionRouter -> IO [NodeInfo]
selectTrampolines or = do
    myThreadId >>= flip labelThread ("OnionRouter.selectTrampolines")
    atomically (selectTrampolines' or) >>= \case
        Left ns -> do
            -- atomically $ writeTChan (routeLog or)
            routeLogger or $ unwords
                ( "ONION Discarding insecure trampolines:" : (map (show . nodeAddr) ns) )
            myThreadId >>= flip labelThread ("OnionRouter.selectTrampolines.sleep")
            case ns of
                [_,_,_] -> threadDelay 1000000 -- wait 1 second if we failed the distinct3by predicate.
                _       -> threadDelay 5000000 -- wait 5 seconds if insufficient nodes.
            myThreadId >>= flip labelThread ("OnionRouter.selectTrampolines")
            selectTrampolines or
        Right ns -> do
            myThreadId >>= flip labelThread ("OnionRouter")
            return ns

choose3 :: (Integral a, DRG drg) => drg -> a -> ([a], drg)
choose3 drg0 cnt = ([a,b,c], drg)
 where
        (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
        [ac,bc] = sort [a,b]
        c1 | c0 < ac    = c0
           | otherwise  = c0 + 1
        c | c1 < bc   = c1
          | otherwise = c1 + 1

-- Select 3 indices into the trampolineNodes set and returns the associated
-- nodes provided they are suitable for use in an onion route.  Otherwise, it
-- returns Left with the nodes that were selected.
--
-- The only write this function does to STM state is that the onionDRG random
-- seed will be updated.
selectTrampolines' :: OnionRouter -> STM (Either [NodeInfo] [NodeInfo])
selectTrampolines' or = do
    cnt <- readTVar (trampolineCount or)
    ts <- readTVar (trampolineNodes or)
    drg0 <- readTVar (onionDRG or)
    let ([a,b,c],drg) = choose3 drg0 cnt
        ns = mapMaybe (\n -> IntMap.lookup n ts) [a,b,c]
    ns' <- case ns of
            [an,bn,cn] | distinct3by nodeClass an bn cn
                       -> return $ Right ns
            _          -> return $ Left ns
    writeTVar (onionDRG or) drg
    return ns'

handleEvent :: (NodeId -> NodeInfo -> IO (Maybe [NodeInfo])) -> OnionRouter -> RouteEvent -> IO ()
handleEvent getnodes or e@(BuildRoute (RouteId rid)) = do
    routeLogger or $ "ONION Rebuilding RouteId " ++ show rid
    mb <- do
        ts <- selectTrampolines or
        join . atomically $ do
        drg <- readTVar (onionDRG or)
        av <- newTVar Nothing
        bv <- newTVar Nothing
        cv <- 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 = fmap (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
                        -- This timeout should be unnecessary... But I'm paranoid.
                        -- Note: 10 seconds should be sufficient for typical get-nodes queries.
                        tm <- timeout 20000000 $ atomically $ do -- Wait for all 3 results.
                            rs <- catMaybes <$> sequence [readTVar av,readTVar bv,readTVar cv]
                            case rs of [_,_,_] -> do
                                                    return $ catMaybes $ catMaybes rs
                                                    -- self <- IntMap.lookup (-1) <$> readTVar (trampolineNodes or)
                                                    -- return $ maybe (catMaybes rs) (\x -> [x,x,x]) self
                                       _       -> retry
                        maybe (routeLogger or "ONION: Unexpected sendq timeout!" >> return [])
                              return
                              tm
                return $ do
                    myThreadId >>= flip labelThread ("OnionRouter.sendqs")
                    nodes <- case ts of
                        [_,_,_] -> sendqs
                        _       -> return []
                    myThreadId >>= flip labelThread ("OnionRouter")
                    routeLogger or $ unlines
                        [ "ONION trampolines:   " ++ show ts
                        , "ONION query results: " ++ show nodes ]
                    case nodes of
                        [a,b,c] | distinct3by nodeClass a b c -> do
                            atomically $ do
                              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
    now <- getPOSIXTime
    atomically $ maybe (return ()) -- writeTVar (pendingRoutes or IntMap.! rid) True)
                       (\r -> do modifyArray (routeMap or)
                                             (freshRoute now r)
                                             rid
                                 v <- routeVersion . fromJust <$> readArray (routeMap or) rid
                                 writeArray (pendingRoutes or) rid v
                                 )
                       mb
    case mb of
        Just _  -> routeLogger or $ "ONION Finished RouteId " ++ show rid
        Nothing -> routeLogger or $ "ONION Failed RouteId " ++ show rid

lookupSender :: OnionRouter -> SockAddr -> Nonce8 -> IO (Maybe (OnionDestination 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
        od <- result
        let nid = nodeId $ onionNodeInfo $ pendingDestination od
        ni <- either (const Nothing) Just $ nodeInfo nid saddr
        Just (OnionDestination (onionAliasSelector $ pendingDestination od)
                               ni
                               (Just $ routeId nid))

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

lookupTimeout :: OnionRouter -> Nonce8 -> OnionDestination r -> STM (OnionDestination RouteId, Int)
lookupTimeout or n8 (OnionDestination asel ni Nothing) = do
    let RouteId rid = routeId (nodeId ni)
    mrr <- readArray (routeMap or) rid
    writeTChan (routeLog or) $ unwords ["ONION lookupTimeout " ,show rid]
    case mrr of
        Just rr -> return ( OnionDestination asel ni (Just $ RouteId rid), timeoutForRoute rr)
        Nothing -> return ( OnionDestination asel ni Nothing             , 0 )

hookQueries :: OnionRouter -> (tid -> Nonce8)
                           -> TransactionMethods d tid (OnionDestination RouteId) x
                           -> TransactionMethods d tid (OnionDestination RouteId) x
hookQueries or t8 tmethods = TransactionMethods
    { dispatchRegister = \mvar od d -> {-# SCC "hookQ.dispatchRegister" #-} do --  :: MVar x -> d -> STM (tid, d)
        let ni = onionNodeInfo od
            rid@(RouteId ridn) = fromMaybe (routeId (nodeId ni)) $ onionRouteSpec od
        wanted <- {-# SCC "hookQ.wanted" #-} (readArray (pendingRoutes or) ridn)
        mr <- {-# SCC "hookQ.mr_action" #-} (readArray (routeMap or) ridn)
        -- Block query until a route is ready.
        check $ fromMaybe False $ do
            RouteRecord{routeVersion=rv} <- {-# SCC "hookQ.mr" #-} mr
            return $ wanted <= rv
        (tid,d') <- dispatchRegister tmethods mvar od d
        let Nonce8 w8 = t8 tid
            od' = case od of
                OnionDestination {} -> od { onionRouteSpec = Just rid }
                OnionToOwner a b    -> OnionToOwner a b -- Type cast.
        let pq = PendingQuery { pendingDestination = od'
                              , pendingVersion = maybe 0 routeVersion mr
                              }
        pqs <- readTVar (pendingQueries or)
        -- check $ W64.size pqs < 20
        modifyTVar' (pendingQueries or) (W64.insert w8 pq)
        writeTChan (routeLog or) $ "ONION query add " ++ unwords [ show (Just $ pendingVersion pq,w8), ":=", show ni ]
        return (tid,d')
    , dispatchResponse = \tid x d -> {-# SCC "hookQ.dispatchResponse" #-} do -- :: tid -> x -> d -> STM (d, IO ())
        let Nonce8 w8 = t8 tid
        mb <- W64.lookup w8 <$> readTVar (pendingQueries or)
        modifyTVar' (pendingQueries or) (W64.delete w8)
        forM_ mb $ \pq -> do
            let od = pendingDestination pq
                RouteId rid = fromMaybe (routeId (nodeId (onionNodeInfo od)))
                                $ onionRouteSpec od
            modifyArray (routeMap or) (fmap gotResponse) rid
        writeTChan (routeLog or) $ "ONION query del " ++ show (fmap pendingVersion mb, w8)
        dispatchResponse tmethods tid x d
    , dispatchCancel = \tid d -> {-# SCC "hookQ.dispatchCancel" #-} do -- :: tid -> d -> STM d
        let Nonce8 w8 = t8 tid
        mb <- W64.lookup w8 <$> readTVar (pendingQueries or)
        modifyTVar' (pendingQueries or) (W64.delete w8)
        forM_ mb $ \pq -> do
            let od = pendingDestination pq
                RouteId rid = fromMaybe (routeId (nodeId (onionNodeInfo od)))
                                $ onionRouteSpec od
            mrr <- readArray (routeMap or) rid
            forM_ mrr $ \rr -> do
            when (routeVersion rr == pendingVersion pq) $ do
            let expireRoute = modifyArray (pendingRoutes or) expire rid
                expire ver | ver <= succ (pendingVersion pq) = succ (pendingVersion pq)
                           | otherwise                       = ver
            modifyArray (routeMap or) (fmap gotTimeout) rid
            case rr of
                RouteRecord{ responseCount = 0
                           , timeoutCount  = c
                           , routeVersion = v } | c >= 5 -> expireRoute
                RouteRecord{ responseCount = 1
                           , timeoutCount = c
                           , routeVersion = v } | c >= 10 -> expireRoute
                RouteRecord{ timeoutCount = c
                           , routeVersion = v } | c >= 20 -> expireRoute
                _ -> return ()
        writeTChan (routeLog or) $ "ONION query can " ++ show (fmap pendingVersion mb, w8)
        dispatchCancel tmethods tid d
    }


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
        do -- Store localhost as trampoline node (-1).
           -- This is not used, but harmless.  I'm leaving it in for
           -- testing purposes.
           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
                     case compare n cnt of
                        EQ -> modifyTVar' (trampolineNodes or) (IntMap.delete n)
                        LT -> do lastnode <- (IntMap.! cnt) <$> readTVar (trampolineNodes or)
                                 modifyTVar' (trampolineNodes or)
                                             (IntMap.insert n lastnode . IntMap.delete cnt)
                                 modifyTVar' (trampolineIds or)
                                             (HashMap.delete (nodeId ni) . HashMap.insert (nodeId lastnode) n)
                        GT -> writeTChan (routeLog or) $ "BUG!! Trampoline maps are out of sync."
                     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= either ipClass' ipClass' . either4or6

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