Sample pseudo random number generator base on DRBG

drbg/common.go

@@ -1,22 +1,168 @@
 package drbg
 
 import (
+	"crypto/cipher"
+	"crypto/rand"
+	"errors"
+	"hash"
+	"io"
 	"time"
+
+	"github.com/emmansun/gmsm/sm3"
+	"github.com/emmansun/gmsm/sm4"
 )
 
+const DRBG_RESEED_COUNTER_INTERVAL_LEVEL_TEST uint64 = 8
 const DRBG_RESEED_COUNTER_INTERVAL_LEVEL2 uint64 = 1 << 10
 const DRBG_RESEED_COUNTER_INTERVAL_LEVEL1 uint64 = 1 << 20
+const DRBG_RESEED_TIME_INTERVAL_LEVEL_TEST = time.Duration(6) * time.Second
 const DRBG_RESEED_TIME_INTERVAL_LEVEL2 = time.Duration(60) * time.Second
 const DRBG_RESEED_TIME_INTERVAL_LEVEL1 = time.Duration(600) * time.Second
 const MAX_BYTES = 1 << 27
+const MAX_BYTES_PER_GENERATE = 1 << 11
+
+var ErrReseedRequired = errors.New("reseed reuqired")
 
 type SecurityLevel byte
 
 const (
-	SECURITY_LEVEL_ONE SecurityLevel = 0x01
+	SECURITY_LEVEL_ONE  SecurityLevel = 0x01
-	SECURITY_LEVEL_TWO SecurityLevel = 0x02
+	SECURITY_LEVEL_TWO  SecurityLevel = 0x02
+	SECURITY_LEVEL_TEST SecurityLevel = 0x99
 )
 
+// DrbgPrng sample pseudo random number generator base on DRBG
+type DrbgPrng struct {
+	entropySource    io.Reader
+	securityStrength int
+	impl             DRBG
+}
+
+// NewCtrDrbgPrng create pseudo random number generator base on CTR DRBG
+func NewCtrDrbgPrng(cipherProvider func(key []byte) (cipher.Block, error), keyLen int, 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
+	nonce := make([]byte, prng.securityStrength/2)
+	err = prng.getEntropy(nonce)
+	if err != nil {
+		return nil, err
+	}
+
+	prng.impl, err = NewCtrDrbg(cipherProvider, keyLen, securityLevel, gm, entropyInput, nonce, personalization)
+	if err != nil {
+		return nil, err
+	}
+
+	return prng, nil
+}
+
+// NewNistCtrDrbgPrng create pseudo random number generator base on CTR DRBG which follows NIST standard
+func NewNistCtrDrbgPrng(cipherProvider func(key []byte) (cipher.Block, error), keyLen int, entropySource io.Reader, securityStrength int, securityLevel SecurityLevel, personalization []byte) (*DrbgPrng, error) {
+	return NewCtrDrbgPrng(cipherProvider, keyLen, entropySource, securityStrength, false, securityLevel, personalization)
+}
+
+// NewNistCtrDrbgPrng create pseudo random number generator base on CTR DRBG which follows GM/T 0105-2021 standard
+func NewGmCtrDrbgPrng(entropySource io.Reader, securityStrength int, securityLevel SecurityLevel, personalization []byte) (*DrbgPrng, error) {
+	return NewCtrDrbgPrng(sm4.NewCipher, 16, entropySource, securityStrength, true, securityLevel, personalization)
+}
+
+// NewHashDrbgPrng create pseudo random number generator base on HASH DRBG
+func NewHashDrbgPrng(md 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
+	nonce := make([]byte, prng.securityStrength/2)
+	err = prng.getEntropy(nonce)
+	if err != nil {
+		return nil, err
+	}
+
+	prng.impl, err = NewHashDrbg(md, securityLevel, gm, entropyInput, nonce, personalization)
+	if err != nil {
+		return nil, err
+	}
+
+	return prng, nil
+}
+
+// NewNistHashDrbgPrng create pseudo random number generator base on hash DRBG which follows NIST standard
+func NewNistHashDrbgPrng(md hash.Hash, entropySource io.Reader, securityStrength int, securityLevel SecurityLevel, personalization []byte) (*DrbgPrng, error) {
+	return NewHashDrbgPrng(md, entropySource, securityStrength, false, securityLevel, personalization)
+}
+
+// NewGmHashDrbgPrng create pseudo random number generator base on hash DRBG which follows GM/T 0105-2021 standard
+func NewGmHashDrbgPrng(entropySource io.Reader, securityStrength int, securityLevel SecurityLevel, personalization []byte) (*DrbgPrng, error) {
+	return NewHashDrbgPrng(sm3.New(), entropySource, securityStrength, true, securityLevel, personalization)
+}
+
+func (prng *DrbgPrng) getEntropy(entropyInput []byte) error {
+	n, err := prng.entropySource.Read(entropyInput)
+	if err != nil {
+		return err
+	}
+	if n != len(entropyInput) {
+		return errors.New("fail to read enough entropy input")
+	}
+	return nil
+}
+
+func (prng *DrbgPrng) Read(data []byte) (int, error) {
+	maxBytesPerRequest := prng.impl.MaxBytesPerRequest()
+	total := 0
+
+	for len(data) > 0 {
+		b := data
+		if len(data) > maxBytesPerRequest {
+			b = data[:maxBytesPerRequest]
+		}
+
+		err := prng.impl.Generate(b, nil)
+		if err == ErrReseedRequired {
+			entropyInput := make([]byte, prng.securityStrength)
+			err := prng.getEntropy(entropyInput)
+			if err != nil {
+				return 0, err
+			}
+			err = prng.impl.Reseed(entropyInput, nil)
+			if err != nil {
+				return 0, err
+			}
+		} else if err != nil {
+			return 0, err
+		}
+		total += len(b)
+		data = data[len(b):]
+	}
+	return total, nil
+}
+
 // DRBG interface for both hash and ctr drbg implementations
 type DRBG interface {
 	// check internal state, return if reseed required
@@ -25,6 +171,8 @@ type DRBG interface {
 	Reseed(entropy, additional []byte) error
 	// generate requrested bytes to b
 	Generate(b, additional []byte) error
+	// MaxBytesPerRequest return max bytes per request
+	MaxBytesPerRequest() int
 }
 
 type BaseDrbg struct {
@@ -42,6 +190,34 @@ func (hd *BaseDrbg) NeedReseed() bool {
 	return (hd.reseedCounter > hd.reseedIntervalInCounter) || (hd.gm && time.Since(hd.reseedTime) > hd.reseedIntervalInTime)
 }
 
+func (hd *BaseDrbg) setSecurityLevel(securityLevel SecurityLevel) {
+	hd.securityLevel = securityLevel
+	switch securityLevel {
+	case SECURITY_LEVEL_TWO:
+		hd.reseedIntervalInCounter = DRBG_RESEED_COUNTER_INTERVAL_LEVEL2
+		hd.reseedIntervalInTime = DRBG_RESEED_TIME_INTERVAL_LEVEL2
+	case SECURITY_LEVEL_TEST:
+		hd.reseedIntervalInCounter = DRBG_RESEED_COUNTER_INTERVAL_LEVEL_TEST
+		hd.reseedIntervalInTime = DRBG_RESEED_TIME_INTERVAL_LEVEL_TEST
+	default:
+		hd.reseedIntervalInCounter = DRBG_RESEED_COUNTER_INTERVAL_LEVEL1
+		hd.reseedIntervalInTime = DRBG_RESEED_TIME_INTERVAL_LEVEL1
+	}
+}
+
+func selectSecurityStrength(requested int) int {
+	switch {
+	case requested <= 14:
+		return 14
+	case requested <= 16:
+		return 16
+	case requested <= 24:
+		return 24
+	default:
+		return 32
+	}
+}
+
 func add(left, right []byte, len int) {
 	var temp uint16 = 0
 	for i := len - 1; i >= 0; i-- {

drbg/common_test.go

@@ -0,0 +1,71 @@
+package drbg
+
+import (
+	"crypto/aes"
+	"crypto/sha256"
+	"testing"
+)
+
+func TestGmCtrDrbgPrng(t *testing.T) {
+	prng, err := NewGmCtrDrbgPrng(nil, 16, SECURITY_LEVEL_TEST, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	data := make([]byte, 33)
+	for i := 0; i < int(DRBG_RESEED_COUNTER_INTERVAL_LEVEL_TEST+1); i++ {
+		n, err := prng.Read(data)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if n != 33 {
+			t.Errorf("not got enough random bytes")
+		}
+	}
+}
+
+func TestNistCtrDrbgPrng(t *testing.T) {
+	prng, err := NewNistCtrDrbgPrng(aes.NewCipher, 16, nil, 16, 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 TestGmHashDrbgPrng(t *testing.T) {
+	prng, err := NewGmHashDrbgPrng(nil, 32, SECURITY_LEVEL_TEST, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	data := make([]byte, 33)
+	for i := 0; i < int(DRBG_RESEED_COUNTER_INTERVAL_LEVEL_TEST+1); i++ {
+		n, err := prng.Read(data)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if n != 33 {
+			t.Errorf("not got enough random bytes")
+		}
+	}
+}
+
+func TestNistHashDrbgPrng(t *testing.T) {
+	prng, err := NewNistHashDrbgPrng(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")
+	}
+}

drbg/ctr_drbg.go

@@ -22,13 +22,7 @@ func NewCtrDrbg(cipherProvider func(key []byte) (cipher.Block, error), keyLen in
 	hd := &CtrDrbg{}
 
 	hd.gm = gm
-	hd.securityLevel = securityLevel
+	hd.setSecurityLevel(securityLevel)
-	hd.reseedIntervalInCounter = DRBG_RESEED_COUNTER_INTERVAL_LEVEL1
-	hd.reseedIntervalInTime = DRBG_RESEED_TIME_INTERVAL_LEVEL1
-	if hd.securityLevel == SECURITY_LEVEL_TWO {
-		hd.reseedIntervalInCounter = DRBG_RESEED_COUNTER_INTERVAL_LEVEL2
-		hd.reseedIntervalInTime = DRBG_RESEED_TIME_INTERVAL_LEVEL2
-	}
 
 	// here for the min length, we just check <=0 now
 	if len(entropy) <= 0 || len(entropy) >= MAX_BYTES {
@@ -109,10 +103,17 @@ func (hd *CtrDrbg) newBlockCipher(key []byte) cipher.Block {
 	return block
 }
 
+func (hd *CtrDrbg) MaxBytesPerRequest() int {
+	if hd.gm {
+		return len(hd.v)
+	}
+	return MAX_BYTES_PER_GENERATE
+}
+
 // Generate CTR DRBG generate process.
 func (hd *CtrDrbg) Generate(b, additional []byte) error {
 	if hd.NeedReseed() {
-		return errors.New("reseed reuqired")
+		return ErrReseedRequired
 	}
 	outlen := len(hd.v)
 	if (hd.gm && len(b) > outlen) || (!hd.gm && len(b) > MAX_BYTES_PER_GENERATE) {

drbg/hash_drbg.go

@@ -11,7 +11,6 @@ import (
 
 const HASH_DRBG_SEED_SIZE = 55
 const HASH_DRBG_MAX_SEED_SIZE = 111
-const MAX_BYTES_PER_GENERATE = 1 << 11
 
 type HashDrbg struct {
 	BaseDrbg
@@ -24,13 +23,7 @@ func NewHashDrbg(md hash.Hash, securityLevel SecurityLevel, gm bool, entropy, no
 	hd := &HashDrbg{}
 
 	hd.gm = gm
-	hd.securityLevel = securityLevel
+	hd.setSecurityLevel(securityLevel)
-	hd.reseedIntervalInCounter = DRBG_RESEED_COUNTER_INTERVAL_LEVEL1
-	hd.reseedIntervalInTime = DRBG_RESEED_TIME_INTERVAL_LEVEL1
-	if hd.securityLevel == SECURITY_LEVEL_TWO {
-		hd.reseedIntervalInCounter = DRBG_RESEED_COUNTER_INTERVAL_LEVEL2
-		hd.reseedIntervalInTime = DRBG_RESEED_TIME_INTERVAL_LEVEL2
-	}
 
 	// here for the min length, we just check <=0 now
 	if len(entropy) <= 0 || len(entropy) >= MAX_BYTES {
@@ -140,11 +133,18 @@ func (hd *HashDrbg) addReseedCounter() {
 	add(t, hd.v, hd.seedLength)
 }
 
+func (hd *HashDrbg) MaxBytesPerRequest() int {
+	if hd.gm {
+		return hd.md.Size()
+	}
+	return MAX_BYTES_PER_GENERATE
+}
+
 // Generate hash DRBG generate process. GM/T 0105-2021 has a little different with NIST.
 // GM/T 0105-2021 can only generate no more than hash.Size bytes once.
 func (hd *HashDrbg) Generate(b, additional []byte) error {
 	if hd.NeedReseed() {
-		return errors.New("reseed reuqired")
+		return ErrReseedRequired
 	}
 	if (hd.gm && len(b) > hd.md.Size()) || (!hd.gm && len(b) > MAX_BYTES_PER_GENERATE) {
 		return errors.New("too many bytes requested")