From cafc51e689948e9bd0755a014a1eeb51bd848fa6 Mon Sep 17 00:00:00 2001
From: zhang_haisong <zhang_haisong@itrus.com.cn>
Date: Tue, 6 Aug 2024 11:18:57 +0800
Subject: [PATCH] Bugfix: Internal state error during multi-packet decryption

---
 examples/ascon_aead_enc.c | 39 ++++++++++++++++++++++++++++++++++++++-
 src/ascon.c               | 15 ++++++++++++---
 2 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/examples/ascon_aead_enc.c b/examples/ascon_aead_enc.c
index 642dfaf..076a140 100644
--- a/examples/ascon_aead_enc.c
+++ b/examples/ascon_aead_enc.c
@@ -17,13 +17,16 @@ int main(void)
 {
     int ret = 1;
     void *ctx = NULL;
-    const char *plaintext = "hello world";
+    const char *plaintext = "hello world, 12345678";
     const char *ad = "0123456789abcdef";
     unsigned char *key = tsm_hex2buf("0123456789abcdef0123456789abcdef");
     unsigned char *iv = tsm_hex2buf("0123456789abcdef0123456789abcdef");
     unsigned char out[1024];
     unsigned char tag[TSM_ASCON_AEAD_TAG_LEN];
     size_t outl, tmplen;
+    void *ctx_dec = NULL;
+    unsigned char dec_out[1024];
+    size_t dec_outl = 0;
 
     if (key == NULL || iv == NULL) {
         goto err;
@@ -64,9 +67,43 @@ int main(void)
 
     printf("\n");
 
+    ctx_dec = tsm_ascon_aead_ctx_new();
+    if (ctx == NULL) {
+        goto err;
+    }
+
+    if (tsm_ascon_aead_init(ctx_dec, TSM_ASCON_AEAD_128, key, iv, TSM_CIPH_FLAG_DECRYPT) != TSM_OK
+        || tsm_ascon_aead_set_tag(ctx_dec, tag)
+        || tsm_ascon_aead_update(ctx_dec, (const unsigned char *)ad, strlen(ad), NULL, NULL) != TSM_OK
+        || tsm_ascon_aead_update(ctx_dec,
+                                 out,
+                                 3,
+                                 dec_out + dec_outl,
+                                 &tmplen)
+               != TSM_OK
+        || (dec_outl += tmplen, 0)
+        || tsm_ascon_aead_update(ctx_dec,
+                                 out + 3,
+                                 outl - 3,
+                                 dec_out + dec_outl,
+                                 &tmplen)
+               != TSM_OK
+        || (dec_outl += tmplen, 0)
+        || tsm_ascon_aead_final(ctx_dec, dec_out + dec_outl, &tmplen) != TSM_OK) {
+        goto err;
+    }
+
+    dec_outl += tmplen;
+
+    dec_out[dec_outl] = 0;
+    printf("ASCON_AEAD_Decrypt=%s", dec_out);
+
+    printf("\n");
+
     ret = 0;
 err:
     tsm_ascon_aead_ctx_free(ctx);
+    tsm_ascon_aead_ctx_free(ctx_dec);
     tsm_free(key);
     tsm_free(iv);
     return ret;
diff --git a/src/ascon.c b/src/ascon.c
index dcaaeeb..18c552f 100644
--- a/src/ascon.c
+++ b/src/ascon.c
@@ -250,14 +250,23 @@ int tsm_ascon_aead_update(void *ctx, const unsigned char *in, size_t inl, unsign
             in += c->block_size - c->buf_len;
             inl -= c->block_size - c->buf_len;
             if (c->mode == TSM_ASCON_AEAD_128) {
-                c->s.x[0] ^= load_u64(c->buf, 8);
+                uint64_t c0 = load_u64(c->buf, 8);
+                c->s.x[0] ^= c0;
                 store_u64(c->s.x[0], 8, out);
+                if (c->flags & TSM_CIPH_FLAG_DECRYPT)
+                    c->s.x[0] = c0;
                 P6(&c->s);
             } else {
-                c->s.x[0] ^= load_u64(c->buf, 8);
-                c->s.x[1] ^= load_u64(c->buf + 8, 8);
+                uint64_t c0 = load_u64(c->buf, 8);
+                uint64_t c1 = load_u64(c->buf + 8, 8);
+                c->s.x[0] ^= c0;
+                c->s.x[1] ^= c1;
                 store_u64(c->s.x[0], 8, out);
                 store_u64(c->s.x[1], 8, out + 8);
+                if (c->flags & TSM_CIPH_FLAG_DECRYPT) {
+                    c->s.x[0] = c0;
+                    c->s.x[1] = c1;
+                }
                 P8(&c->s);
             }
             c->buf_len = 0;