From 4d75f337bb7c5c6e741cf65768fe572ad7726b1f Mon Sep 17 00:00:00 2001
From: Sean Parkinson <sparki@Seans-MBP.fios-router.home>
Date: Wed, 24 Jan 2018 16:00:47 -0800
Subject: [PATCH] Fix AVX2 final func to reset state

---
 wolfcrypt/benchmark/benchmark.c | 60 ++++++++++++++++++++++-----------
 wolfcrypt/src/poly1305.c        |  4 +++
 2 files changed, 45 insertions(+), 19 deletions(-)

diff --git a/wolfcrypt/benchmark/benchmark.c b/wolfcrypt/benchmark/benchmark.c
index 10971ef55..1a9f36fc1 100644
--- a/wolfcrypt/benchmark/benchmark.c
+++ b/wolfcrypt/benchmark/benchmark.c
@@ -2047,27 +2047,49 @@ void bench_poly1305()
     Poly1305 enc;
     byte     mac[16];
     double   start;
-    int      ret, i, count;
+    int      ret = 0, i, count;
 
-    ret = wc_Poly1305SetKey(&enc, bench_key, 32);
-    if (ret != 0) {
-        printf("Poly1305SetKey failed, ret = %d\n", ret);
-        return;
-    }
-
-    bench_stats_start(&count, &start);
-    do {
-        for (i = 0; i < numBlocks; i++) {
-            ret = wc_Poly1305Update(&enc, bench_plain, BENCH_SIZE);
-            if (ret != 0) {
-                printf("Poly1305Update failed: %d\n", ret);
-                break;
-            }
+    if (digest_stream) {
+        ret = wc_Poly1305SetKey(&enc, bench_key, 32);
+        if (ret != 0) {
+            printf("Poly1305SetKey failed, ret = %d\n", ret);
+            return;
         }
-        wc_Poly1305Final(&enc, mac);
-        count += i;
-    } while (bench_stats_sym_check(start));
-    bench_stats_sym_finish("POLY1305", 0, count, bench_size, start, ret);
+
+        bench_stats_start(&count, &start);
+        do {
+            for (i = 0; i < numBlocks; i++) {
+                ret = wc_Poly1305Update(&enc, bench_plain, BENCH_SIZE);
+                if (ret != 0) {
+                    printf("Poly1305Update failed: %d\n", ret);
+                    break;
+                }
+            }
+            wc_Poly1305Final(&enc, mac);
+            count += i;
+        } while (bench_stats_sym_check(start));
+        bench_stats_sym_finish("POLY1305", 0, count, bench_size, start, ret);
+    }
+    else {
+        bench_stats_start(&count, &start);
+        do {
+            for (i = 0; i < numBlocks; i++) {
+                ret = wc_Poly1305SetKey(&enc, bench_key, 32);
+                if (ret != 0) {
+                    printf("Poly1305SetKey failed, ret = %d\n", ret);
+                    return;
+                }
+                ret = wc_Poly1305Update(&enc, bench_plain, BENCH_SIZE);
+                if (ret != 0) {
+                    printf("Poly1305Update failed: %d\n", ret);
+                    break;
+                }
+                wc_Poly1305Final(&enc, mac);
+            }
+            count += i;
+        } while (bench_stats_sym_check(start));
+        bench_stats_sym_finish("POLY1305", 0, count, bench_size, start, ret);
+    }
 }
 #endif /* HAVE_POLY1305 */
 
diff --git a/wolfcrypt/src/poly1305.c b/wolfcrypt/src/poly1305.c
index d04ecccc3..eef1aaba2 100644
--- a/wolfcrypt/src/poly1305.c
+++ b/wolfcrypt/src/poly1305.c
@@ -1008,6 +1008,10 @@ static void poly1305_final_avx2(Poly1305* ctx, byte* mac)
           [r3] "r" (ctx->r3), [r4] "r" (ctx->r4)
         : "memory", "ymm0"
     );
+
+    ctx->leftover = 0;
+    ctx->finished = 0;
+    ctx->started = 0;
 }
 #endif