Warp/server/internal/federation/discovery_test.go

package federation

import (
	"context"
	"errors"
	"net"
	"net/http"
	"net/http/httptest"
	"slices"
	"strings"
	"sync/atomic"
	"testing"
	"time"
)

func TestParseServerName(t *testing.T) {
	cases := []struct {
		in   string
		want ServerName
		bad  bool
	}{
		{in: "example.org", want: ServerName{Host: "example.org"}},
		{in: "example.org:8448", want: ServerName{Host: "example.org", Port: 8448}},
		{in: "1.2.3.4", want: ServerName{Host: "1.2.3.4", IsIP: true}},
		{in: "1.2.3.4:80", want: ServerName{Host: "1.2.3.4", Port: 80, IsIP: true}},
		{in: "[::1]", want: ServerName{Host: "::1", IsIP: true}},
		{in: "[2001:db8::1]:8448", want: ServerName{Host: "2001:db8::1", Port: 8448, IsIP: true}},
		{in: "", bad: true},
		{in: "::1", bad: true},
		{in: "example.org:", bad: true},
		{in: "example.org:0", bad: true},
		{in: "example.org:70000", bad: true},
		{in: "exa mple.org", bad: true},
		{in: "[1.2.3.4]", bad: true},
		{in: "[::1", bad: true},
	}
	for _, tc := range cases {
		got, err := ParseServerName(tc.in)
		if tc.bad {
			if !errors.Is(err, ErrInvalidServerName) {
				t.Errorf("ParseServerName(%q) = %+v, %v; want an error", tc.in, got, err)
			}
			continue
		}
		if err != nil || got != tc.want {
			t.Errorf("ParseServerName(%q) = %+v, %v; want %+v", tc.in, got, err, tc.want)
		}
	}
}

// fakeDNS answers SRV lookups from a table.
type fakeDNS map[string][]*net.SRV

func (f fakeDNS) lookup(_ context.Context, service, proto, name string) (string, []*net.SRV, error) {
	records, ok := f["_"+service+"._"+proto+"."+name]
	if !ok {
		return "", nil, &net.DNSError{Err: "no such host", Name: name, IsNotFound: true}
	}
	return "", records, nil
}

// wellKnownServer serves /.well-known/matrix/server for any Host, from a
// table, counting requests.
func wellKnownServer(t *testing.T, docs map[string]string, hits *atomic.Int32) *Resolver {
	t.Helper()
	srv := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		hits.Add(1)
		doc, ok := docs[strings.Split(r.Host, ":")[0]]
		if !ok || r.URL.Path != "/.well-known/matrix/server" {
			http.NotFound(w, r)
			return
		}
		w.Header().Set("Cache-Control", "max-age=60")
		_, _ = w.Write([]byte(doc))
	}))
	t.Cleanup(srv.Close)

	// Every hostname is dialled at the test server, whose certificate the
	// client trusts, so the well-known fetch is real HTTPS to a fake name.
	transport := srv.Client().Transport.(*http.Transport).Clone()
	addr := srv.Listener.Addr().String()
	transport.DialContext = func(ctx context.Context, network, _ string) (net.Conn, error) {
		return (&net.Dialer{}).DialContext(ctx, network, addr)
	}
	transport.TLSClientConfig.InsecureSkipVerify = true //nolint:gosec // test server with a fixed certificate
	r := NewResolver(&http.Client{Transport: transport})
	return r
}

func TestResolve(t *testing.T) {
	var hits atomic.Int32
	r := wellKnownServer(t, map[string]string{
		"delegated.example": `{"m.server":"matrix.delegated.example"}`,
		"withport.example":  `{"m.server":"matrix.withport.example:443"}`,
		"toip.example":      `{"m.server":"10.0.0.1"}`,
		"broken.example":    `{"m.server":"not a name"}`,
	}, &hits)
	r.LookupSRV = fakeDNS{
		"_matrix-fed._tcp.matrix.delegated.example": {{Target: "fed.delegated.example.", Port: 8443}},
		"_matrix-fed._tcp.srv.example":              {{Target: "a.srv.example.", Port: 1}, {Target: "b.srv.example.", Port: 2}},
		"_matrix._tcp.legacy.example":               {{Target: "old.legacy.example.", Port: 9}},
	}.lookup

	cases := []struct {
		name      string
		addresses []string
		host      string
	}{
		{"1.2.3.4", []string{"1.2.3.4:8448"}, "1.2.3.4"},
		{"[::1]:9000", []string{"[::1]:9000"}, "[::1]:9000"},
		{"explicit.example:1234", []string{"explicit.example:1234"}, "explicit.example:1234"},
		// Delegation by well-known, then SRV on the delegated name: the
		// SRV target is dialled, the delegated name is what TLS checks.
		{"delegated.example", []string{"fed.delegated.example:8443"}, "matrix.delegated.example"},
		{"withport.example", []string{"matrix.withport.example:443"}, "matrix.withport.example:443"},
		{"toip.example", []string{"10.0.0.1:8448"}, "10.0.0.1"},
		{"broken.example", []string{"broken.example:8448"}, "broken.example"},
		{"srv.example", []string{"a.srv.example:1", "b.srv.example:2"}, "srv.example"},
		{"legacy.example", []string{"old.legacy.example:9"}, "legacy.example"},
		{"plain.example", []string{"plain.example:8448"}, "plain.example"},
	}
	for _, tc := range cases {
		got, err := r.Resolve(t.Context(), tc.name)
		if err != nil {
			t.Errorf("Resolve(%s): %v", tc.name, err)
			continue
		}
		if !slices.Equal(got.Addresses, tc.addresses) || got.Host != tc.host {
			t.Errorf("Resolve(%s) = %+v, want %v via Host %s", tc.name, got, tc.addresses, tc.host)
		}
	}

	if got, _ := r.Resolve(t.Context(), "delegated.example"); got.TLSServerName() != "matrix.delegated.example" {
		t.Errorf("TLSServerName = %q", got.TLSServerName())
	}
	if got, _ := r.Resolve(t.Context(), "[::1]:9000"); got.TLSServerName() != "::1" {
		t.Errorf("TLSServerName of an IPv6 literal = %q", got.TLSServerName())
	}
}

func TestWellKnownIsCached(t *testing.T) {
	var hits atomic.Int32
	r := wellKnownServer(t, map[string]string{"cached.example": `{"m.server":"other.example:443"}`}, &hits)
	r.LookupSRV = fakeDNS{}.lookup
	now := time.Unix(1_700_000_000, 0)
	r.Now = func() time.Time { return now }

	for range 3 {
		if _, err := r.Resolve(t.Context(), "cached.example"); err != nil {
			t.Fatal(err)
		}
	}
	// A server with no document is cached too, as a failure.
	for range 3 {
		if _, err := r.Resolve(t.Context(), "missing.example"); err != nil {
			t.Fatal(err)
		}
	}
	if n := hits.Load(); n != 2 {
		t.Errorf("%d well-known fetches, want one per name", n)
	}

	now = now.Add(2 * time.Minute) // past max-age=60
	if _, err := r.Resolve(t.Context(), "cached.example"); err != nil {
		t.Fatal(err)
	}
	if n := hits.Load(); n != 3 {
		t.Errorf("%d well-known fetches after expiry, want 3", n)
	}
}

func TestCacheTTL(t *testing.T) {
	cases := map[string]time.Duration{
		"":                    wellKnownDefaultTTL,
		"max-age=3600":        time.Hour,
		"public, max-age=120": 2 * time.Minute,
		"max-age=999999999":   wellKnownMaxTTL,
		"max-age=nonsense":    wellKnownDefaultTTL,
	}
	for header, want := range cases {
		if got := cacheTTL(header); got != want {
			t.Errorf("cacheTTL(%q) = %v, want %v", header, got, want)
		}
	}
}