2 files changed, 26 insertions, 13 deletions

diff --git a/dht-client.cabal b/dht-client.cabal
index 02b5cdbe..9e3c78ae 100644
--- a/dht-client.cabal
+++ b/dht-client.cabal
@@ -73,6 +73,7 @@ library
                     Data.Wrapper.PSQ
                     Data.Wrapper.PSQInt
                     Data.MinMaxPSQ
+                    Data.InOrOut
                     Network.Address
                     Network.Kademlia.Bootstrap
                     Network.Kademlia.Routing
diff --git a/src/Network/Tox/Crypto/Handlers.hs b/src/Network/Tox/Crypto/Handlers.hs
index 0ad868a7..b227e5ee 100644
--- a/src/Network/Tox/Crypto/Handlers.hs
+++ b/src/Network/Tox/Crypto/Handlers.hs
@@ -753,7 +753,10 @@ sendOffline crypto session = sendCrypto crypto session (return ()) (OneByte OFFL
 
 
 sendKill :: TransportCrypto -> NetCryptoSession -> IO (Either String ())
-sendKill crypto session = sendCrypto crypto session (return ()) (OneByte KillPacket)
+sendKill crypto session = do
+    let cm=OneByte KillPacket
+    addMsgToLastN False (cm ^. messageType) session cm
+    sendCrypto crypto session (return ()) cm
 
 setNick :: TransportCrypto -> NetCryptoSession -> ByteString -> IO (Either String ())
 setNick crypto session nick = do
@@ -766,6 +769,7 @@ setNick crypto session nick = do
                 view <- readTVar viewVar
                 writeTVar (svNick view) nick
         let cm = UpToN NICKNAME nick
+        addMsgToLastN False (cm ^. messageType) session cm
         sendCrypto crypto session updateLocal cm
 
 setTyping :: TransportCrypto -> NetCryptoSession -> TypingStatus -> IO (Either String ())
@@ -774,6 +778,7 @@ setTyping crypto session status = do
         view <- readTVar (ncView session)
         writeTVar (svTyping view) status
    let cm = TwoByte TYPING (fromEnum8 status)
+   addMsgToLastN False (cm ^. messageType) session cm
    sendCrypto crypto session updateLocal cm
 
 setNoSpam :: TransportCrypto -> NetCryptoSession -> Maybe NoSpam -> IO (Either String ())
@@ -790,6 +795,7 @@ setStatus crypto session status = do
         view <- readTVar (ncView session)
         writeTVar (svStatus view) status
    let cm = TwoByte USERSTATUS (fromEnum8 status)
+   addMsgToLastN False (cm ^. messageType) session cm
    sendCrypto crypto session updateLocal cm
 
 setStatusMsg :: TransportCrypto -> NetCryptoSession -> ByteString -> IO (Either String ())
@@ -802,6 +808,7 @@ setStatusMsg crypto session msg = do
                 view <- readTVar (ncView session)
                 writeTVar (svStatusMsg view) msg
         let cm = UpToN STATUSMESSAGE msg
+        addMsgToLastN False (cm ^. messageType) session cm
         sendCrypto crypto session updateLocal cm
 
 -- | handles nothings
@@ -859,18 +866,7 @@ hookHelper _ typ session cm | any ($ typ) [isKillPacket, isOFFLINE] = atomically
     return Nothing
 
 hookHelper handledFlg typ  session  cm = do
-    let msgQ = ncLastNMsgs session
-        msgNumVar = ncMsgNumVar session
-        dropCntVar = ncDropCntVar session
-    atomically $ do
-        num <- readTVar msgNumVar
-        view <- readTVar (ncView session)
-        snapshot <- viewSnapshot view
-        (wraps,offset) <- PQ.enqueue msgQ num (handledFlg,(snapshot,Out cm))
-        capacity <- PQ.getCapacity msgQ
-        let dropped = wraps * capacity + offset
-        modifyTVar' msgNumVar (+1)
-        writeTVar dropCntVar dropped
+    addMsgToLastN handledFlg typ session cm
     atomically $ do
         idtmchans <- IntMap.assocs <$> readTVar (ncListeners session)
         mbChans
@@ -885,6 +881,22 @@ hookHelper handledFlg typ  session  cm = do
             writeTMChan chan cm
     return Nothing
 
+addMsgToLastN :: Bool -> MessageType -> NetCryptoSession -> CryptoMessage -> IO ()
+addMsgToLastN handledFlg typ session cm = do
+    let lastNQ = ncLastNMsgs session
+        msgNumVar = ncMsgNumVar session
+        dropCntVar = ncDropCntVar session
+    atomically $ do
+        num <- readTVar msgNumVar
+        view <- readTVar (ncView session)
+        snapshot <- viewSnapshot view
+        (wraps,offset) <- PQ.enqueue lastNQ num (handledFlg,(snapshot,Out cm))
+        capacity <- PQ.getCapacity lastNQ
+        let dropped = wraps * capacity + offset
+        modifyTVar' msgNumVar (+1)
+        writeTVar dropCntVar dropped
+
+
 -- | use to add a single hook to a specific session.
 addCryptoDataHook1 :: Map.Map MessageType [NetCryptoHook] -> MessageType -> NetCryptoHook -> Map.Map MessageType [NetCryptoHook]
 addCryptoDataHook1 mp typ hook = case Map.lookup typ mp of