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

2 files changed, 137 insertions, 12 deletions

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)