plcdns / plcdns_test.go

package main

import (
	"encoding/json"
	"net"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/miekg/dns"
)

// Mock DID document for testing
var mockDIDDocument = DIDDocument{
	ID: "did:plc:test123",
	AlsoKnownAs: []string{
		"at://test.bsky.social",
	},
	VerificationMethod: []VerificationMethod{
		{
			ID:                 "did:plc:test123#atproto",
			Type:               "Multikey",
			Controller:         "did:plc:test123",
			PublicKeyMultibase: "zQ3shXjHeiBuRCKmM36cuYnm7YEMzhGnCmCyW92sRJ9przSMS",
		},
	},
	Service: []Service{
		{
			ID:              "#atproto_pds",
			Type:            "AtprotoPersonalDataServer",
			ServiceEndpoint: "https://bsky.social",
		},
		{
			ID:              "#atproto_labeler",
			Type:            "AtprotoLabeler",
			ServiceEndpoint: "https://mod.bsky.app",
		},
	},
}

// Test domain parsing
func TestParseDomain(t *testing.T) {
	handler := NewPLCHandler("https://plc.directory")

	tests := []struct {
		name          string
		domain        string
		expectedDID   string
		expectedType  QueryType
		expectedValid bool
	}{
		{
			name:          "Valid handle query",
			domain:        "_handle.test123.plc.atscan.net",
			expectedDID:   "did:plc:test123",
			expectedType:  QueryHandle,
			expectedValid: true,
		},
		{
			name:          "Valid PDS query",
			domain:        "_pds.test123.plc.atscan.net",
			expectedDID:   "did:plc:test123",
			expectedType:  QueryPDS,
			expectedValid: true,
		},
		{
			name:          "Valid labeler query",
			domain:        "_labeler.test123.plc.atscan.net",
			expectedDID:   "did:plc:test123",
			expectedType:  QueryLabeler,
			expectedValid: true,
		},
		{
			name:          "Valid pubkey query",
			domain:        "_pubkey.test123.plc.atscan.net",
			expectedDID:   "did:plc:test123",
			expectedType:  QueryPubKey,
			expectedValid: true,
		},
		{
			name:          "Invalid prefix",
			domain:        "_invalid.test123.plc.atscan.net",
			expectedDID:   "",
			expectedType:  QueryInvalid,
			expectedValid: false,
		},
		{
			name:          "Too few parts",
			domain:        "_handle.test123.plc",
			expectedDID:   "",
			expectedType:  QueryInvalid,
			expectedValid: false,
		},
		{
			name:          "Domain with trailing dot",
			domain:        "_handle.test123.plc.atscan.net.",
			expectedDID:   "did:plc:test123",
			expectedType:  QueryHandle,
			expectedValid: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			did, queryType, valid := handler.parseDomain(tt.domain)

			if valid != tt.expectedValid {
				t.Errorf("expected valid=%v, got %v", tt.expectedValid, valid)
			}

			if did != tt.expectedDID {
				t.Errorf("expected DID=%s, got %s", tt.expectedDID, did)
			}

			if queryType != tt.expectedType {
				t.Errorf("expected queryType=%v, got %v", tt.expectedType, queryType)
			}
		})
	}
}

// Test getting handle from DID document
func TestGetHandle(t *testing.T) {
	handler := NewPLCHandler("https://plc.directory")

	tests := []struct {
		name     string
		doc      *DIDDocument
		expected string
	}{
		{
			name:     "Valid handle",
			doc:      &mockDIDDocument,
			expected: "test.bsky.social",
		},
		{
			name: "No handle",
			doc: &DIDDocument{
				AlsoKnownAs: []string{},
			},
			expected: "",
		},
		{
			name: "Non-AT protocol handle",
			doc: &DIDDocument{
				AlsoKnownAs: []string{"https://example.com"},
			},
			expected: "",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := handler.getHandle(tt.doc)
			if result != tt.expected {
				t.Errorf("expected %s, got %s", tt.expected, result)
			}
		})
	}
}

// Test getting PDS from DID document
func TestGetPDS(t *testing.T) {
	handler := NewPLCHandler("https://plc.directory")

	tests := []struct {
		name     string
		doc      *DIDDocument
		expected string
	}{
		{
			name:     "Valid PDS",
			doc:      &mockDIDDocument,
			expected: "https://bsky.social",
		},
		{
			name: "No PDS",
			doc: &DIDDocument{
				Service: []Service{},
			},
			expected: "",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := handler.getPDS(tt.doc)
			if result != tt.expected {
				t.Errorf("expected %s, got %s", tt.expected, result)
			}
		})
	}
}

// Test getting labeler from DID document
func TestGetLabeler(t *testing.T) {
	handler := NewPLCHandler("https://plc.directory")

	tests := []struct {
		name     string
		doc      *DIDDocument
		expected string
	}{
		{
			name:     "Valid labeler",
			doc:      &mockDIDDocument,
			expected: "https://mod.bsky.app",
		},
		{
			name: "No labeler",
			doc: &DIDDocument{
				Service: []Service{},
			},
			expected: "",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := handler.getLabeler(tt.doc)
			if result != tt.expected {
				t.Errorf("expected %s, got %s", tt.expected, result)
			}
		})
	}
}

// Test getting pubkey from DID document
func TestGetPubKey(t *testing.T) {
	handler := NewPLCHandler("https://plc.directory")

	tests := []struct {
		name     string
		doc      *DIDDocument
		expected string
	}{
		{
			name:     "Valid pubkey",
			doc:      &mockDIDDocument,
			expected: "zQ3shXjHeiBuRCKmM36cuYnm7YEMzhGnCmCyW92sRJ9przSMS",
		},
		{
			name: "No pubkey",
			doc: &DIDDocument{
				VerificationMethod: []VerificationMethod{},
			},
			expected: "",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := handler.getPubKey(tt.doc)
			if result != tt.expected {
				t.Errorf("expected %s, got %s", tt.expected, result)
			}
		})
	}
}

// Test creating TXT records
func TestCreateTXTRecord(t *testing.T) {
	handler := NewPLCHandler("https://plc.directory")

	tests := []struct {
		name          string
		doc           *DIDDocument
		qname         string
		queryType     QueryType
		expectedValue string
		shouldBeEmpty bool
	}{
		{
			name:          "Handle record",
			doc:           &mockDIDDocument,
			qname:         "_handle.test123.plc.atscan.net.",
			queryType:     QueryHandle,
			expectedValue: "test.bsky.social",
			shouldBeEmpty: false,
		},
		{
			name:          "PDS record",
			doc:           &mockDIDDocument,
			qname:         "_pds.test123.plc.atscan.net.",
			queryType:     QueryPDS,
			expectedValue: "https://bsky.social",
			shouldBeEmpty: false,
		},
		{
			name:          "Labeler record",
			doc:           &mockDIDDocument,
			qname:         "_labeler.test123.plc.atscan.net.",
			queryType:     QueryLabeler,
			expectedValue: "https://mod.bsky.app",
			shouldBeEmpty: false,
		},
		{
			name:          "Pubkey record",
			doc:           &mockDIDDocument,
			qname:         "_pubkey.test123.plc.atscan.net.",
			queryType:     QueryPubKey,
			expectedValue: "zQ3shXjHeiBuRCKmM36cuYnm7YEMzhGnCmCyW92sRJ9przSMS",
			shouldBeEmpty: false,
		},
		{
			name: "Empty handle",
			doc: &DIDDocument{
				AlsoKnownAs: []string{},
			},
			qname:         "_handle.test123.plc.atscan.net.",
			queryType:     QueryHandle,
			expectedValue: "",
			shouldBeEmpty: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			records := handler.createTXTRecord(tt.doc, tt.qname, tt.queryType)

			if tt.shouldBeEmpty {
				if len(records) != 0 {
					t.Errorf("expected empty records, got %d records", len(records))
				}
				return
			}

			if len(records) != 1 {
				t.Fatalf("expected 1 record, got %d", len(records))
			}

			txtRecord, ok := records[0].(*dns.TXT)
			if !ok {
				t.Fatal("record is not TXT type")
			}

			if txtRecord.Hdr.Name != tt.qname {
				t.Errorf("expected name %s, got %s", tt.qname, txtRecord.Hdr.Name)
			}

			if len(txtRecord.Txt) != 1 {
				t.Fatalf("expected 1 TXT value, got %d", len(txtRecord.Txt))
			}

			if txtRecord.Txt[0] != tt.expectedValue {
				t.Errorf("expected value %s, got %s", tt.expectedValue, txtRecord.Txt[0])
			}

			if txtRecord.Hdr.Ttl != 300 {
				t.Errorf("expected TTL 300, got %d", txtRecord.Hdr.Ttl)
			}
		})
	}
}

// Test fetching DID document with mock server
func TestFetchDIDDocument(t *testing.T) {
	// Create mock HTTP server
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/did:plc:test123" {
			w.WriteHeader(http.StatusOK)
			json.NewEncoder(w).Encode(mockDIDDocument)
		} else if r.URL.Path == "/did:plc:notfound" {
			w.WriteHeader(http.StatusNotFound)
		} else {
			w.WriteHeader(http.StatusInternalServerError)
		}
	}))
	defer server.Close()

	handler := NewPLCHandler(server.URL)

	tests := []struct {
		name        string
		did         string
		expectError bool
	}{
		{
			name:        "Valid DID",
			did:         "did:plc:test123",
			expectError: false,
		},
		{
			name:        "Not found DID",
			did:         "did:plc:notfound",
			expectError: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			doc, err := handler.fetchDIDDocument(tt.did)

			if tt.expectError {
				if err == nil {
					t.Error("expected error, got nil")
				}
				return
			}

			if err != nil {
				t.Fatalf("unexpected error: %v", err)
			}

			if doc.ID != mockDIDDocument.ID {
				t.Errorf("expected ID %s, got %s", mockDIDDocument.ID, doc.ID)
			}
		})
	}
}

// Test caching
func TestCaching(t *testing.T) {
	callCount := 0

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		callCount++
		w.WriteHeader(http.StatusOK)
		json.NewEncoder(w).Encode(mockDIDDocument)
	}))
	defer server.Close()

	handler := NewPLCHandler(server.URL)

	// First call
	_, err := handler.fetchDIDDocument("did:plc:test123")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if callCount != 1 {
		t.Errorf("expected 1 HTTP call, got %d", callCount)
	}

	// Second call (should use cache)
	_, err = handler.fetchDIDDocument("did:plc:test123")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if callCount != 1 {
		t.Errorf("expected 1 HTTP call (cached), got %d", callCount)
	}

	// Invalidate cache by modifying timestamp
	handler.cache["did:plc:test123"].Timestamp = time.Now().Add(-10 * time.Minute)

	// Third call (cache expired, should fetch again)
	_, err = handler.fetchDIDDocument("did:plc:test123")
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if callCount != 2 {
		t.Errorf("expected 2 HTTP calls (cache expired), got %d", callCount)
	}
}

// Integration test for DNS server
func TestDNSServerIntegration(t *testing.T) {
	// Create mock HTTP server
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		json.NewEncoder(w).Encode(mockDIDDocument)
	}))
	defer server.Close()

	handler := NewPLCHandler(server.URL)

	tests := []struct {
		name          string
		qname         string
		qtype         uint16
		expectedValue string
		expectError   bool
	}{
		{
			name:          "Query handle",
			qname:         "_handle.test123.plc.atscan.net.",
			qtype:         dns.TypeTXT,
			expectedValue: "test.bsky.social",
			expectError:   false,
		},
		{
			name:          "Query PDS",
			qname:         "_pds.test123.plc.atscan.net.",
			qtype:         dns.TypeTXT,
			expectedValue: "https://bsky.social",
			expectError:   false,
		},
		{
			name:          "Query labeler",
			qname:         "_labeler.test123.plc.atscan.net.",
			qtype:         dns.TypeTXT,
			expectedValue: "https://mod.bsky.app",
			expectError:   false,
		},
		{
			name:          "Query pubkey",
			qname:         "_pubkey.test123.plc.atscan.net.",
			qtype:         dns.TypeTXT,
			expectedValue: "zQ3shXjHeiBuRCKmM36cuYnm7YEMzhGnCmCyW92sRJ9przSMS",
			expectError:   false,
		},
		{
			name:        "Invalid query type",
			qname:       "_handle.test123.plc.atscan.net.",
			qtype:       dns.TypeA,
			expectError: true,
		},
		{
			name:        "Invalid domain",
			qname:       "invalid.domain.com.",
			qtype:       dns.TypeTXT,
			expectError: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Create DNS request
			m := new(dns.Msg)
			m.SetQuestion(tt.qname, tt.qtype)

			// Create response writer
			rw := &testResponseWriter{msg: new(dns.Msg)}

			// Handle request
			handler.ServeDNS(rw, m)

			if tt.expectError {
				if rw.msg.Rcode == dns.RcodeSuccess && len(rw.msg.Answer) > 0 {
					t.Error("expected error response, got success")
				}
				return
			}

			if rw.msg.Rcode != dns.RcodeSuccess {
				t.Errorf("expected success, got rcode %d", rw.msg.Rcode)
			}

			if len(rw.msg.Answer) != 1 {
				t.Fatalf("expected 1 answer, got %d", len(rw.msg.Answer))
			}

			txtRecord, ok := rw.msg.Answer[0].(*dns.TXT)
			if !ok {
				t.Fatal("answer is not TXT record")
			}

			if len(txtRecord.Txt) != 1 {
				t.Fatalf("expected 1 TXT value, got %d", len(txtRecord.Txt))
			}

			if txtRecord.Txt[0] != tt.expectedValue {
				t.Errorf("expected value %s, got %s", tt.expectedValue, txtRecord.Txt[0])
			}
		})
	}
}

// Test response writer for DNS testing
type testResponseWriter struct {
	msg *dns.Msg
}

func (w *testResponseWriter) LocalAddr() net.Addr {
	return nil
}

func (w *testResponseWriter) RemoteAddr() net.Addr {
	return nil
}

func (w *testResponseWriter) WriteMsg(m *dns.Msg) error {
	w.msg = m
	return nil
}

func (w *testResponseWriter) Write([]byte) (int, error) {
	return 0, nil
}

func (w *testResponseWriter) Close() error {
	return nil
}

func (w *testResponseWriter) TsigStatus() error {
	return nil
}

func (w *testResponseWriter) TsigTimersOnly(bool) {}

func (w *testResponseWriter) Hijack() {}