refine tests for better coverage

mv entropy related test to entropy_test.go

crypt_test.go

@@ -26,7 +26,6 @@ import (
 	"bytes"
 	"crypto/aes"
 	"crypto/cipher"
-	"crypto/md5"
 	"crypto/rand"
 	"encoding/binary"
 	"hash/crc32"
@@ -428,31 +427,35 @@ func BenchmarkAEAD_Chacha20_Poly1035(b *testing.B) {
 	}
 }
 
-func BenchmarkCsprngSystem(b *testing.B) {
-	data := make([]byte, md5.Size)
-	b.SetBytes(int64(len(data)))
+func TestCryptErrors(t *testing.T) {
+	invalidKey := []byte("invalid")
 
-	for b.Loop() {
-		io.ReadFull(rand.Reader, data)
+	if _, err := NewSM4BlockCrypt(invalidKey); err == nil {
+		t.Error("NewSM4BlockCrypt should fail with invalid key")
 	}
-}
-
-func BenchmarkCsprngAES128(b *testing.B) {
-	var data [aes.BlockSize]byte
-	b.SetBytes(aes.BlockSize)
-
-	r := NewEntropyAES()
-	for b.Loop() {
-		io.ReadFull(r, data[:])
+	if _, err := NewTwofishBlockCrypt(invalidKey); err == nil {
+		t.Error("NewTwofishBlockCrypt should fail with invalid key")
+	}
+	if _, err := NewTripleDESBlockCrypt(invalidKey); err == nil {
+		t.Error("NewTripleDESBlockCrypt should fail with invalid key")
+	}
+	if _, err := NewCast5BlockCrypt(invalidKey); err == nil {
+		t.Error("NewCast5BlockCrypt should fail with invalid key")
+	}
+	// Blowfish supports variable key length, so "invalid" (7 bytes) might be valid.
+	// Blowfish key size: 1-56 bytes. So 7 bytes is valid.
+	// Let's try empty key or very long key.
+	if _, err := NewBlowfishBlockCrypt(nil); err == nil {
+		t.Error("NewBlowfishBlockCrypt should fail with nil key")
 	}
-}
-
-func BenchmarkCsprngChacha8(b *testing.B) {
-	var data [8]byte
-	b.SetBytes(8)
 
-	r := NewEntropyChacha8()
-	for b.Loop() {
-		io.ReadFull(r, data[:])
+	if _, err := NewAESBlockCrypt(invalidKey); err == nil {
+		t.Error("NewAESBlockCrypt should fail with invalid key")
+	}
+	if _, err := NewTEABlockCrypt(invalidKey); err == nil {
+		t.Error("NewTEABlockCrypt should fail with invalid key")
+	}
+	if _, err := NewXTEABlockCrypt(invalidKey); err == nil {
+		t.Error("NewXTEABlockCrypt should fail with invalid key")
 	}
 }

entropy_test.go

@@ -0,0 +1,142 @@
+package kcp
+
+import (
+	"crypto/aes"
+	"crypto/md5"
+	"crypto/rand"
+	"io"
+	"testing"
+)
+
+func TestEntropyAES(t *testing.T) {
+	r := NewEntropyAES()
+	buf := make([]byte, 16)
+	n, err := r.Read(buf)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if n != 16 {
+		t.Fatalf("expected 16 bytes, got %d", n)
+	}
+
+	// Test empty read
+	n, err = r.Read(nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if n != 0 {
+		t.Fatalf("expected 0 bytes, got %d", n)
+	}
+}
+
+func TestEntropyChacha8(t *testing.T) {
+	r := NewEntropyChacha8()
+	buf := make([]byte, 32)
+	n, err := r.Read(buf)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if n != 32 {
+		t.Fatalf("expected 32 bytes, got %d", n)
+	}
+
+	// Test empty read
+	n, err = r.Read(nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if n != 0 {
+		t.Fatalf("expected 0 bytes, got %d", n)
+	}
+}
+
+func TestSetEntropy(t *testing.T) {
+	// Save original entropy
+	orig := entropy
+	defer SetEntropy(orig)
+
+	// Set custom entropy
+	custom := NewEntropyChacha8()
+	SetEntropy(custom)
+
+	if entropy != custom {
+		t.Fatal("SetEntropy failed")
+	}
+
+	// Test fillRand with custom entropy
+	buf := make([]byte, 10)
+	fillRand(buf)
+	// We can't easily verify the content is from custom, but we can verify it didn't panic
+}
+
+func TestFillRand(t *testing.T) {
+	buf := make([]byte, 100)
+	fillRand(buf)
+
+	// Check if buffer is not all zeros (probabilistic, but highly likely)
+	allZero := true
+	for _, b := range buf {
+		if b != 0 {
+			allZero = false
+			break
+		}
+	}
+	if allZero {
+		t.Fatal("fillRand produced all zeros")
+	}
+
+	// Test empty buffer
+	fillRand(nil)
+}
+
+type mockReader struct{}
+
+func (m *mockReader) Read(p []byte) (n int, err error) {
+	for i := range p {
+		p[i] = 0xAA
+	}
+	return len(p), nil
+}
+
+func TestSetEntropyMock(t *testing.T) {
+	orig := entropy
+	defer SetEntropy(orig)
+
+	SetEntropy(&mockReader{})
+	buf := make([]byte, 10)
+	fillRand(buf)
+	for _, b := range buf {
+		if b != 0xAA {
+			t.Fatalf("expected 0xAA, got %x", b)
+		}
+	}
+}
+
+func BenchmarkCsprngSystem(b *testing.B) {
+	data := make([]byte, md5.Size)
+	b.SetBytes(int64(len(data)))
+
+	for b.Loop() {
+		io.ReadFull(rand.Reader, data)
+	}
+}
+
+func BenchmarkCsprngAES128(b *testing.B) {
+	var data [aes.BlockSize]byte
+	b.SetBytes(aes.BlockSize)
+
+	r := NewEntropyAES()
+	for b.Loop() {
+		io.ReadFull(r, data[:])
+	}
+}
+
+func BenchmarkCsprngChacha8(b *testing.B) {
+	var data [8]byte
+	b.SetBytes(8)
+
+	r := NewEntropyChacha8()
+	for b.Loop() {
+		io.ReadFull(r, data[:])
+	}
+}

sess_test.go

@@ -1269,3 +1269,52 @@ func TestSealAllocated(t *testing.T) {
 	b := make([]byte, 100*1024*1024) // 100 MB
 	cli.Write(b)
 }
+
+func TestSessionGetters(t *testing.T) {
+	sess := new(UDPSession)
+	sess.kcp = NewKCP(1, func(buf []byte, size int) {})
+
+	if sess.GetConv() != 1 {
+		t.Error("GetConv failed")
+	}
+	// RTO, SRTT, SRTTVar are dynamic, just check they don't panic
+	sess.GetRTO()
+	sess.GetSRTT()
+	sess.GetSRTTVar()
+}
+
+func TestTimedSchedClose(t *testing.T) {
+	ts := NewTimedSched(1)
+	ts.Close()
+}
+
+func TestListenDial(t *testing.T) {
+	l, err := Listen("127.0.0.1:0")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer l.Close()
+	t.Logf("Listening on %s", l.Addr().String())
+
+	ch := make(chan struct{})
+	go func() {
+		conn, err := l.Accept()
+		if err != nil {
+			return
+		}
+		t.Log("Accepted connection")
+		conn.Close()
+		close(ch)
+	}()
+
+	conn, err := Dial(l.Addr().String())
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Log("Dialed")
+	conn.Write([]byte("hello"))
+	t.Log("Wrote data")
+	time.Sleep(100 * time.Millisecond)
+	conn.Close()
+	<-ch
+}