4 files changed, 195 insertions, 8 deletions

diff --git a/src/Crypto/Tox.hs b/src/Crypto/Tox.hs
index 18cdb5d2..624da233 100644
--- a/src/Crypto/Tox.hs
+++ b/src/Crypto/Tox.hs
@@ -34,11 +34,13 @@ module Crypto.Tox
     , decodePlain
     -- , computeSharedSecret
     , lookupSharedSecret
+    , lookupNonceFunction
     , encrypt
     , decrypt
     , Nonce8(..)
     , Nonce24(..)
     , incrementNonce24
+    , nonce24ToWord16
     , addtoNonce24
     , Nonce32(..)
     , getRemainingEncrypted
@@ -300,10 +302,19 @@ lookupSharedSecret :: TransportCrypto -> SecretKey -> PublicKey -> Nonce24 -> IO
 lookupSharedSecret crypto sk recipient nonce
     = ($ nonce) <$> lookupNonceFunction crypto sk recipient
 
+{-# INLINE lookupNonceFunction #-}
 lookupNonceFunction :: TransportCrypto -> SecretKey -> PublicKey -> IO (Nonce24 -> State)
-lookupNonceFunction TransportCrypto{secretsCache} sk recipient = do
+lookupNonceFunction c@(TransportCrypto{secretsCache}) sk recipient = do
     now <- getPOSIXTime
-    atomically $ do
+    atomically $ lookupNonceFunctionSTM now c sk recipient
+
+{-# INLINE lookupNonceFunctionSTM #-}
+-- | This version of 'lookupNonceFunction' is STM instead of IO, this means if some later part of
+--   of the transaction fails, we may end up forgoing a computation that could have been cached.
+--   Use with care. In most circumstances you probably want 'lookupNonceFunction'. It also commits
+--   us to using TVars to store the cache.
+lookupNonceFunctionSTM :: POSIXTime -> TransportCrypto -> SecretKey -> PublicKey -> STM (Nonce24 -> State)
+lookupNonceFunctionSTM now TransportCrypto{secretsCache} sk recipient = do
     mm <- readTVar $ sharedSecret secretsCache
     case MM.lookup' recipient mm of
         Nothing -> do
@@ -332,7 +343,10 @@ hsalsa20 k n = BA.append a b
 
 
 newtype Nonce24 = Nonce24 ByteString
- deriving (Eq, Ord, ByteArrayAccess,Data)
+ deriving (Eq, Ord, ByteArrayAccess, Data)
+
+nonce24ToWord16 :: Nonce24 -> Word16
+nonce24ToWord16 (Nonce24 n24) = fromIntegral (B.index n24 23) + 256 * fromIntegral (B.index n24 22)
 
 addtoNonce24 :: Nonce24 -> Word -> IO Nonce24
 addtoNonce24 (Nonce24 n24) n = Nonce24 <$> BA.copy n24 init
@@ -376,6 +390,7 @@ addtoNonce24 (Nonce24 n24) n = Nonce24 <$> BA.copy n24 init
             pokeElemOff ptr 5 $ tBE32 (W# sum_)
         init _ = error "incrementNonce24: I only support 64 and 32 bits"
 
+{-# INLINE incrementNonce24 #-}
 incrementNonce24 :: Nonce24 -> IO Nonce24
 incrementNonce24 nonce24 = addtoNonce24 nonce24 1
 
diff --git a/src/Data/PacketQueue.hs b/src/Data/PacketQueue.hs
index b349bf4b..927d6c53 100644
--- a/src/Data/PacketQueue.hs
+++ b/src/Data/PacketQueue.hs
@@ -3,12 +3,19 @@
 -- be stored out of order, but from which they are extracted in the proper
 -- sequence.
 {-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE FlexibleContexts #-}
 module Data.PacketQueue
     ( PacketQueue
     , new
     , dequeue
     , enqueue
     , observeOutOfBand
+    , PacketOutQueue
+    , newOutGoing
+    , tryAppendQueueOutgoing
+    , dequeueOutgoing
+    , mapOutGoing
+    , OutGoingResult(..)
     ) where
 
 import Control.Concurrent.STM
@@ -21,7 +28,7 @@ data PacketQueue a = PacketQueue
     { pktq    :: TArray Word32 (Maybe a)
     , seqno   :: TVar Word32
     , qsize   :: Word32
-    , buffend :: TVar Word32
+    , buffend :: TVar Word32 -- on incoming, highest packet number handled + 1
     }
 
 -- | Create a new PacketQueue.
@@ -77,3 +84,87 @@ enqueue PacketQueue{ pktq, seqno, qsize, buffend } no x = do
 
 -- lookup :: PacketQueue a -> Word32 -> STM (Maybe a)
 -- lookup PacketQueue{ pktq, seqno, qsize } no = _todo
+
+-----------------------------------------------------
+-- * PacketOutQueue
+--
+
+data PacketOutQueue extra msg toWire fromWire = PacketOutQueue
+    { pktoInPQ   :: PacketQueue fromWire -- ^ reference to the incoming 'PacketQueue'
+    , pktoOutPQ  :: PacketQueue (Word32,toWire)
+    , pktoPacketNo :: TVar Word32
+    , pktoToWireIO :: IO (STM extra)
+    , pktoToWire :: STM extra
+                 -> Word32{-packet number we expect to recieve-}
+                 -> Word32{- buffer_end -}
+                 -> Word32{- packet number -}
+                 -> msg
+                 -> STM (Maybe (toWire,Word32{-next packet no-}))
+    }
+
+mapOutGoing :: ((Word32,towire) -> Maybe (Word32,towire)) -> PacketOutQueue extra msg towire fromwire -> STM ()
+mapOutGoing f q@(PacketOutQueue { pktoOutPQ=PacketQueue{ pktq } }) = do
+    (z,n) <- getBounds pktq
+    let ff i = do
+        e <- readArray pktq i
+        writeArray pktq i (e>>=f)
+    mapM_ ff [z .. n]
+
+newOutGoing :: PacketQueue fromwire
+                      -- ^ Incoming queue
+            -> (STM io -> Word32 {-packet number we expect to recieve-} -> Word32{-buffer_end-} -> Word32{-packet number-} -> msg -> STM (Maybe (wire,Word32{-next packet no-})))
+                      -- ^ toWire callback
+            -> IO (STM io)
+                      -- ^ io action to get extra parameter
+            -> Word32 -- ^ packet number of first outgoing packet
+            -> Word32 -- ^ Capacity of queue.
+            -> Word32 -- ^ Initial sequence number.
+            -> STM (PacketOutQueue io msg wire fromwire)
+newOutGoing inq towire toWireIO num capacity seqstart = do
+    outq <- new capacity seqstart
+    numVar <- newTVar num
+    return $ PacketOutQueue
+                { pktoInPQ = inq
+                , pktoOutPQ = outq
+                , pktoPacketNo = numVar
+                , pktoToWireIO = toWireIO
+                , pktoToWire = towire
+                }
+
+data OutGoingResult = OGSuccess | OGFull | OGEncodeFail
+    deriving (Eq,Show)
+
+-- | Convert a message to packet format and append it to the front of a queue
+--   used for outgoing messages. (Note that ‘front‛ usually means the higher
+--   index in this implementation.)
+tryAppendQueueOutgoing :: STM extra -> PacketOutQueue extra msg wire fromwire -> msg -> STM OutGoingResult
+tryAppendQueueOutgoing getExtra q@(PacketOutQueue { pktoInPQ, pktoOutPQ, pktoPacketNo, pktoToWireIO, pktoToWire }) msg = do
+    be <- readTVar (buffend pktoOutPQ)
+    let i = be `mod` (qsize pktoOutPQ)
+    mbPkt <- readArray (pktq pktoOutPQ) i
+    pktno <- readTVar pktoPacketNo
+    nextno <- readTVar (seqno pktoInPQ)
+    mbWire <- pktoToWire getExtra nextno be pktno msg
+    case mbWire of
+        Just (pkt,pktno')
+            -> case mbPkt of
+                    -- slot is free, insert element
+                    Nothing -> do
+                        modifyTVar' (buffend pktoOutPQ) (+1)
+                        writeTVar pktoPacketNo $! pktno'
+                        writeArray (pktq pktoOutPQ) i (Just (pktno,pkt))
+                        return OGSuccess
+                    -- queue is full, block until its not
+                    _       -> return OGFull
+        -- don't know how to send this message
+        Nothing -> return OGEncodeFail
+
+dequeueOutgoing :: PacketOutQueue extra msg wire fromwire -> STM (Word32,wire)
+dequeueOutgoing (PacketOutQueue {pktoOutPQ=PacketQueue { pktq, seqno, qsize }}) = do
+    i0 <- readTVar seqno
+    let i = i0 `mod` qsize
+    x <- maybe retry return =<< readArray pktq i
+    -- writeArray pktq i Nothing -- not cleaning
+    modifyTVar' seqno   succ
+    return x
+
diff --git a/src/Network/Tox/Crypto/Handlers.hs b/src/Network/Tox/Crypto/Handlers.hs
index edfcb260..6a79da1b 100644
--- a/src/Network/Tox/Crypto/Handlers.hs
+++ b/src/Network/Tox/Crypto/Handlers.hs
@@ -1,5 +1,6 @@
 {-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE TupleSections #-}
+{-# LANGUAGE TypeOperators #-}
 module Network.Tox.Crypto.Handlers where
 
 import Network.Tox.NodeId
@@ -22,6 +23,7 @@ import qualified Data.PacketQueue as PQ
          ;import Data.PacketQueue (PacketQueue)
 import Data.Serialize             as S
 import Data.Word
+import qualified Data.Word64Map as W64
 import GHC.Conc (unsafeIOToSTM)
 import qualified Data.Set         as Set
 import qualified Data.Array.Unboxed as A
@@ -32,6 +34,7 @@ import System.Environment
 import System.Directory
 import Control.Concurrent
 import GHC.Conc (labelThread)
+import System.IO.Unsafe(unsafeDupablePerformIO {- unsafeIOToSTM -})
 
 -- util, todo: move to another module
 maybeToEither :: Maybe b -> Either String b
@@ -45,7 +48,10 @@ data NetCryptoSessionStatus = Unaccepted | Accepted | Confirmed
 type IOHook addr x = addr -> x -> IO (Maybe (x -> x))
 type NetCryptoHook = IOHook NetCryptoSession CryptoData
 type MsgTypeArray = A.UArray Word8 Word16
-
+type MsgOutMap = W64.Word64Map Word8
+-- type MsgOutMap = A.UArray Word64 Word8 -- if above is too slow, switch to this, but use reasonable bounds
+msgOutMapLookup :: Word64 -> MsgOutMap -> Maybe Word8
+msgOutMapLookup = W64.lookup
 
 -- | Information, that may be made visible in multiple sessions, as well
 -- as displayed in some way to the user via mutiple views.
@@ -86,7 +92,7 @@ data NetCryptoSession = NCrypto { ncState          :: TVar NetCryptoSessionStatu
                                                                       --   may not be in ncHooks yet, but they should appear
                                                                       --   here if ncUnrecognizedHook will add them to ncHooks
                                                                       --   on an as-need basis.
-                                , ncOutgoingTypeArray :: TVar MsgTypeArray
+                                , ncOutgoingIdMap :: TVar MsgOutMap
                                 , ncAllSessions :: NetCryptoSessions -- ^ may be needed if one net-crypto session
                                                                      --   needs to possibly start another, as is
                                                                      --   the case in group chats
@@ -94,6 +100,7 @@ data NetCryptoSession = NCrypto { ncState          :: TVar NetCryptoSessionStatu
                                 , ncPacketQueue :: PacketQueue CryptoData
                                 , ncBufferStart :: TVar Word32
                                 , ncDequeueThread :: Maybe ThreadId
+                                , ncOutgoingQueue :: PQ.PacketOutQueue (State,Nonce24,TVar MsgOutMap) CryptoMessage (CryptoPacket Encrypted) CryptoData
                                 }
 
 data NetCryptoSessions = NCSessions { netCryptoSessions :: TVar (Map.Map SockAddr NetCryptoSession)
@@ -104,6 +111,7 @@ data NetCryptoSessions = NCSessions { netCryptoSessions :: TVar (Map.Map SockAdd
                                     , sessionView :: SessionView
                                     , msgTypeArray :: MsgTypeArray
                                     , inboundQueueCapacity :: Word32
+                                    , outboundQueueCapacity :: Word32
                                     , nextSessionId :: TVar SessionID
                                     }
 
@@ -168,6 +176,7 @@ newSessionsState crypto unrechook hooks = do
                                                   }
                       , msgTypeArray = allMsgTypes -- todo make this a parameter
                       , inboundQueueCapacity = 200
+                      , outboundQueueCapacity = 400
                       , nextSessionId = nextSessionId0
                       }
 
@@ -208,6 +217,53 @@ newHandShakeData crypto basenonce (HParam {hpOtherCookie,hpMySecretKey,hpCookieR
                 , otherCookie = freshCookie'
                 }) freshCookie
 
+type XMessage = CryptoMessage -- todo
+
+ncToWire :: STM (State,Nonce24,TVar MsgOutMap)
+         -> Word32{- packet number we expect to recieve -}
+         -> Word32{- buffer_end -}
+         -> Word32{- packet number -}
+         -> XMessage
+         -> STM (Maybe (CryptoPacket Encrypted,Word32{-next packet no-}))
+ncToWire getState seqno bufend pktno msg = do
+    let typ = getMessageType msg
+        typ64 = toWord64 typ
+    let lsness msg =
+            case typ of
+                    Msg mid -> lossyness mid
+                    GrpMsg KnownLossy _ -> Lossy
+                    GrpMsg KnownLossless _ -> Lossless
+    (state,n24,msgOutMapVar) <- getState
+    msgOutMap <- readTVar msgOutMapVar
+    case msgOutMapLookup typ64 msgOutMap of
+        Just outid -> do
+            let setMessageId (OneByte _)   mid = OneByte (toEnum8 mid)
+                setMessageId (TwoByte _ x) mid = TwoByte (toEnum8 mid) x
+                setMessageId (UpToN _ x)   mid = UpToN (toEnum8 mid) x
+                msg' = setMessageId msg outid
+                in case lsness msg of
+                        UnknownLossyness -> return Nothing
+                        Lossy    -> let cd =
+                                         CryptoData
+                                            { bufferStart = seqno
+                                            , bufferEnd = bufend
+                                            , bufferData = msg'
+                                            }
+                                        plain = encodePlain cd
+                                        encrypted = encrypt state plain
+                                        pkt = CryptoPacket { pktNonce = nonce24ToWord16 n24, pktData = encrypted }
+                                        in return (Just (pkt, pktno))
+                        Lossless -> let cd =
+                                         CryptoData
+                                            { bufferStart = seqno
+                                            , bufferEnd = pktno
+                                            , bufferData = msg'
+                                            }
+                                        plain = encodePlain cd
+                                        encrypted = encrypt state plain
+                                        pkt = CryptoPacket { pktNonce = nonce24ToWord16 n24, pktData = encrypted }
+                                        in return (Just (pkt, pktno+1))
+
 -- | called when we recieve a crypto handshake with valid cookie
 -- TODO set priority on contact addr to 0 if it is older than ForgetPeriod,
 -- then increment it regardless. (Keep addr in MinMaxPSQ in Roster.Contact)
@@ -250,10 +306,18 @@ freshCryptoSession  sessions
     ncHooks0 <- atomically $ newTVar (defaultHooks sessions)
     ncUnrecognizedHook0 <- atomically $ newTVar (defaultUnrecognizedHook sessions)
     ncIncomingTypeArray0 <- atomically $ newTVar (msgTypeArray sessions)
-    ncOutgoingTypeArray0 <- atomically $ newTVar allMsgTypes
+    ncOutgoingIdMap0 <- atomically $ newTVar W64.empty
     ncView0 <- atomically $ newTVar (sessionView sessions)
     pktq <- atomically $ PQ.new (inboundQueueCapacity sessions) 0
     bufstart <- atomically $ newTVar 0
+    let toWireIO = do
+            f <- lookupNonceFunction crypto newsession theirSessionKey
+            atomically $ do
+                n24 <- readTVar ncMyPacketNonce0
+                let n24plus1 = unsafeDupablePerformIO (incrementNonce24 n24)
+                writeTVar ncMyPacketNonce0 n24plus1
+                return (return (f n24, n24, ncOutgoingIdMap0))
+    pktoq <- atomically $ PQ.newOutGoing pktq ncToWire toWireIO 0 (outboundQueueCapacity sessions) 0
     let netCryptoSession0 =
             NCrypto { ncState         = ncState0
                     , ncSessionId     = sessionId
@@ -270,11 +334,12 @@ freshCryptoSession  sessions
                     , ncUnrecognizedHook = ncUnrecognizedHook0
                     , ncAllSessions = sessions
                     , ncIncomingTypeArray = ncIncomingTypeArray0
-                    , ncOutgoingTypeArray = ncOutgoingTypeArray0
+                    , ncOutgoingIdMap     = ncOutgoingIdMap0
                     , ncView = ncView0
                     , ncPacketQueue        = pktq
                     , ncBufferStart        = bufstart
                     , ncDequeueThread      = Nothing -- error "you want the NetCrypto-Dequeue thread id, but is it started?"
+                    , ncOutgoingQueue      = pktoq
                     }
     threadid <- forkIO $ do
         tid <- myThreadId
diff --git a/src/Network/Tox/Crypto/Transport.hs b/src/Network/Tox/Crypto/Transport.hs
index 386c2766..b89bde46 100644
--- a/src/Network/Tox/Crypto/Transport.hs
+++ b/src/Network/Tox/Crypto/Transport.hs
@@ -23,6 +23,7 @@ module Network.Tox.Crypto.Transport
     , MessageType(..)
     , KnownLossyness(..)
     , AsWord16(..)
+    , AsWord64(..)
     -- feild name classes
     , HasGroupChatID(..)
     , HasGroupNumber(..)
@@ -420,6 +421,11 @@ class AsWord16 a where
    toWord16 :: a -> Word16
    fromWord16 :: Word16 -> a
 
+class AsWord64 a where
+   toWord64 :: a -> Word64
+   fromWord64 :: Word64 -> a
+
+
 toEnum8 :: (Enum a, Integral word8) => word8 -> a
 toEnum8 = toEnum . fromIntegral
 fromEnum8 :: Enum a => a -> Word8
@@ -428,6 +434,9 @@ fromEnum8 = fromIntegral . fromEnum
 fromEnum16 :: Enum a => a -> Word16
 fromEnum16 = fromIntegral . fromEnum
 
+fromEnum64 :: Enum a => a -> Word64
+fromEnum64 = fromIntegral . fromEnum
+
 
 -- MessageType, for our client keep it inside 16 bits
 -- but we should extend it to 32 or even 64 on the wire.
@@ -440,6 +449,13 @@ instance AsWord16 MessageType where
     fromWord16 x | x < 1024, x .|. 0x0200 == 0x0200 = GrpMsg (toEnum8 ((x - 512) `div` 256)) (toEnum8 x)
     fromWord16 x = error "Not clear how to convert Word16 to MessageType"
 
+instance AsWord64 MessageType where
+    toWord64 (Msg msgID) = fromIntegral (fromIntegral (fromEnum16 msgID) :: Word8)
+    toWord64 (GrpMsg lsy msgName) = 512 + 256 * (fromEnum64 lsy) + fromIntegral (fromEnum8 msgName)
+    fromWord64 x | x < 256                          = Msg (toEnum $ fromIntegral x)
+    fromWord64 x | x < 1024, x .|. 0x0200 == 0x0200 = GrpMsg (toEnum8 ((x - 512) `div` 256)) (toEnum8 x)
+    fromWord64 x = error "Not clear how to convert Word64 to MessageType"
+
 word16 :: (Functor f, AsWord16 x) => (Word16 -> f Word16) -> (x -> f x)
 word16  = lens toWord16 (\_ x -> fromWord16 x)