1 files changed, 14 insertions, 13 deletions

diff --git a/receiver.cpp b/receiver.cpp
index 9e622d9..a3a0dc9 100644
--- a/receiver.cpp
+++ b/receiver.cpp
@@ -131,6 +131,13 @@ public:
         const uint8_t *payload = packet + GCM_IV_LENGTH + GCM_TAG_LENGTH;
         len -= GCM_IV_LENGTH + GCM_TAG_LENGTH;
 
+        std::string ivs(packet, packet +GCM_IV_LENGTH);
+        if (_used_ivs.find(ivs) != _used_ivs.end()) {
+            std::cerr << "Error: Session " << std::hex << _session_id << " reused IV. Dropping packet" << std::endl;
+            return;
+        }
+        _used_ivs.insert(ivs);
+
         // Create output file if it doesn't exist
         if (_fd < 0)
             _fd = ::open(_filename.c_str(), O_WRONLY | O_APPEND | O_CREAT, 0755);
@@ -168,16 +175,16 @@ public:
     }
 
 private:
-    uint64_t            _session_id;
-    uint8_t             _key[AES_KEY_LENGTH];
-    int                 _fd = -1;
-    time_t              _last_access = 0;
-    std::string         _filename;
-    mbedtls_gcm_context _ctx;
+    uint64_t                _session_id;
+    uint8_t                 _key[AES_KEY_LENGTH];
+    int                     _fd = -1;
+    time_t                  _last_access = 0;
+    std::string             _filename;
+    std::set<std::string>   _used_ivs;
+    mbedtls_gcm_context     _ctx;
 };
 
 std::map<uint64_t, std::unique_ptr<Session>> g_sessions;
-std::set<std::string> g_used_ivs;
 
 static uint8_t hex2nyble(char c)
 {
@@ -286,12 +293,6 @@ int main() {
                 g_sessions[session_id] = std::make_unique<Session>(session_id, rsa_plain_text);
             break;
         case 1: {
-            std::string sessid_iv(packet + 1, packet + 1 + SESSION_ID_LENGTH + GCM_IV_LENGTH);
-            if (g_used_ivs.find(sessid_iv) != g_used_ivs.end()) {
-                std::cerr << "Error: Session " << std::hex << session_id << " reused IV. Dropping packet" << std::endl;
-                break;
-            }
-            g_used_ivs.insert(sessid_iv);
             if (session != g_sessions.end())
                 session->second->write_log(packet + 1 + SESSION_ID_LENGTH, len - 1 - SESSION_ID_LENGTH);
             else