path: root/ToxMessage.hs

{-# LANGUAGE PatternSynonyms            #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DeriveFunctor, DeriveTraversable,DeriveDataTypeable #-}
{-# LANGUAGE GADTs #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}
{-# LANGUAGE ExistentialQuantification #-}
{-# LANGUAGE DataKinds, KindSignatures #-}
module ToxMessage where

import Debug.Trace
import Data.ByteString (ByteString)
import qualified Crypto.MAC.Poly1305 as Poly1305 (Auth(..))
import qualified Crypto.PubKey.Curve25519 as Curve25519
import Data.ByteArray as BA (ByteArrayAccess, Bytes)
import qualified Data.ByteArray as BA
import qualified Data.ByteString as B
import qualified Data.ByteString.Char8 as C8
import qualified Data.ByteString.Base16 as Base16
import Data.Bits
import Data.Hashable
import Data.Bits.ByteString ()
import Data.Word
import Data.Data
import Data.Ord
import Data.Serialize
import Foreign.Ptr
import Foreign.Marshal.Alloc
import System.Endian
import Foreign.Storable
import GHC.TypeLits
import Data.Tuple

newtype Auth = Auth Poly1305.Auth
 deriving (Eq, ByteArrayAccess)

instance Ord Auth where
    compare (Auth a) (Auth b) = comparing (BA.convert :: Poly1305.Auth -> Bytes) a b

instance Data Auth where
    gfoldl k z x = z x

    -- Well, this is a little wonky... XXX
    gunfold k z c = k (z (Auth . Poly1305.Auth . (BA.convert :: ByteString -> Bytes)))

    toConstr _ = con_Auth

    dataTypeOf _ = mkDataType "ToxMessage" [con_Auth]

con_Auth :: Constr
con_Auth = mkConstr (dataTypeOf (Auth (error "con_Auth"))) "Auth" [] Prefix

instance Serialize Auth where
    get = Auth . Poly1305.Auth . BA.convert <$> getBytes 16
    put (Auth (Poly1305.Auth bs)) = putByteString $ BA.convert bs
--
-- | An 'Aliased' 'PubKey' is one that is not the DHT key and so should not go
-- into the kademlia routing table buckets.
--
-- Note: This includes the long-term tox-id key that humans use to friend each
-- other and is often refered to as your "real public key" by the Tox
-- documents.  For the purposes of the DHT, it is an alias.
newtype Aliased a = Aliased a
 deriving (Eq,Ord,Show,Data,Functor,Foldable,Traversable)

newtype Nonce24 = Nonce24 ByteString
 deriving (Eq, Ord, ByteArrayAccess,Data)

quoted :: ShowS -> ShowS
quoted shows s = '"':shows ('"':s)

bin2hex :: ByteArrayAccess bs => bs -> String
bin2hex = C8.unpack . Base16.encode . BA.convert

instance Show Nonce24 where
    showsPrec d nonce = quoted (mappend $ bin2hex nonce)

instance Serialize Nonce24 where
    get = Nonce24 <$> getBytes 24
    put (Nonce24 bs) = putByteString bs

newtype Nonce8 = Nonce8 Word64
 deriving (Eq, Ord, Data, Serialize)

instance ByteArrayAccess Nonce8 where
    length _ = 8
    withByteArray (Nonce8 w64) kont =
        allocaBytes 8 $ \p -> do
            poke (castPtr p :: Ptr Word64) $ toBE64 w64
            kont p

instance Show Nonce8 where
    showsPrec d nonce = quoted (mappend $ bin2hex nonce)




-- TODO: This should probably be represented by Curve25519.PublicKey, but
-- ByteString has more instances...
newtype PubKey = PubKey ByteString
 deriving (Eq,Ord,Data, ByteArrayAccess, Bits, Hashable)

instance Serialize PubKey where
    get = PubKey <$> getBytes 32
    put (PubKey bs) = putByteString bs

instance Show PubKey where
    show (PubKey bs) = C8.unpack $ Base16.encode bs

instance FiniteBits PubKey where
    finiteBitSize _ = 256

instance Read PubKey where
    readsPrec _ str
        | (bs, xs) <- Base16.decode $ C8.pack str
        , B.length bs == 32
                    = [ (PubKey bs, drop 64 str) ]
        | otherwise = []



-- | A chunk of data encrypted with public-key cryptography.
data ImplicitAssymetric = ImplicitAssymetric
    { assymetricAuth  :: Auth
    , assymetricBytes :: ByteString
    }
 deriving (Eq, Ord,Data)

getRemaining :: Get ByteString
getRemaining = remaining >>= getBytes

instance Serialize ImplicitAssymetric where
    get = ImplicitAssymetric <$> get <*> getRemaining
    put (ImplicitAssymetric auth bs) = put auth >> putByteString bs

-- | Like ImplicitAssymetric, but includes the nonce used to encrypt.
data UnclaimedAssymetric = UnclaimedAssymetric
    { assymetricNonce :: Nonce24
    , assymetricData  :: !ImplicitAssymetric
    }
 deriving (Eq, Ord, Data)

instance Serialize UnclaimedAssymetric where
    get = UnclaimedAssymetric <$> get <*> get
    put (UnclaimedAssymetric nonce dta) = put nonce >> put dta

-- | Like UnclaimedAssymetric, but includes the public key of the sender.
data Assymetric = Assymetric
    { senderKey :: PubKey
    , sent      :: !UnclaimedAssymetric
    }
 deriving (Eq, Ord,Data)

-- get requires isolate.
-- sender key, then nonce
instance Serialize Assymetric where
    get = Assymetric <$> get <*> get
    put (Assymetric key dta) = put key >> put dta

-- Aliased packets have the sender key and nonce reversed.
instance Serialize (Aliased Assymetric) where
    get = do
        nonce <- get
        key <- get
        dta <- get
        return $ Aliased (Assymetric key (UnclaimedAssymetric nonce dta))

    put (Aliased (Assymetric key (UnclaimedAssymetric nonce dta))) = do
        put nonce
        put key
        put dta


newtype Cookie = Cookie UnclaimedAssymetric
 deriving (Eq, Ord,Data)

newtype ReturnPath (n::Nat) = ReturnPath ByteString
 deriving (Eq, Ord,Data)

emptyReturnPath :: ReturnPath 0
emptyReturnPath = ReturnPath B.empty

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 Symmetric = Symmetric
    { symmetricNonce :: Nonce24
    , symmetricAuth  :: Auth
    , symmetricBytes :: ByteString
    }
 deriving (Eq, Ord,Data)

-- Test run histogram:
-- 377    PongType
-- 387    DataRequestType
-- 3238   PingType
-- 9231   DHTRequestType
-- 10299  PacketKind 130
-- 12626  PacketKind 129
-- 16596  OnionRequest0
-- 16742  SendNodesType
-- 41877  Announce
-- 81793  GetNodesType
data Packet where
    Ping          :: Assymetric -> Packet -- 0x00 -- Assymetric query
    Pong          :: Assymetric -> Packet -- 0x01 -- Assymetric response

    GetNodes      :: Assymetric -> Packet -- 0x02 -- Assymetric query
    SendNodes     :: Assymetric -> Packet -- 0x04 -- Assymetric response

    CookieRequest :: Assymetric -> Packet -- 0x18
    CookieResponse :: UnclaimedAssymetric -> Packet -- 0x19

    OnionRequest0 :: Assymetric -> Packet -- 0x80


    CryptoHandshake :: Cookie -> UnclaimedAssymetric -> Packet -- 0x1a

    CryptoData :: Word16 -> ImplicitAssymetric -> Packet -- 0x1b

    DHTRequest :: PubKey -> Assymetric -> Packet -- 0x20 -- Sometimes Assymetric query

    DataToRoute :: PubKey -> Aliased Assymetric -> ReturnPath 3 -> Packet
    DataToRouteResponse :: Aliased Assymetric -> Packet

    LanDiscovery :: PubKey -> Packet  -- 0x21

    OnionRequest1 :: Aliased Assymetric -> Symmetric -> Packet -- 0x81
    OnionRequest2 :: Aliased Assymetric -> Symmetric -> Packet -- 0x82

    OnionRequest3 :: ByteString -> Symmetric -> Packet -- 0x82

    Announce :: Aliased Assymetric -> ReturnPath 3 -> Packet --0x83
    AnnounceResponse :: Nonce8 -> UnclaimedAssymetric -> Packet -- 0x84

    OnionResponse3 :: ReturnPath 3 -> Packet -> Packet -- 0x8c
    OnionResponse2 :: Symmetric -> ByteString -> Packet -- 0x8d
    OnionResponse1 :: Symmetric -> ByteString -> Packet -- 0x8e


    GetBootstrapInfo :: ByteString -> Packet -- 0xf0 + 77 bytes                         -- ByteString query
    BootstrapInfo :: Word32 -> ByteString -> Packet -- 0xf0 + version + (256 byte motd) -- ByteSTring response

 deriving (Eq, Ord,Data)

class KnownNat n => OnionPacket n where
    mkOnion :: ReturnPath n -> Packet -> Packet

instance OnionPacket 0 where mkOnion _ = id
instance OnionPacket 3 where mkOnion = OnionResponse3

newtype PacketKind = PacketKind Word8
 deriving (Eq, Ord, Serialize)

-- TODO: Auth fail:
pattern OnionRequest0Type = PacketKind 128 -- 0x80     Onion Request 0
pattern OnionRequest1Type = PacketKind 129 -- 0x81     Onion Request 1
pattern OnionRequest2Type = PacketKind 130 -- 0x82     Onion Request 2
pattern AnnounceType  = PacketKind 131 -- 0x83     Announce Request
pattern AnnounceResponseType = PacketKind 132 -- 0x84     Announce Response

pattern DataRequestType = PacketKind 133 -- 0x85     Onion Data Request (data to route request packet)
pattern DataResponseType = PacketKind 134 -- 0x86     Onion Data Response (data to route response packet)
-- 0x8c     Onion Response 3
-- 0x8d     Onion Response 2
pattern OnionResponse3Type = PacketKind 140 -- 0x8c     Onion Response 3
pattern OnionResponse2Type = PacketKind 141 -- 0x8d     Onion Response 2
pattern OnionResponse1Type = PacketKind 142 -- 0x8e     Onion Response 1
-- 0xf0     Bootstrap Info

--  TODO Fix these fails...
--  GetNodesType decipherAndAuth: auth fail
--  MessageType 128 decipherAndAuth: auth fail
--  MessageType 129 decipherAndAuth: auth fail
--  MessageType 130 decipherAndAuth: auth fail
--  MessageType 131 decipherAndAuth: auth fail
--  MessageType 32 decipherAndAuth: auth fail

-- TODO: Auth fail:
pattern DHTRequestType = PacketKind 32 -- 0x20     DHT Request

pattern PingType      = PacketKind 0 -- 0x00     Ping Request
pattern PongType      = PacketKind 1 -- 0x01     Ping Response
pattern GetNodesType  = PacketKind 2 -- 0x02     Nodes Request
pattern SendNodesType = PacketKind 4 -- 0x04     Nodes Response


instance Show PacketKind where
    showsPrec d PingType = mappend "PingType"
    showsPrec d PongType = mappend "PongType"
    showsPrec d GetNodesType = mappend "GetNodesType"
    showsPrec d SendNodesType = mappend "SendNodesType"
    showsPrec d DHTRequestType = mappend "DHTRequestType"
    showsPrec d OnionRequest0Type = mappend "OnionRequest0"
    showsPrec d OnionResponse1Type = mappend "OnionResponse1"
    showsPrec d OnionResponse3Type = mappend "OnionResponse3"
    showsPrec d AnnounceType = mappend "Announce"
    showsPrec d AnnounceResponseType = mappend "AnnounceResponse"
    showsPrec d DataRequestType = mappend "DataRequestType"
    showsPrec d DataResponseType = mappend "DataResponseType"
    showsPrec d (PacketKind x) = mappend "PacketKind " . showsPrec (d+1) x

pktKind :: Packet -> PacketKind

-- These are (Assymetric -> Assymetric) queries.
pktKind Ping             {} = PacketKind 0x00
pktKind Pong             {} = PacketKind 0x01
pktKind GetNodes         {} = PacketKind 0x02
pktKind SendNodes        {} = PacketKind 0x04


-- This is a (Assymetric -> UnclaimedAssymetric) query
pktKind CookieRequest    {} = PacketKind 0x18
pktKind CookieResponse   {} = PacketKind 0x19

-- Query (Assymetric -> (Nonce8,UnclaimedAssymetric))
pktKind Announce         {} = PacketKind 0x83
pktKind AnnounceResponse {} = PacketKind 0x84

-- Query (Assymetric -> ByteString)
pktKind OnionRequest0    {} = PacketKind 0x80


-- This is a (ByteString -> ByteString) query
pktKind GetBootstrapInfo {} = PacketKind 0xf0
pktKind BootstrapInfo    {} = PacketKind 0xf0


-- Trigering event. No direct response. (PubKey -> ())
pktKind LanDiscovery     {} = PacketKind 0x21

-- Two cases:
--  Half-established:    (Cookie,UnclaimedAssymetric) -> (Cookie,UnclaimedAssymetric)
--  Session established: (Cookie,UnclaimedAssymetric) -> (Word16,ImplicitAssymetric)
pktKind CryptoHandshake  {} = PacketKind 0x1a

-- Session data, no direct response.
--   (reponse to CryptoHandshake, or other data)
pktKind CryptoData       {} = PacketKind 0x1b

-- Two cases:
--    ( (PubKey, Assymetric) -> response )
--    ( (PubKey, Assymetric) -> ()       )
pktKind DHTRequest       {} = PacketKind 0x20


-- Query ( (PubKey,Aliased Assymetric) -> Aliased Assymetric)
pktKind DataToRoute      {}    = PacketKind 0x85
pktKind DataToRouteResponse {} = PacketKind 0x86

-- 3 Queries ( (Aliased Assymetric, Symmetric )
--            -> ( Symmetric, ByteString )     )
pktKind OnionRequest1    {} = PacketKind 0x81
pktKind OnionResponse1   {} = PacketKind 0x8e

pktKind OnionRequest2    {} = PacketKind 0x82
pktKind OnionResponse2   {} = PacketKind 0x8d

pktKind OnionRequest3    {} = PacketKind 0x82
pktKind OnionResponse3   {} = PacketKind 0x8c

data PacketClass =
    AssymetricClass (Assymetric -> Packet) (Packet -> Assymetric)
    | forall n. OnionPacket n => AliasedClass ((Aliased Assymetric,ReturnPath n) -> Packet) (Packet -> (Aliased Assymetric,ReturnPath n))
    | forall n. OnionPacket n => ToRouteClass ((PubKey,(Aliased Assymetric,ReturnPath n)) -> Packet) (Packet -> (PubKey,(Aliased Assymetric,ReturnPath n)))
    | forall n. OnionPacket n => OnionClass ((Packet,ReturnPath n) -> Packet) (Packet -> (Packet,ReturnPath n))
    | NoncedUnclaimedClass (Nonce8 -> UnclaimedAssymetric -> Packet)
                           (Packet -> (Nonce8, UnclaimedAssymetric))
    | Unclassified
{-
data Packet' where
    :: Assymetric -> Packet
    :: UnclaimedAssymetric -> Packet
    :: Word16 -> ImplicitAssymetric -> Packet
    :: PubKey -> Assymetric -> Packet
    :: PubKey -> Packet
    :: Aliased Assymetric -> Symmetric -> Packet
    :: ByteString -> Symmetric -> Packet
    :: Aliased Assymetric -> Packet
    :: Symmetric -> ByteString -> Packet
    :: ByteString -> Packet
    :: Word32 -> ByteString -> Packet
-}

pktClass :: PacketKind -> PacketClass
pktClass (PacketKind 0) = AssymetricClass Ping (\(Ping a) -> a)
pktClass (PacketKind 1) = AssymetricClass Pong (\(Pong a) -> a)
pktClass (PacketKind 2) = AssymetricClass GetNodes (\(GetNodes a) -> a)
pktClass (PacketKind 4) = AssymetricClass SendNodes (\(SendNodes a) -> a)

pktClass (PacketKind 0x18) = AssymetricClass CookieRequest (\(CookieRequest a) -> a)
pktClass (PacketKind 0x80) = AssymetricClass OnionRequest0 (\(OnionRequest0 a) -> a)
pktClass (PacketKind 0x86) = AssymetricClass (DataToRouteResponse . Aliased) (\(DataToRouteResponse (Aliased a)) -> a)

pktClass (PacketKind 0x83) = AliasedClass (uncurry Announce) (\(Announce a r)-> (a,r))
pktClass (PacketKind 0x84) = NoncedUnclaimedClass AnnounceResponse (\(AnnounceResponse n8 uncl)-> (n8,uncl))

pktClass (PacketKind 0x8c) = OnionClass (uncurry OnionResponse3 . swap) (\(OnionResponse3 r a)-> (a,r))

pktClass DataRequestType = ToRouteClass (\(k,(a,r))-> DataToRoute k a r)
                                        (\(DataToRoute k a r) -> (k,(a,r)))

pktClass _ = Unclassified


instance Serialize Packet where
    get = getPacket
    put = putPacket

getPacket :: Get Packet
getPacket = do
    typ <- get
    case pktClass typ of
        AssymetricClass toPacket _      -> toPacket <$> get
        AliasedClass    toPacket _      -> do
            trace ("PARSE "++show typ) $ return ()
            cnt <- remaining
            a <- isolate (cnt - 59*3) get
            r <- get
            trace ("PARSED "++show typ) $ return ()
            return $ toPacket (a,r)
        ToRouteClass toPacket _      -> do
            trace ("R-PARSE "++show typ) $ return ()
            cnt <- remaining
            (pub,a) <- isolate (cnt - 59*3) get
            r <- get
            trace ("R-PARSED "++show typ) $ return ()
            return $ toPacket (pub,(a,r))
        OnionClass      toPacket _      -> do
            trace ("ONION-PARSE "++show typ) $ return ()
            p <- get
            trace ("ONION-PARSED "++show typ) $ return ()
            return $ toPacket p
        NoncedUnclaimedClass toPacket _ -> toPacket <$> get <*> get
        Unclassified -> fail $ "todo: unserialize packet "++show typ

putPacket :: Packet -> PutM ()
putPacket p = do
    put $ pktKind p
    case pktClass (pktKind p) of
        AssymetricClass _ fromPacket      -> put $ fromPacket p
        AliasedClass    _ fromPacket      -> put $ fromPacket p
        ToRouteClass    _ fromPacket      -> put $ fromPacket p
        OnionClass      _ fromPacket      -> put $ swap $ fromPacket p
        NoncedUnclaimedClass _ fromPacket -> put $ fromPacket p -- putting a pair.
        Unclassified -> fail $ "todo: serialize packet "++show (pktKind p)