From 10a97b4c51696be94c1d4bd2b7ffd4df2a8368c4 Mon Sep 17 00:00:00 2001
From: Sun Yimin <emmansun@users.noreply.github.com>
Date: Wed, 5 Jun 2024 13:01:28 +0800
Subject: [PATCH] drbg: support NIST HMAC-DRBG 2

---
 drbg/common.go      | 37 +++++++++++++++++++++++++++++++++++++
 drbg/common_test.go | 16 ++++++++++++++++
 2 files changed, 53 insertions(+)

diff --git a/drbg/common.go b/drbg/common.go
index 24d1a9d1..9ae6bd7c 100644
--- a/drbg/common.go
+++ b/drbg/common.go
@@ -132,6 +132,43 @@ func NewGmHashDrbgPrng(entropySource io.Reader, securityStrength int, securityLe
 	return NewHashDrbgPrng(sm3.New, entropySource, securityStrength, true, securityLevel, personalization)
 }
 
+// NewHmacDrbgPrng create pseudo random number generator base on hash mac DRBG
+func NewHmacDrbgPrng(newHash func() hash.Hash, entropySource io.Reader, securityStrength int, gm bool, securityLevel SecurityLevel, personalization []byte) (*DrbgPrng, error) {
+	prng := new(DrbgPrng)
+	if entropySource != nil {
+		prng.entropySource = entropySource
+	} else {
+		prng.entropySource = rand.Reader
+	}
+	prng.securityStrength = selectSecurityStrength(securityStrength)
+
+	// Get entropy input
+	entropyInput := make([]byte, prng.securityStrength)
+	err := prng.getEntropy(entropyInput)
+	if err != nil {
+		return nil, err
+	}
+
+	// Get nonce from entropy source here
+	nonce := make([]byte, prng.securityStrength/2)
+	err = prng.getEntropy(nonce)
+	if err != nil {
+		return nil, err
+	}
+
+	prng.impl, err = NewHmacDrbg(newHash, securityLevel, gm, entropyInput, nonce, personalization)
+	if err != nil {
+		return nil, err
+	}
+
+	return prng, nil
+}
+
+// NewNistHmacDrbgPrng create pseudo random number generator base on hash mac DRBG which follows NIST standard
+func NewNistHmacDrbgPrng(newHash func() hash.Hash, entropySource io.Reader, securityStrength int, securityLevel SecurityLevel, personalization []byte) (*DrbgPrng, error) {
+	return NewHmacDrbgPrng(newHash, entropySource, securityStrength, false, securityLevel, personalization)
+}
+
 func (prng *DrbgPrng) getEntropy(entropyInput []byte) error {
 	n, err := prng.entropySource.Read(entropyInput)
 	if err != nil {
diff --git a/drbg/common_test.go b/drbg/common_test.go
index 0b2aff4a..f1cfe608 100644
--- a/drbg/common_test.go
+++ b/drbg/common_test.go
@@ -95,6 +95,22 @@ func TestNistHashDrbgPrng(t *testing.T) {
 	}
 }
 
+
+func TestNistHmacDrbgPrng(t *testing.T) {
+	prng, err := NewNistHmacDrbgPrng(sha256.New, nil, 32, SECURITY_LEVEL_TEST, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	data := make([]byte, MAX_BYTES_PER_GENERATE+1)
+	n, err := prng.Read(data)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if n != MAX_BYTES_PER_GENERATE+1 {
+		t.Errorf("not got enough random bytes")
+	}
+}
+
 func TestGMSecurityStrengthValidation(t *testing.T) {
 	_, err := NewGmHashDrbgPrng(nil, 24, SECURITY_LEVEL_TEST, nil)
 	if err == nil {