From 207fd1e7a4789c28479052a34957f0bcba10a958 Mon Sep 17 00:00:00 2001
From: Sun Yimin <emmansun@users.noreply.github.com>
Date: Wed, 7 Jun 2023 08:32:46 +0800
Subject: [PATCH] drbg: fix drbg prng read issue when reseed

---
 drbg/common.go      |  5 +++--
 drbg/common_test.go | 25 +++++++++++++++++++++++++
 2 files changed, 28 insertions(+), 2 deletions(-)

diff --git a/drbg/common.go b/drbg/common.go
index 0213cc4..e85a4b4 100644
--- a/drbg/common.go
+++ b/drbg/common.go
@@ -166,9 +166,10 @@ func (prng *DrbgPrng) Read(data []byte) (int, error) {
 			}
 		} else if err != nil {
 			return 0, err
+		} else {
+			total += len(b)
+			data = data[len(b):]
 		}
-		total += len(b)
-		data = data[len(b):]
 	}
 	return total, nil
 }
diff --git a/drbg/common_test.go b/drbg/common_test.go
index ff578d7..0b2aff4 100644
--- a/drbg/common_test.go
+++ b/drbg/common_test.go
@@ -1,6 +1,7 @@
 package drbg
 
 import (
+	"bytes"
 	"crypto/aes"
 	"crypto/sha256"
 	"testing"
@@ -23,6 +24,30 @@ func TestGmCtrDrbgPrng(t *testing.T) {
 	}
 }
 
+func TestGmCtrDrbgPrngReseedCase(t *testing.T) {
+	prng, err := NewGmCtrDrbgPrng(nil, 32, SECURITY_LEVEL_TEST, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	data := make([]byte, 64)
+	for i := 0; i < int(DRBG_RESEED_COUNTER_INTERVAL_LEVEL_TEST+1); i++ {
+		for j := 0; j < 64; j++ {
+			data[j] = 0
+		}
+		n, err := prng.Read(data)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if n != 64 {
+			t.Errorf("not got enough random bytes")
+		}
+		if bytes.Contains(data, []byte{0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0}) {
+			t.Fatal("failed, it's a bug")
+		}
+	}
+}
+
 func TestNistCtrDrbgPrng(t *testing.T) {
 	prng, err := NewNistCtrDrbgPrng(aes.NewCipher, 16, nil, 16, SECURITY_LEVEL_TEST, nil)
 	if err != nil {