test: improve coverage across package

analyzer/udp/wireguard_test.go

@@ -0,0 +1,187 @@
+package udp
+
+import (
+	"encoding/binary"
+	"testing"
+
+	"git.difuse.io/Difuse/Mellaris/analyzer"
+)
+
+func makeWireGuardPacket(msgType byte, body []byte) []byte {
+	buf := make([]byte, 1+3+len(body))
+	buf[0] = msgType
+	copy(buf[4:], body)
+	return buf
+}
+
+func TestWireGuardUDPStream_Feed_HandshakeInitiation(t *testing.T) {
+	s := newWireGuardUDPStream(nil)
+	body := make([]byte, wireguardSizeHandshakeInitiation-4)
+	binary.LittleEndian.PutUint32(body[0:4], 0x01020304)
+
+	u, done := s.Feed(false, makeWireGuardPacket(wireguardTypeHandshakeInitiation, body))
+	if u == nil {
+		t.Fatal("Feed returned nil update")
+	}
+	if done {
+		t.Error("should not be done")
+	}
+	if u.Type != analyzer.PropUpdateReplace {
+		t.Errorf("Type = %d, want PropUpdateReplace", u.Type)
+	}
+	initMap, ok := u.M["handshake_initiation"].(analyzer.PropMap)
+	if !ok {
+		t.Fatal("handshake_initiation missing")
+	}
+	if idx := initMap["sender_index"].(uint32); idx != 0x01020304 {
+		t.Errorf("sender_index = %d, want %d", idx, 0x01020304)
+	}
+}
+
+func TestWireGuardUDPStream_Feed_HandshakeResponse(t *testing.T) {
+	s := newWireGuardUDPStream(nil)
+	body := make([]byte, wireguardSizeHandshakeResponse-4)
+	binary.LittleEndian.PutUint32(body[0:4], 0x01020304)
+	binary.LittleEndian.PutUint32(body[4:8], 0x05060708)
+
+	u, _ := s.Feed(false, makeWireGuardPacket(wireguardTypeHandshakeResponse, body))
+	if u == nil {
+		t.Fatal("Feed returned nil update")
+	}
+	respMap, ok := u.M["handshake_response"].(analyzer.PropMap)
+	if !ok {
+		t.Fatal("handshake_response missing")
+	}
+	if idx := respMap["sender_index"].(uint32); idx != 0x01020304 {
+		t.Errorf("sender_index = %d", idx)
+	}
+	if idx := respMap["receiver_index"].(uint32); idx != 0x05060708 {
+		t.Errorf("receiver_index = %d", idx)
+	}
+}
+
+func TestWireGuardUDPStream_Feed_PacketData(t *testing.T) {
+	s := newWireGuardUDPStream(nil)
+	body := make([]byte, wireguardMinSizePacketData-4)
+	binary.LittleEndian.PutUint32(body[0:4], 0x0a0b0c0d)
+	binary.LittleEndian.PutUint64(body[4:12], 42)
+
+	u, _ := s.Feed(false, makeWireGuardPacket(wireguardTypeData, body))
+	if u == nil {
+		t.Fatal("Feed returned nil update")
+	}
+	dataMap, ok := u.M["packet_data"].(analyzer.PropMap)
+	if !ok {
+		t.Fatal("packet_data missing")
+	}
+	if idx := dataMap["receiver_index"].(uint32); idx != 0x0a0b0c0d {
+		t.Errorf("receiver_index = %d", idx)
+	}
+	if ctr := dataMap["counter"].(uint64); ctr != 42 {
+		t.Errorf("counter = %d", ctr)
+	}
+}
+
+func TestWireGuardUDPStream_Feed_CookieReply(t *testing.T) {
+	s := newWireGuardUDPStream(nil)
+	body := make([]byte, wireguardSizePacketCookieReply-4)
+	binary.LittleEndian.PutUint32(body[0:4], 0x11111111)
+
+	u, _ := s.Feed(false, makeWireGuardPacket(wireguardTypeCookieReply, body))
+	if u == nil {
+		t.Fatal("Feed returned nil update")
+	}
+	crMap, ok := u.M["packet_cookie_reply"].(analyzer.PropMap)
+	if !ok {
+		t.Fatal("packet_cookie_reply missing")
+	}
+	if idx := crMap["receiver_index"].(uint32); idx != 0x11111111 {
+		t.Errorf("receiver_index = %d", idx)
+	}
+}
+
+func TestWireGuardUDPStream_Feed_InvalidLength(t *testing.T) {
+	s := newWireGuardUDPStream(nil)
+	u, done := s.Feed(false, []byte{wireguardTypeHandshakeInitiation, 0, 0, 0, 0x01})
+	if u == nil || !done {
+		u, done = s.Feed(false, []byte{wireguardTypeHandshakeInitiation, 0, 0, 0, 0x01})
+	}
+	if u == nil || !done {
+		u, _ = s.Feed(false, []byte{wireguardTypeHandshakeInitiation, 0, 0, 0, 0x01})
+	}
+	u, done = s.Feed(false, []byte{wireguardTypeHandshakeInitiation, 0, 0, 0, 0x01})
+	if u != nil || !done {
+		t.Error("should return done after 4 invalid packets")
+	}
+}
+
+func TestWireGuardUDPStream_Feed_TooShort(t *testing.T) {
+	s := newWireGuardUDPStream(nil)
+	u, _ := s.Feed(false, []byte{1, 0, 0})
+	if u != nil {
+		t.Error("should return nil for too-short packet")
+	}
+}
+
+func TestWireGuardUDPStream_Feed_NonZeroReserved(t *testing.T) {
+	s := newWireGuardUDPStream(nil)
+	u, _ := s.Feed(false, []byte{wireguardTypeHandshakeInitiation, 0, 0, 1, 0, 0, 0, 0})
+	if u != nil {
+		t.Error("should return nil when reserved bytes are non-zero")
+	}
+}
+
+func TestWireGuardUDPStream_Feed_HandshakeInitiation_WrongSize(t *testing.T) {
+	s := newWireGuardUDPStream(nil)
+	body := make([]byte, wireguardSizeHandshakeInitiation-4+1)
+	binary.LittleEndian.PutUint32(body[0:4], 1)
+	u, _ := s.Feed(false, makeWireGuardPacket(wireguardTypeHandshakeInitiation, body))
+	if u != nil {
+		t.Error("should return nil for wrong init size")
+	}
+}
+
+func TestWireGuardUDPStream_Close(t *testing.T) {
+	s := newWireGuardUDPStream(nil)
+	u := s.Close(false)
+	if u != nil {
+		t.Error("Close() should return nil")
+	}
+}
+
+func TestWireGuardUDPStream_PutSenderIndex_MatchReceiverIndex(t *testing.T) {
+	s := newWireGuardUDPStream(nil)
+
+	s.putSenderIndex(false, 0xdeadbeef)
+	if !s.matchReceiverIndex(true, 0xdeadbeef) {
+		t.Error("should match reverse direction")
+	}
+	if s.matchReceiverIndex(false, 0xdeadbeef) {
+		t.Error("should not match same direction")
+	}
+	if s.matchReceiverIndex(true, 0x12345678) {
+		t.Error("should not match wrong index")
+	}
+}
+
+func TestWireGuardUDPStream_RememberedIndexRing(t *testing.T) {
+	s := newWireGuardUDPStream(nil)
+
+	for i := uint32(0); i < wireguardRememberedIndexCount+2; i++ {
+		s.putSenderIndex(false, i)
+	}
+
+	if s.matchReceiverIndex(true, 0) {
+		t.Error("index 0 should have been evicted from ring")
+	}
+	if !s.matchReceiverIndex(true, uint32(wireguardRememberedIndexCount+1)) {
+		t.Error("latest index should still be in ring")
+	}
+}
+
+func TestWireGuardAnalyzer_Name(t *testing.T) {
+	a := &WireGuardAnalyzer{}
+	if a.Name() != "wireguard" {
+		t.Errorf("Name() = %q, want wireguard", a.Name())
+	}
+}