1 files changed, 91 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