add SetRandFunc function

1 files changed, 36 insertions, 0 deletions

diff --git a/ffi/go/bls/bls_test.go b/ffi/go/bls/bls_test.go
index 84c8be2..a2789f3 100644
--- a/ffi/go/bls/bls_test.go
+++ b/ffi/go/bls/bls_test.go
@@ -5,6 +5,8 @@ import "strconv"
 import "crypto/sha256"
 import "crypto/sha512"
 import "fmt"
+import "io"
+import "crypto/rand"
 
 var unitN = 0
 
@@ -448,6 +450,39 @@ func testAggregateHashes(t *testing.T) {
 	}
 }
 
+type SeqRead struct {
+}
+
+func (self *SeqRead) Read(buf []byte) (int, error) {
+	n := len(buf)
+	for i := 0; i < n; i++ {
+		buf[i] = byte(i)
+	}
+	return n, nil
+}
+
+func testReadRand(t *testing.T) {
+	var s1 io.Reader = &SeqRead{}
+	SetRandFunc(&s1)
+	var sec SecretKey
+	sec.SetByCSPRNG()
+	buf := sec.GetLittleEndian()
+	fmt.Printf("(SeqRead) buf=%x\n", buf)
+	for i := 0; i < len(buf); i++ {
+		if buf[i] != byte(i) {
+			t.Fatal("buf")
+		}
+	}
+	SetRandFunc(&rand.Reader)
+	sec.SetByCSPRNG()
+	buf = sec.GetLittleEndian()
+	fmt.Printf("(rand.Reader) buf=%x\n", buf)
+	SetRandFunc(nil)
+	sec.SetByCSPRNG()
+	buf = sec.GetLittleEndian()
+	fmt.Printf("(default) buf=%x\n", buf)
+}
+
 func test(t *testing.T, c int) {
 	err := Init(c)
 	if err != nil {
@@ -455,6 +490,7 @@ func test(t *testing.T, c int) {
 	}
 	unitN = GetOpUnitSize()
 	t.Logf("unitN=%d\n", unitN)
+	testReadRand(t)
 	testPre(t)
 	testRecoverSecretKey(t)
 	testAdd(t)