add AES128 based nonce generator

crypt_test.go

@@ -2,6 +2,7 @@ package kcp
 
 import (
 	"bytes"
+	"crypto/aes"
 	"crypto/md5"
 	"crypto/rand"
 	"crypto/sha1"
@@ -267,10 +268,21 @@ func BenchmarkCsprngSHA1(b *testing.B) {
 
 func BenchmarkCsprngNonceMD5(b *testing.B) {
 	var ng nonceMD5
-
+	ng.Init()
 	b.SetBytes(md5.Size)
 	data := make([]byte, md5.Size)
 	for i := 0; i < b.N; i++ {
 		ng.Fill(data)
 	}
 }
+
+func BenchmarkCsprngNonceAES128(b *testing.B) {
+	var ng nonceAES128
+	ng.Init()
+
+	b.SetBytes(aes.BlockSize)
+	data := make([]byte, aes.BlockSize)
+	for i := 0; i < b.N; i++ {
+		ng.Fill(data)
+	}
+}

entropy.go

@@ -1,24 +1,51 @@
 package kcp
 
 import (
+	"crypto/aes"
+	"crypto/cipher"
 	"crypto/md5"
 	"crypto/rand"
 	"io"
 )
 
-// nonceMD5 is a nonce generator for each packet header
-// which took the advantages of both MD5 and CSPRNG(like /dev/urandom).
-// The benchmark shows it's faster than previous CSPRNG only method.
+type Entropy interface {
+	Init()
+	Fill(nonce []byte)
+}
+
+// nonceMD5 nonce generator for packet header
 type nonceMD5 struct {
-	data [md5.Size]byte
+	seed [md5.Size]byte
 }
 
-// Fill fills a nonce into the provided slice with no more than md5.Size bytes
-// the entropy will be updated whenever a leading 0 appears
+func (n *nonceMD5) Init() { /*nothing required*/ }
+
 func (n *nonceMD5) Fill(nonce []byte) {
-	if n.data[0] == 0 { // 1/256 chance for entropy update
-		io.ReadFull(rand.Reader, n.data[:])
+	if n.seed[0] == 0 { // entropy update
+		io.ReadFull(rand.Reader, n.seed[:])
+	}
+	n.seed = md5.Sum(n.seed[:])
+	copy(nonce, n.seed[:])
+}
+
+// nonceAES128 nonce generator for packet headers
+type nonceAES128 struct {
+	seed  [aes.BlockSize]byte
+	block cipher.Block
+}
+
+func (n *nonceAES128) Init() {
+	var key [16]byte //aes-128
+	io.ReadFull(rand.Reader, key[:])
+	io.ReadFull(rand.Reader, n.seed[:])
+	block, _ := aes.NewCipher(key[:])
+	n.block = block
+}
+
+func (n *nonceAES128) Fill(nonce []byte) {
+	if n.seed[0] == 0 { // entropy update
+		io.ReadFull(rand.Reader, n.seed[:])
 	}
-	n.data = md5.Sum(n.data[:])
-	copy(nonce, n.data[:])
+	n.block.Encrypt(n.seed[:], n.seed[:])
+	copy(nonce, n.seed[:])
 }

sess.go

@@ -97,7 +97,7 @@ type (
 		chErrorEvent chan error    // notify Read() have an error
 
 		// nonce generator
-		nonce nonceMD5
+		nonce Entropy
 
 		isClosed bool // flag the session has Closed
 		mu       sync.Mutex
@@ -116,6 +116,8 @@ type (
 func newUDPSession(conv uint32, dataShards, parityShards int, l *Listener, conn net.PacketConn, remote net.Addr, block BlockCrypt) *UDPSession {
 	sess := new(UDPSession)
 	sess.die = make(chan struct{})
+	sess.nonce = new(nonceAES128)
+	sess.nonce.Init()
 	sess.chReadEvent = make(chan struct{}, 1)
 	sess.chWriteEvent = make(chan struct{}, 1)
 	sess.chErrorEvent = make(chan error, 1)