From bf8f8312b4025e04617420d23aea1b5f919e03b3 Mon Sep 17 00:00:00 2001
From: Max Lv <max.c.lv@gmail.com>
Date: Fri, 21 Jul 2017 19:54:00 +0800
Subject: [PATCH] Refine the ppbloom filter

---
 src/aead.c   |  9 ++++++---
 src/stream.c | 13 ++++++++-----
 2 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/src/aead.c b/src/aead.c
index e0f60d54..61bac8db 100644
--- a/src/aead.c
+++ b/src/aead.c
@@ -622,9 +622,6 @@ aead_decrypt(buffer_t *ciphertext, cipher_ctx_t *cipher_ctx, size_t capacity)
 
         cipher_ctx->init = 1;
 
-    } else if (cipher_ctx->init == 1) {
-        ppbloom_add((void *)cipher_ctx->salt, salt_len);
-        cipher_ctx->init = 2;
     }
 
     size_t plen = 0;
@@ -652,6 +649,12 @@ aead_decrypt(buffer_t *ciphertext, cipher_ctx_t *cipher_ctx, size_t capacity)
     memcpy(ciphertext->data, plaintext->data, plaintext->len);
     ciphertext->len = plaintext->len;
 
+    // Add the salt to bloom filter
+    if (cipher_ctx->init == 1) {
+        ppbloom_add((void *)cipher_ctx->salt, salt_len);
+        cipher_ctx->init = 2;
+    }
+
     return CRYPTO_OK;
 }
 
diff --git a/src/stream.c b/src/stream.c
index e05d28d0..b79259dc 100644
--- a/src/stream.c
+++ b/src/stream.c
@@ -553,11 +553,6 @@ stream_decrypt(buffer_t *ciphertext, cipher_ctx_t *cipher_ctx, size_t capacity)
                 return CRYPTO_ERROR;
             }
         }
-    } else if (cipher_ctx->init == 1) {
-        if (cipher->method >= RC4_MD5) {
-            ppbloom_add((void *)cipher_ctx->nonce, cipher->nonce_len);
-            cipher_ctx->init = 2;
-        }
     }
 
     if (ciphertext->len <= 0)
@@ -601,6 +596,14 @@ stream_decrypt(buffer_t *ciphertext, cipher_ctx_t *cipher_ctx, size_t capacity)
     memcpy(ciphertext->data, plaintext->data, plaintext->len);
     ciphertext->len = plaintext->len;
 
+    // Add to bloom filter
+    if (cipher_ctx->init == 1) {
+        if (cipher->method >= RC4_MD5) {
+            ppbloom_add((void *)cipher_ctx->nonce, cipher->nonce_len);
+            cipher_ctx->init = 2;
+        }
+    }
+
     return CRYPTO_OK;
 }