Merge branch 'dht-rewrite' of jerkface.net:repo/bittorrent into dht-rewrite

4 files changed, 150 insertions, 19 deletions

diff --git a/src/Crypto/Tox.hs b/src/Crypto/Tox.hs
index a25f9f4f..9f86f6a4 100644
--- a/src/Crypto/Tox.hs
+++ b/src/Crypto/Tox.hs
@@ -7,7 +7,7 @@
 {-# LANGUAGE DeriveTraversable #-}
 {-# LANGUAGE ExplicitNamespaces #-}
 {-# LANGUAGE TypeOperators #-}
-{-# LANGUAGE MagicHash, UnboxedTuples #-}
+{-# LANGUAGE MagicHash, UnboxedTuples, BangPatterns #-}
 module Crypto.Tox
     ( PublicKey
     , publicKey
@@ -35,6 +35,7 @@ module Crypto.Tox
     , Nonce8(..)
     , Nonce24(..)
     , incrementNonce24
+    , addtoNonce24
     , Nonce32(..)
     , getRemainingEncrypted
     , putEncrypted
@@ -258,17 +259,18 @@ hsalsa20 k n = BA.append a b
 newtype Nonce24 = Nonce24 ByteString
  deriving (Eq, Ord, ByteArrayAccess,Data)
 
-incrementNonce24 :: Nonce24 -> IO Nonce24
-incrementNonce24 (Nonce24 n24) = Nonce24 <$> BA.copy n24 init
+addtoNonce24 :: Nonce24 -> Word -> IO Nonce24
+addtoNonce24 (Nonce24 n24) n = Nonce24 <$> BA.copy n24 init
     where
         init :: Ptr Word -> IO ()
         init ptr | fitsInInt (Proxy :: Proxy Word64) = do
             let frmBE64 = fromIntegral . fromBE64 . fromIntegral
                 tBE64   = fromIntegral . toBE64   . fromIntegral
+                !(W# input) = n
             W# w1 <- frmBE64 <$> peek ptr
             W# w2 <- frmBE64 <$> peekElemOff ptr 1
             W# w3 <- frmBE64 <$> peekElemOff ptr 2
-            let (# overflw,  sum   #) = plusWord2# w3 (int2Word# 1#)
+            let (# overflw,  sum   #) = plusWord2# w3 input
                 (# overflw', sum'  #) = plusWord2# w2 overflw
                 (# discard,  sum'' #) = plusWord2# w1 overflw'
             poke        ptr   $ tBE64 (W# sum'')
@@ -278,13 +280,14 @@ incrementNonce24 (Nonce24 n24) = Nonce24 <$> BA.copy n24 init
         init ptr | fitsInInt (Proxy :: Proxy Word32) = do
             let frmBE32 = fromIntegral . fromBE32 . fromIntegral
                 tBE32   = fromIntegral . toBE32   . fromIntegral
+                !(W# input) = n
             W# w1 <- frmBE32 <$> peek ptr
             W# w2 <- frmBE32 <$> peekElemOff ptr 1
             W# w3 <- frmBE32 <$> peekElemOff ptr 2
             W# w4 <- frmBE32 <$> peekElemOff ptr 3
             W# w5 <- frmBE32 <$> peekElemOff ptr 4
             W# w6 <- frmBE32 <$> peekElemOff ptr 5
-            let (# overflw_,   sum_     #) = plusWord2# w6 (int2Word# 1#)
+            let (# overflw_,   sum_     #) = plusWord2# w6 input
                 (# overflw__,  sum__    #) = plusWord2# w5 overflw_
                 (# overflw___, sum___   #) = plusWord2# w6 overflw__
                 (# overflw,  sum   #) = plusWord2# w3 overflw___
@@ -298,6 +301,9 @@ incrementNonce24 (Nonce24 n24) = Nonce24 <$> BA.copy n24 init
             pokeElemOff ptr 5 $ tBE32 (W# sum_)
         init _ = error "incrementNonce24: I only support 64 and 32 bits"
 
+incrementNonce24 :: Nonce24 -> IO Nonce24
+incrementNonce24 nonce24 = addtoNonce24 nonce24 1
+
 quoted :: ShowS -> ShowS
 quoted shows s = '"':shows ('"':s)
 
diff --git a/src/Network/Tox.hs b/src/Network/Tox.hs
index 87835769..2f778874 100644
--- a/src/Network/Tox.hs
+++ b/src/Network/Tox.hs
@@ -89,7 +89,7 @@ import Crypto.Tox
 import Data.Word64Map                        (fitsInInt)
 import qualified Data.Word64Map              (empty)
 import Network.Tox.Crypto.Transport          (NetCrypto)
-import Network.Tox.Crypto.Handlers           (newSessionsState, cryptoNetHandler)
+import Network.Tox.Crypto.Handlers           (newSessionsState, cryptoNetHandler, cryptoDefaultHooks)
 import qualified Network.Tox.DHT.Handlers    as DHT
 import qualified Network.Tox.DHT.Transport   as DHT
 import Network.Tox.NodeId
@@ -309,7 +309,7 @@ newTox keydb addr = do
                     (const id)
 
     roster <- newRoster
-    sessionsState <- newSessionsState crypto
+    sessionsState <- newSessionsState crypto cryptoDefaultHooks
     return Tox
         { toxDHT = dhtclient
         , toxOnion = onionclient
diff --git a/src/Network/Tox/Crypto/Handlers.hs b/src/Network/Tox/Crypto/Handlers.hs
index 29f55e54..c5476371 100644
--- a/src/Network/Tox/Crypto/Handlers.hs
+++ b/src/Network/Tox/Crypto/Handlers.hs
@@ -12,27 +12,48 @@ import Crypto.Hash
 import Control.Applicative
 import Control.Monad
 import Data.Time.Clock.POSIX
+import qualified Data.ByteString  as B
+import Control.Lens
+import Data.Function
+import Data.Serialize             as S
+import Data.Word
+import GHC.Conc (unsafeIOToSTM)
+
+-- util, todo: move to another module
+maybeToEither (Just x) = Right x
+maybeToEither Nothing  = Left "maybeToEither"
 
 data NetCryptoSessionStatus = Unaccepted | Accepted | Confirmed
     deriving (Eq,Ord,Show,Enum)
 
 
+type IOHook addr x = addr -> x -> IO (Maybe (x -> x))
+type NetCryptoHook = IOHook SockAddr CryptoData
+
+
 data NetCryptoSession = NCrypto { ncState         :: TVar NetCryptoSessionStatus
-                                , ncTheirPacketNonce:: TVar Nonce24 -- base nonce + packet number
-                                , ncMyPacketNonce   :: TVar Nonce24 -- base nonce + packet number
+                                , ncTheirBaseNonce :: TVar Nonce24 -- base nonce + packet number
+                                , ncMyPacketNonce :: TVar Nonce24 -- base nonce + packet number
                                 , ncHandShake     :: TVar (Maybe (Handshake Encrypted))
                                 , ncCookie        :: TVar (Maybe Cookie)
                                 , ncTheirSessionPublic :: Maybe PublicKey
                                 , ncSessionSecret :: SecretKey
                                 , ncSockAddr      :: SockAddr
+                                , ncHooks         :: TVar (Map.Map MessageType [NetCryptoHook])
                                 }
 
 data NetCryptoSessions = NCSessions { netCryptoSessions :: TVar (Map.Map SockAddr NetCryptoSession)
                                     , transportCrypto :: TransportCrypto
+                                    , defaultHooks :: Map.Map MessageType [NetCryptoHook]
                                     }
 
-newSessionsState :: TransportCrypto -> IO NetCryptoSessions
-newSessionsState crypto = error "todo"
+newSessionsState :: TransportCrypto -> Map.Map MessageType [NetCryptoHook] ->  IO NetCryptoSessions
+newSessionsState crypto hooks = do
+    x <- atomically $ newTVar Map.empty
+    return NCSessions { netCryptoSessions = x
+                      , transportCrypto = crypto
+                      , defaultHooks = hooks
+                      }
 
 data HandshakeParams 
         = HParam
@@ -48,6 +69,7 @@ newHandShakeData = error "todo"
 
 cryptoNetHandler :: NetCryptoSessions -> SockAddr -> NetCrypto -> IO (Maybe (NetCrypto -> NetCrypto))
 cryptoNetHandler sessions addr (NetHandshake (Handshake (Cookie n24 ecookie) nonce24 encrypted)) = do
+    -- Handle Handshake Message
     let crypto = transportCrypto sessions
         allsessions = netCryptoSessions sessions
         anyRight xs f = foldr1 (<|>) $ map f xs
@@ -81,17 +103,18 @@ cryptoNetHandler sessions addr (NetHandshake (Handshake (Cookie n24 ecookie) non
         Left _ -> return ()
         Right hp@(HParam
                       { hpTheirBaseNonce = Just theirBaseNonce
-                      , hpOtherCookie    = Just otherCookie             
+                      , hpOtherCookie    = Just otherCookie
                       , hpTheirSessionKeyPublic = theirSessionKey
                       , hpMySecretKey        = key
                       , hpCookieRemotePubkey = remotePublicKey
                       , hpCookieRemoteDhtkey = remoteDhtPublicKey
                       }) -> do
             sessionsmap <- atomically $ readTVar allsessions
+            -- Do a lookup, in case we decide to handle the update case differently
             case Map.lookup addr sessionsmap of
-                Nothing -> do -- create new session
+                _ -> do -- create new session
                     ncState0 <- atomically $ newTVar Accepted
-                    ncTheirPacketNonce0 <- atomically $ newTVar theirBaseNonce
+                    ncTheirBaseNonce0 <- atomically $ newTVar theirBaseNonce
                     n24 <- atomically $ transportNewNonce crypto
                     let myhandshakeData = newHandShakeData crypto hp
                         plain = encodePlain myhandshakeData
@@ -105,22 +128,78 @@ cryptoNetHandler sessions addr (NetHandshake (Handshake (Cookie n24 ecookie) non
                     ncHandShake0 <- atomically $ newTVar (Just myhandshake)
                     cookie0 <- atomically $ newTVar (Just otherCookie)
                     newsession <- generateSecretKey
+                    ncHooks0 <- atomically $ newTVar (defaultHooks sessions)
                     let netCryptoSession = 
                             NCrypto { ncState         = ncState0
-                                    , ncTheirPacketNonce= ncTheirPacketNonce0
+                                    , ncTheirBaseNonce= ncTheirBaseNonce0
                                     , ncMyPacketNonce   = ncMyPacketNonce0
                                     , ncHandShake     = ncHandShake0
                                     , ncCookie        = cookie0
                                     , ncTheirSessionPublic = Just theirSessionKey
                                     , ncSessionSecret = newsession
                                     , ncSockAddr      = addr
+                                    , ncHooks         = ncHooks0
                                     }
                     atomically $ modifyTVar allsessions (Map.insert addr netCryptoSession)
-                Just netCryptoSession -> return () -- TODO: UPdate existing session
     return Nothing 
 cryptoNetHandler sessions addr (NetCrypto (CryptoPacket nonce16 encrypted)) = do
     let crypto = transportCrypto sessions
+        allsessions = netCryptoSessions sessions
+    sessionsmap <- atomically $ readTVar allsessions
     -- Handle Encrypted Message
-    -- TODO
-    return Nothing
--- cryptoNetHandler _ _ _ = return $ Just id
+    case Map.lookup addr sessionsmap of
+        Nothing -> return Nothing -- drop packet, we have no session
+        Just (NCrypto {ncState, ncHooks,ncSessionSecret,ncTheirSessionPublic,ncTheirBaseNonce}) -> do
+            theirBaseNonce <- atomically $ readTVar ncTheirBaseNonce
+            -- Try to decrypt message
+            let diff :: Word16
+                diff = nonce16 - fromIntegral (last2Bytes theirBaseNonce) -- truncating to Word16
+            tempNonce <- addtoNonce24 theirBaseNonce (fromIntegral diff) -- expanding to Word
+            let lr = do -- Either Monad --
+                        pubkey <- maybeToEither ncTheirSessionPublic
+                        decodePlain =<< decrypt (computeSharedSecret ncSessionSecret pubkey tempNonce) encrypted
+            case lr of
+                Left _ -> return Nothing -- decryption failed, ignore packet
+                Right cd@(CryptoData {bufferStart, bufferEnd, bufferData=cm}) -> do -- decryption succeeded,
+                    -- TODO: Why do I need bufferStart & bufferEnd?
+                    --
+                    --   buffer_start = highest packet number handled + 1
+                    --                , recvbuffers buffer_start
+                    --
+                    --   bufferEnd = sendbuffer buffer_end if lossy, otherwise packet number
+                    -- update ncTheirBaseNonce if necessary
+                    when (diff > 2 * dATA_NUM_THRESHOLD)$
+                        atomically $ do
+                            y <- readTVar ncTheirBaseNonce
+                            -- all because Storable forces IO...
+                            x <- unsafeIOToSTM $ addtoNonce24 y (fromIntegral dATA_NUM_THRESHOLD)
+                            writeTVar ncTheirBaseNonce y
+                    -- then set session confirmed,
+                    atomically $ writeTVar ncState Confirmed
+                    hookmap <- atomically $ readTVar ncHooks
+                    -- run hook
+                    flip fix cd $ \lookupAgain cd -> do
+                        let msgTyp = cd ^. messageType
+                        case Map.lookup msgTyp hookmap of
+                                Nothing -> return Nothing -- discarding, because no hooks
+                                Just hooks -> flip fix (hooks,cd,msgTyp) $ \loop (hooks,cd,typ) -> do
+                                                let _ = cd :: CryptoData
+                                                case (hooks,cd) of
+                                                    ([],_) ->  return Nothing
+                                                    (hook:more,cd) -> do
+                                                        r <- hook addr cd :: IO (Maybe (CryptoData -> CryptoData))
+                                                        case r of
+                                                            Just f -> let newcd = f cd
+                                                                          newtyp = newcd ^. messageType
+                                                                          in if newtyp == typ then loop (more,newcd,newtyp)
+                                                                                              else lookupAgain newcd
+                                                            Nothing -> return Nothing -- message consumed
+    where
+        last2Bytes :: Nonce24 -> Word
+        last2Bytes (Nonce24 bs) = case S.decode (B.drop 22 bs) of
+                                    Right n -> n
+                                    _ -> error "unreachable-last2Bytes"
+        dATA_NUM_THRESHOLD = 21845 -- = 65535 / 3 
+
+
+cryptoDefaultHooks = Map.empty
diff --git a/src/Network/Tox/Crypto/Transport.hs b/src/Network/Tox/Crypto/Transport.hs
index 634a7a43..8739c853 100644
--- a/src/Network/Tox/Crypto/Transport.hs
+++ b/src/Network/Tox/Crypto/Transport.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE ViewPatterns #-}
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE LambdaCase #-}
@@ -43,6 +44,8 @@ module Network.Tox.Crypto.Transport
     , sizedN
     , sizedAtLeastN
     , isIndirectGrpChat
+    , LossyOrLossless(..)
+    , lossyness
     ) where
 
 import Crypto.Tox
@@ -148,6 +151,10 @@ data CryptoData = CryptoData
     , bufferData :: CryptoMessage
     }
 
+instance Serialize CryptoData where
+    get = CryptoData <$> get <*> get <*> get
+    put (CryptoData start end dta) = put start >> put end >> put dta
+
 -- The 'UserStatus' equivalent in Presence is:
 --
 -- data JabberShow = Offline
@@ -176,6 +183,25 @@ instance Sized CryptoMessage where
         TwoByte {}              -> 2
         UpToN { msgBytes = bs } -> 1 + B.length bs
 
+instance Serialize CryptoMessage where
+    get = do
+        i <- get :: Get MessageID
+        n <- remaining
+        case msgSizeParam i of
+            Just (True,1) -> return $ OneByte i
+            Just (True,2) -> TwoByte i <$> get
+            _ -> UpToN i <$> getByteString n
+
+    put (OneByte i)   =    putWord8 (fromIntegral . fromEnum $ i)
+    put (TwoByte i b) = do putWord8 (fromIntegral . fromEnum $ i)
+                           putWord8 b 
+    put (UpToN i x)   = do putWord8 (fromIntegral . fromEnum $ i)
+                           putByteString x
+
+instance Serialize MessageID where
+    get = toEnum . fromIntegral <$> getWord8
+    put x = putWord8 (fromIntegral . fromEnum $ x)
+    
 erCompat :: String -> a
 erCompat lens = error $ "Use of '" ++ lens ++ "'  lens on incompatible CryptoMessage type"
 
@@ -381,6 +407,12 @@ data MessageType = Msg MessageID
                  | GrpMsg MessageName
     deriving (Eq,Show)
 
+instance Ord MessageType where
+    compare (Msg x) (Msg y) = compare x y
+    compare (GrpMsg x) (GrpMsg y) = compare x y
+    compare (Msg _) (GrpMsg _) = LT
+    compare (GrpMsg _) (Msg _) = GT
+
 class HasMessageType x where
     getMessageType :: x -> MessageType
     setMessageType :: x -> MessageType -> x
@@ -403,6 +435,10 @@ instance HasMessageType CryptoMessage where
     setMessageType (TwoByte mid0 x) (Msg mid) = UpToN mid (B.singleton x)
     setMessageType (UpToN mid0 x)   (Msg mid) = UpToN mid x
 
+instance HasMessageType CryptoData where
+    getMessageType (CryptoData { bufferData })     = getMessageType bufferData
+    setMessageType cd@(CryptoData { bufferData=bd }) typ = cd { bufferData=setMessageType bd typ }
+
 -- | This lens should always succeed on CryptoMessage
 messageType :: (Functor f, HasMessageType x) => (MessageType -> f MessageType) -> (x -> f x)
 messageType = lens getMessageType setMessageType
@@ -564,6 +600,16 @@ isIndirectGrpChat MESSAGE_GROUPCHAT = True
 isIndirectGrpChat LOSSY_GROUPCHAT = True
 isIndirectGrpChat _ = False
 
+data LossyOrLossless = UnknownLossyness | Lossless | Lossy 
+    deriving (Eq,Ord,Enum,Show,Bounded)
+
+lossyness :: MessageID -> LossyOrLossless
+lossyness (fromEnum -> x) | x < 3             = Lossy
+lossyness (fromEnum -> x) | x >= 16, x < 192  = Lossless
+lossyness (fromEnum -> x) | x >= 192, x < 255 = Lossy
+lossyness (fromEnum -> 255)                   = Lossless
+lossyness _ = UnknownLossyness
+
 -- TODO: Flesh this out.
 data MessageID              -- First byte indicates data
     = Padding               -- ^ 0 padding (skipped until we hit a non zero (data id) byte)