thefeed/internal/protocol/dns_test.go

package protocol

import (
	"strings"
	"testing"
)

func TestEncodeDecodeQuerySingleLabel(t *testing.T) {
	qk, _, err := DeriveKeys("test-key")
	if err != nil {
		t.Fatalf("DeriveKeys: %v", err)
	}
	domain := "t.example.com"
	tests := []struct {
		channel uint16
		block   uint16
	}{
		{0, 0},
		{1, 0},
		{1, 5},
		{255, 99},
	}
	for _, tt := range tests {
		qname, err := EncodeQuery(qk, tt.channel, tt.block, domain, QuerySingleLabel)
		if err != nil {
			t.Fatalf("EncodeQuery(%d, %d): %v", tt.channel, tt.block, err)
		}
		if !strings.HasSuffix(qname, "."+domain) {
			t.Errorf("query %q should end with .%s", qname, domain)
		}
		// Single label: only one dot before domain
		subdomain := qname[:len(qname)-len(domain)-1]
		if strings.Contains(subdomain, ".") {
			t.Errorf("single-label query should not have dots in subdomain, got %q", subdomain)
		}
		ch, blk, err := DecodeQuery(qk, qname, domain)
		if err != nil {
			t.Fatalf("DecodeQuery: %v", err)
		}
		if ch != tt.channel || blk != tt.block {
			t.Errorf("got ch=%d blk=%d, want ch=%d blk=%d", ch, blk, tt.channel, tt.block)
		}
	}
}

func TestEncodeDecodeQueryDoubleLabel(t *testing.T) {
	qk, _, err := DeriveKeys("test-key")
	if err != nil {
		t.Fatalf("DeriveKeys: %v", err)
	}
	domain := "t.example.com"
	qname, err := EncodeQuery(qk, 3, 7, domain, QueryDoubleLabel)
	if err != nil {
		t.Fatal(err)
	}
	// Double label: two hex labels before domain
	subdomain := qname[:len(qname)-len(domain)-1]
	parts := strings.Split(subdomain, ".")
	if len(parts) != 2 {
		t.Errorf("double-label query should have 2 parts, got %d: %q", len(parts), subdomain)
	}
	ch, blk, err := DecodeQuery(qk, qname, domain)
	if err != nil {
		t.Fatalf("DecodeQuery: %v", err)
	}
	if ch != 3 || blk != 7 {
		t.Errorf("got ch=%d blk=%d, want ch=3 blk=7", ch, blk)
	}
}

func TestDecodeQueryWrongKey(t *testing.T) {
	qk1, _, _ := DeriveKeys("key1")
	qk2, _, _ := DeriveKeys("key2")
	qname, _ := EncodeQuery(qk1, 1, 0, "t.example.com", QuerySingleLabel)
	_, _, err := DecodeQuery(qk2, qname, "t.example.com")
	if err == nil {
		t.Error("expected error when decoding with wrong key")
	}
}

func TestDecodeQueryWrongDomain(t *testing.T) {
	qk, _, _ := DeriveKeys("key")
	qname, _ := EncodeQuery(qk, 1, 0, "t.example.com", QuerySingleLabel)
	_, _, err := DecodeQuery(qk, qname, "t.other.com")
	if err == nil {
		t.Error("expected error for wrong domain")
	}
}

func TestEncodeDecodeResponse(t *testing.T) {
	_, rk, err := DeriveKeys("test-key")
	if err != nil {
		t.Fatal(err)
	}
	data := []byte("Hello World!")
	encoded, err := EncodeResponse(rk, data, DefaultMaxPadding)
	if err != nil {
		t.Fatalf("EncodeResponse: %v", err)
	}
	decoded, err := DecodeResponse(rk, encoded)
	if err != nil {
		t.Fatalf("DecodeResponse: %v", err)
	}
	if string(decoded) != string(data) {
		t.Errorf("got %q, want %q", decoded, data)
	}
}

func TestEncodeDecodeResponseNoPadding(t *testing.T) {
	_, rk, err := DeriveKeys("test-key")
	if err != nil {
		t.Fatal(err)
	}
	data := []byte("No padding test")
	encoded, err := EncodeResponse(rk, data, 0)
	if err != nil {
		t.Fatalf("EncodeResponse: %v", err)
	}
	decoded, err := DecodeResponse(rk, encoded)
	if err != nil {
		t.Fatalf("DecodeResponse: %v", err)
	}
	if string(decoded) != string(data) {
		t.Errorf("got %q, want %q", decoded, data)
	}
}

func TestResponseVaryingSize(t *testing.T) {
	_, rk, _ := DeriveKeys("test-key")
	data := []byte("fixed data")
	sizes := make(map[int]bool)
	for i := 0; i < 50; i++ {
		encoded, err := EncodeResponse(rk, data, 32)
		if err != nil {
			t.Fatal(err)
		}
		sizes[len(encoded)] = true
	}
	if len(sizes) < 2 {
		t.Error("expected varying response sizes with padding, got uniform")
	}
}

func TestDecodeResponseWrongKey(t *testing.T) {
	_, rk1, _ := DeriveKeys("key1")
	_, rk2, _ := DeriveKeys("key2")
	encoded, _ := EncodeResponse(rk1, []byte("data"), 0)
	_, err := DecodeResponse(rk2, encoded)
	if err == nil {
		t.Error("expected error for wrong key")
	}
}

func TestQueryDomainWithTrailingDot(t *testing.T) {
	qk, _, _ := DeriveKeys("key")
	qname, _ := EncodeQuery(qk, 1, 0, "t.example.com", QuerySingleLabel)
	ch, blk, err := DecodeQuery(qk, qname+".", "t.example.com.")
	if err != nil {
		t.Fatalf("DecodeQuery with trailing dot: %v", err)
	}
	if ch != 1 || blk != 0 {
		t.Errorf("got ch=%d blk=%d, want ch=1 blk=0", ch, blk)
	}
}