path: root/OnionTransport.hs

{-# LANGUAGE DataKinds                  #-}
{-# LANGUAGE FlexibleContexts           #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE KindSignatures             #-}
{-# LANGUAGE StandaloneDeriving #-}
{-# LANGUAGE LambdaCase                 #-}
{-# LANGUAGE RankNTypes                 #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# LANGUAGE TupleSections              #-}
{-# LANGUAGE TypeOperators              #-}
{-# LANGUAGE UndecidableInstances       #-}
module OnionTransport
    ( parseOnionAddr
    , encodeOnionAddr
    , forwardOnions
    , OnionToOwner(..)
    , OnionMessage(..)
    , DataToRoute(..)
    , AnnounceResponse(..)
    , AnnounceRequest(..)
    , Forwarding(..)
    , ReturnPath(..)
    , OnionRequest(..)
    , OnionResponse(..)
    , Addressed(..)
    , UDPTransport
    , KeyRecord(..)
    , encrypt
    , decrypt
    ) where

import Network.QueryResponse
import ToxCrypto hiding (encrypt,decrypt)
import ToxAddress
import qualified ToxCrypto
import DHTTransport (NodeInfo(..),NodeId(..),SendNodes,nodeInfo)

import Control.Arrow
import qualified Data.ByteString as B
         ;import Data.ByteString (ByteString)
import Data.Functor.Identity
import Data.Serialize            as S
import Data.Typeable
import Data.Word
import GHC.TypeLits
import Network.Socket
import GHC.Generics

type HandleLo a = Maybe (Either String (ByteString, SockAddr)) -> IO a

type UDPTransport = Transport String SockAddr ByteString


getOnionAssym :: Get (Assym (Encrypted DataToRoute))
getOnionAssym = getAliasedAssym

putOnionAssym :: Serialize a => Word8 -> Put -> Assym a -> Put
putOnionAssym typ p a = put typ >> p >> putAliasedAssym a

data OnionMessage (f :: * -> *)
    = OnionAnnounce (Assym (f (AnnounceRequest,Nonce8)))
    | OnionAnnounceResponse Nonce8 Nonce24 (f AnnounceResponse)
    | OnionToRoute PublicKey (Assym (f DataToRoute)) -- destination key, aliased Assym
    | OnionToRouteResponse (Assym (f DataToRoute))

data OnionToOwner = OnionToOwner NodeInfo (ReturnPath 3)
                  | OnionToMe SockAddr -- SockAddr is immediate peer in route
    deriving Show


onionToOwner :: Assym a -> ReturnPath 3 -> SockAddr -> Either String OnionToOwner
onionToOwner assym ret3 saddr = do
    ni <- nodeInfo (NodeId $ senderKey assym) saddr
    return $ OnionToOwner ni ret3
-- data CookieAddress = WithoutCookie NodeInfo | CookieAddress Cookie SockAddr


onion :: Sized msg =>
               ByteString
               -> SockAddr
               -> Get (Assym (Encrypted msg) -> t)
               -> Either String (t, OnionToOwner)
onion bs saddr getf = do (f,(assym,ret3)) <- runGet ((,) <$> getf <*> getOnionRequest) bs
                         oaddr <- onionToOwner assym ret3 saddr
                         return (f assym, oaddr)


parseOnionAddr :: (ByteString, SockAddr) -> Either (OnionMessage Encrypted,OnionToOwner) (ByteString,SockAddr)
parseOnionAddr (msg,saddr)
    | Just (typ,bs) <- B.uncons msg
    , let right = Right (msg,saddr)
          query = either (const right) Left
          response = either (const right) (Left . (, OnionToMe saddr))
    = case typ of
    0x83 -> query $ onion bs saddr (pure OnionAnnounce)                          -- Announce Request
    0x85 -> query $ onion bs saddr (OnionToRoute <$> getPublicKey)               -- Onion Data Request
    0x84 -> response $ runGet (OnionAnnounceResponse <$> get <*> get <*> get) bs -- Announce Response
    0x86 -> response $ runGet (OnionToRouteResponse <$> getOnionAssym) bs        -- Onion Data Response
    _ -> right

encodeOnionAddr :: (OnionMessage Encrypted,OnionToOwner) -> (ByteString, SockAddr)
encodeOnionAddr (msg,addr) = ( runPut (putmsg >> putpath), saddr )
 where
    (saddr,putpath) | OnionToOwner ni p <- addr = (nodeAddr ni, put p)
                    | OnionToMe a       <- addr = (a, return ())

    putmsg | OnionAnnounce a                <- msg = putOnionAssym 0x83 (return ()) a
           | OnionToRoute pubkey a          <- msg = putOnionAssym 0x85 (putPublicKey pubkey) a
           | OnionToRouteResponse a         <- msg = putOnionAssym 0x86 (return ()) a
           | OnionAnnounceResponse n8 n24 x <- msg = put (0x84 :: Word8) >> put n8 >> put n24 >> put x

forwardOnions :: TransportCrypto -> UDPTransport -> UDPTransport -- HandleHi a -> IO a -> HandleLo a
forwardOnions crypto udp = udp { awaitMessage = await' }
 where
    -- forMe :: HandleHi
    -- forThem :: handleLo
    await' :: HandleLo a -> IO a
    await' forThem = awaitMessage udp $ \case
        m@(Just (Right (bs,saddr))) -> case B.head bs of
            0x80 -> forward forThem bs $ handleOnionRequest (Proxy :: Proxy 0) crypto saddr (forThem m)
            0x81 -> forward forThem bs $ handleOnionRequest (Proxy :: Proxy 1) crypto saddr (forThem m)
            0x82 -> forward forThem bs $ handleOnionRequest (Proxy :: Proxy 2) crypto saddr (forThem m)
            0x8c -> forward forThem bs $ handleOnionResponse (Proxy :: Proxy 3) crypto saddr (forThem m)
            0x8d -> forward forThem bs $ handleOnionResponse (Proxy :: Proxy 2) crypto saddr (forThem m)
            0x8e -> forward forThem bs $ handleOnionResponse (Proxy :: Proxy 1) crypto saddr (forThem m)
            _    -> forThem m
        m -> forThem m

forward :: forall c b b1.
          Serialize b =>
              (Maybe (Either String b1) -> c) -> ByteString -> (b -> c) -> c
forward forMe bs f = either (forMe . Just . Left) f $ decode $ B.tail bs

-- n = 0, 1, 2
data OnionRequest (n :: Nat) = OnionRequest
    { onionNonce    :: Nonce24
    , onionForward  :: Forwarding (3 - n) (OnionMessage Encrypted)
    , pathFromOwner :: ReturnPath n
    }

instance ( Serialize (Forwarding (3 - n) (OnionMessage Encrypted))
         , Serialize (ReturnPath n)
         ) => Serialize (OnionRequest n) where
    get = OnionRequest <$> get <*> get <*> get
    put (OnionRequest n f p) = put n >> put f >> put p

-- n = 1, 2, 3
-- Attributed (Encrypted (

data OnionResponse (n :: Nat) = OnionResponse
    { pathToOwner :: ReturnPath n
    , msgToOwner  :: OnionMessage Encrypted
    }

instance ( KnownNat n, Serialize (ReturnPath n) ) => Serialize (OnionResponse n) where
    get = OnionResponse <$> get <*> get
    put (OnionResponse p m) = put p >> put m


data Addressed a = Addressed { sockAddr :: SockAddr, unaddressed :: a }

data ReturnPath (n :: Nat) where
    NoReturnPath :: ReturnPath 0
    ReturnPath :: Nonce24 -> Encrypted (Addressed (ReturnPath (n - 1))) -> ReturnPath n

instance KnownNat n => Sized (Addressed (ReturnPath n)) where size = _todo
--     -- Size: 59 = 1(family) + 16(ip) + 2(port) +16(mac) + 24(nonce)

instance (Serialize (Encrypted (Addressed (ReturnPath (n - 1))))) => Serialize (ReturnPath n) where
    get = ReturnPath <$> get <*> get
    put (ReturnPath n24 p) = put n24 >> put p

instance KnownNat n => Show (ReturnPath n) where
    show rpath = "ReturnPath" ++ show (natVal rpath)


-- instance KnownNat n => Serialize (ReturnPath n) where
--     -- Size: 59 = 1(family) + 16(ip) + 2(port) +16(mac) + 24(nonce)
--     get = ReturnPath <$> getBytes ( 59 * (fromIntegral $ natVal $ Proxy @n) )
--     put (ReturnPath bs) = putByteString bs


data Forwarding (n :: Nat) msg where
    NotForwarded :: msg -> Forwarding 0 msg
    Forwarding :: Assym (Encrypted (Addressed (Forwarding (n - 1) msg))) -> Forwarding n msg

instance (KnownNat n, Sized msg) => Sized (Addressed (Forwarding n msg)) where size = _todo

instance (Serialize msg, Serialize (Encrypted (Addressed (Forwarding (n - 1) msg)))) => Serialize (Forwarding n msg) where
    get = Forwarding <$> getAliasedAssym
    put (Forwarding x) = putAliasedAssym x


handleOnionRequest :: KnownNat n => proxy n -> TransportCrypto -> SockAddr -> IO a -> OnionRequest n -> IO a
handleOnionRequest = _todo

handleOnionResponse :: KnownNat n => proxy n -> TransportCrypto -> SockAddr -> IO a -> OnionResponse n -> IO a
handleOnionResponse = _todo

data AnnounceRequest = AnnounceRequest
    { announcePingId  :: Nonce32 -- Ping ID
    , announceSeeking :: NodeId  -- Public key we are searching for
    , announceKey     :: NodeId  -- Public key that we want those sending back data packets to use
    }

instance Sized AnnounceRequest where size = _todo

instance S.Serialize AnnounceRequest where
    get = AnnounceRequest <$> S.get <*> S.get <*> S.get
    put (AnnounceRequest p s k) = S.put (p,s,k)

getOnionRequest :: Sized msg => Get (Assym (Encrypted msg), ReturnPath 3)
getOnionRequest = (,) <$> getAliasedAssym <*> _todo

data KeyRecord = NotStored    Nonce32
               | SendBackKey  PublicKey
               | Acknowledged Nonce32

instance S.Serialize KeyRecord where
    get = do
        is_stored <- S.get :: S.Get Word8
        case is_stored of
            1 -> SendBackKey  <$> getPublicKey
            2 -> Acknowledged <$> S.get
            _ -> NotStored    <$> S.get
    put (NotStored n32)    = S.put (0 :: Word8) >> S.put n32
    put (SendBackKey key)  = S.put (1 :: Word8) >> putPublicKey key
    put (Acknowledged n32) = S.put (2 :: Word8) >> S.put n32

data AnnounceResponse = AnnounceResponse
    { is_stored     :: KeyRecord
    , announceNodes :: SendNodes
    }

instance Sized AnnounceResponse where
    size = VarSize $ \AnnounceResponse {} -> _todo

instance S.Serialize AnnounceResponse where
    get = AnnounceResponse <$> S.get <*> S.get
    put (AnnounceResponse st ns) = S.put st >> S.put ns

data DataToRoute = DataToRoute
    { dataFromKey :: PublicKey
    , dataToRoute :: Encrypted (Word8,ByteString)
    }

instance Sized DataToRoute where
    size = VarSize $ \DataToRoute {} -> _todo

instance Serialize DataToRoute where
    get = return $ DataToRoute _todo _todo
    put _ = return () -- todo

encrypt :: TransportCrypto -> OnionMessage Identity -> OnionToOwner -> (OnionMessage Encrypted, OnionToOwner)
encrypt crypto msg rpath = (transcode (encryptMessage crypto) msg, rpath)

encryptMessage :: Serialize a =>
                  TransportCrypto -> Nonce24 -> Either (Identity a) (Assym (Identity a)) -> Encrypted a
encryptMessage crypto n (Right a) = ToxCrypto.encrypt secret plain
 where
    secret = computeSharedSecret (transportSecret crypto) (senderKey a) n
    plain  = encodePlain $ runIdentity $ assymData a
encryptMessage crypto n (Left x) = ToxCrypto.encrypt secret plain
 where
    secret = computeSharedSecret (transportSecret crypto) _todo n
    plain  = encodePlain $ runIdentity $ x

decrypt :: TransportCrypto -> OnionMessage Encrypted -> OnionToOwner -> Either String (OnionMessage Identity, OnionToOwner)
decrypt crypto msg addr = (, addr) <$> (sequenceMessage $ transcode (decryptMessage crypto) msg)

decryptMessage :: Serialize x =>
                  TransportCrypto
                  -> Nonce24
                  -> Either (Encrypted x) (Assym (Encrypted x))
                  -> (Either String ∘ Identity) x
decryptMessage crypto n (Right assymE) = plain $ ToxCrypto.decrypt secret e
 where
    secret = computeSharedSecret (transportSecret crypto) (senderKey assymE) n
    e      = assymData assymE
    plain  = Composed . fmap Identity . (>>= decodePlain)
decryptMessage crypto n (Left e) = _todo


sequenceMessage :: Applicative m => OnionMessage (m ∘ f) -> m (OnionMessage f)
sequenceMessage (OnionAnnounce a)        = fmap OnionAnnounce        $ sequenceA $ fmap uncomposed a
sequenceMessage (OnionToRoute pub a)     = fmap (OnionToRoute pub)   $ sequenceA $ fmap uncomposed a
sequenceMessage (OnionToRouteResponse a) = fmap OnionToRouteResponse $ sequenceA $ fmap uncomposed a
sequenceMessage (OnionAnnounceResponse n8 n24 dta) = OnionAnnounceResponse n8 n24 <$> uncomposed dta

transcode :: forall f g. (forall a. Serialize a => Nonce24 -> Either (f a) (Assym (f a)) -> g a) -> OnionMessage f -> OnionMessage g
transcode f (OnionAnnounce a)        = OnionAnnounce        $ a { assymData = f (assymNonce a) (Right a) }
transcode f (OnionToRoute pub a)     = OnionToRoute pub     $ a { assymData = f (assymNonce a) (Right a) }
transcode f (OnionToRouteResponse a) = OnionToRouteResponse $ a { assymData = f (assymNonce a) (Right a) }
transcode f (OnionAnnounceResponse n8 n24 dta) = OnionAnnounceResponse n8 n24 $ f n24 $ Left dta