add SetRandFunc function

3 files changed, 82 insertions, 0 deletions

diff --git a/ffi/go/bls/bls.go b/ffi/go/bls/bls.go
index a5c657c..2374d2e 100644
--- a/ffi/go/bls/bls.go
+++ b/ffi/go/bls/bls.go
@@ -10,11 +10,14 @@ package bls
 #cgo LDFLAGS:-lbls384
 #cgo LDFLAGS:-lcrypto -lgmp -lgmpxx -lstdc++
 #include "config.h"
+typedef unsigned int (*ReadRandFunc)(void *, void *, unsigned int);
+int wrapReadRandCgo(void *self, void *buf, unsigned int n);
 #include <bls/bls.h>
 */
 import "C"
 import "fmt"
 import "unsafe"
+import "io"
 
 // Init --
 // call this function before calling all the other operations
@@ -384,3 +387,34 @@ func (sign *Sign) VerifyAggregateHashes(pubVec []PublicKey, hash [][]byte) bool
 	}
 	return C.blsVerifyAggregatedHashes(sign.getPointer(), pubVec[0].getPointer(), unsafe.Pointer(&h[0]), C.size_t(hashByte), C.size_t(n)) == 1
 }
+
+///
+
+var s_randReader *io.Reader
+
+func createSlice(buf *C.char, n C.uint) []byte {
+	size := int(n)
+	return (*[1 << 30]byte)(unsafe.Pointer(buf))[:size:size]
+}
+
+// this function can't be put in callback.go
+//export wrapReadRandGo
+func wrapReadRandGo(buf *C.char, n C.uint) C.uint {
+	slice := createSlice(buf, n)
+	ret, err := (*s_randReader).Read(slice)
+	if ret == int(n) && err == nil {
+		return n
+	}
+	return 0
+}
+
+// SetRandFunc --
+func SetRandFunc(randReader *io.Reader) {
+	s_randReader = randReader
+	if randReader != nil {
+		C.blsSetRandFunc(nil, C.ReadRandFunc(unsafe.Pointer(C.wrapReadRandCgo)))
+	} else {
+		// use default random generator
+		C.blsSetRandFunc(nil, C.ReadRandFunc(unsafe.Pointer(nil)))
+	}
+}
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)
diff --git a/ffi/go/bls/callback.go b/ffi/go/bls/callback.go
new file mode 100644
index 0000000..ba73a5e
--- /dev/null
+++ b/ffi/go/bls/callback.go
@@ -0,0 +1,12 @@
+package bls
+
+/*
+// exported from bls.go
+unsigned int wrapReadRandGo(void *buf, unsigned int n);
+int wrapReadRandCgo(void *self, void *buf, unsigned int n)
+{
+	(void)self;
+	return wrapReadRandGo(buf, n);
+}
+*/
+import "C"