diff --git a/crypt_test.go b/crypt_test.go index b00a4e1..db0a857 100644 --- a/crypt_test.go +++ b/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") } } diff --git a/entropy_test.go b/entropy_test.go new file mode 100644 index 0000000..2e066e2 --- /dev/null +++ b/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[:]) + } +} diff --git a/kcp_test.go b/kcp_test.go index 4d48f54..964e442 100644 --- a/kcp_test.go +++ b/kcp_test.go @@ -154,7 +154,7 @@ func BenchmarkFlush(b *testing.B) { for range kcp.snd_buf.MaxLen() { kcp.snd_buf.Push(segment{xmit: 1, resendts: currentMs() + 10000}) } - + b.ReportAllocs() var mu sync.Mutex for b.Loop() { diff --git a/sess_test.go b/sess_test.go index e6049fd..a31114b 100644 --- a/sess_test.go +++ b/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 +}