msa-agent/internal/transport/transport_test.go

package transport

import (
	"context"
	"crypto/ed25519"
	"crypto/x509"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"
)

// The relay is represented by a real local HTTP server in these tests, so the TLS requirement is
// satisfied by rewriting the URL rather than by weakening the check: what is under test is the
// handshake and the status mapping, not the certificate.
func newClient(t *testing.T, handler http.Handler, token string) (*Client, *httptest.Server) {
	t.Helper()
	server := httptest.NewServer(handler)
	t.Cleanup(server.Close)
	_, signer, err := ed25519.GenerateKey(nil)
	if err != nil {
		t.Fatal(err)
	}
	client, err := New(Config{
		BaseURL:   strings.Replace(server.URL, "http://", "https://", 1),
		NodeID:    "nd_test0001",
		Token:     token,
		Signer:    signer,
		Timeout:   2 * time.Second,
		UserAgent: "msa-agent/test",
	})
	if err != nil {
		t.Fatal(err)
	}
	// Point the client at the plain-HTTP test server. Only the scheme check is bypassed here; every
	// other behaviour is the production one.
	client.http.Transport = server.Client().Transport
	client.baseURL = server.URL
	return client, server
}

func TestNewRejectsAnythingButAnHTTPSRelay(t *testing.T) {
	_, signer, err := ed25519.GenerateKey(nil)
	if err != nil {
		t.Fatal(err)
	}
	for _, base := range []string{
		"http://relay.example",
		"ws://relay.example",
		"https://relay.example/some/path",
		"https://",
		"not a url",
	} {
		if _, err := New(Config{BaseURL: base, NodeID: "nd_x", Token: "t", Signer: signer}); err == nil {
			t.Fatalf("accepted relay url %q", base)
		}
	}
	if _, err := New(Config{BaseURL: "https://relay.example", NodeID: "nd_x", Token: "t", Signer: signer}); err != nil {
		t.Fatalf("rejected a valid relay url: %v", err)
	}
	if _, err := New(Config{BaseURL: "https://relay.example", NodeID: "nd_x", Token: "t"}); err == nil {
		t.Fatal("accepted a missing signing key")
	}
}

// The handshake is challenge-response, so a stolen bearer token alone is not enough to impersonate
// an agent. This test is the one that would fail if someone replaced the signature with the token.
func TestHelloSignsTheChallengeBoundToTheNodeID(t *testing.T) {
	var (
		requests []HelloRequest
		issued   = "challenge-abc"
	)
	client, _ := newClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.Header.Get("Authorization") == "" {
			t.Error("handshake was sent without a bearer token")
		}
		var body HelloRequest
		if err := json.NewDecoder(r.Body).Decode(&body); err != nil {
			t.Fatal(err)
		}
		requests = append(requests, body)
		out := HelloResponse{}
		if body.Challenge == "" {
			out.Challenge = issued
		} else {
			out.Session = Session{Token: "session-token", ExpiresIn: 3600}
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(out)
	}), "bootstrap-token")

	response, err := client.Hello(context.Background(), "test-build", "deadbeef")
	if err != nil {
		t.Fatalf("Hello: %v", err)
	}
	if response.Session.Token != "session-token" {
		t.Fatalf("session token is %q", response.Session.Token)
	}
	if len(requests) != 2 {
		t.Fatalf("handshake took %d round trips, want 2", len(requests))
	}
	if requests[0].Challenge != "" {
		t.Fatal("the first handshake request must not carry a challenge")
	}
	if requests[1].Challenge != issued {
		t.Fatalf("the agent echoed challenge %q, want %q", requests[1].Challenge, issued)
	}
	public := client.signer.Public().(ed25519.PublicKey)
	if !ed25519.Verify(public, challengeBytes("nd_test0001", issued), requests[1].Signature) {
		t.Fatal("the challenge signature does not verify")
	}
	// Bound to the node id: the same signature must not verify under a different node.
	if ed25519.Verify(public, challengeBytes("nd_other001", issued), requests[1].Signature) {
		t.Fatal("the challenge signature is not bound to the node id")
	}
}

func TestStatusMapping(t *testing.T) {
	cases := []struct {
		status int
		want   error
	}{
		{http.StatusUnauthorized, ErrUnauthorized},
		{http.StatusNoContent, ErrNoWork},
		{http.StatusRequestTimeout, ErrNoWork},
		{http.StatusTooManyRequests, ErrRateLimited},
	}
	for _, c := range cases {
		client, _ := newClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.WriteHeader(c.status)
		}), "t")
		_, err := client.Poll(context.Background())
		if err == nil || !errors.Is(err, c.want) {
			t.Fatalf("status %d: got %v, want %v", c.status, err, c.want)
		}
	}
}

func TestPollReturnsTheEnvelopesTheRelayIsHolding(t *testing.T) {
	want := []string{`{"v":1,"node":"nd_test0001"}`, `{"v":1,"node":"nd_test0001"}`}
	client, _ := newClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != PathPoll {
			t.Fatalf("polled %q", r.URL.Path)
		}
		w.Header().Set("Content-Type", "application/json")
		json.NewEncoder(w).Encode(Batch{Envelopes: []json.RawMessage{
			json.RawMessage(want[0]), json.RawMessage(want[1]),
		}})
	}), "t")
	batch, err := client.Poll(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	if len(batch.Envelopes) != 2 {
		t.Fatalf("got %d envelopes, want 2", len(batch.Envelopes))
	}
	for i := range want {
		if string(batch.Envelopes[i]) != want[i] {
			t.Fatalf("envelope %d is %s", i, batch.Envelopes[i])
		}
	}
}

// An empty relay answer is the common case for a long poll, and must not be an error or the agent
// would back off every time nothing happened.
func TestAnEmptyPollIsNotAnError(t *testing.T) {
	client, _ := newClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}), "t")
	if _, err := client.Poll(context.Background()); !errors.Is(err, ErrNoWork) {
		t.Fatalf("got %v, want %v", err, ErrNoWork)
	}
}

func TestResponsesAreBoundedBeforeDecoding(t *testing.T) {
	client, _ := newClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Write([]byte(strings.Repeat("x", 4096)))
	}), "t")
	client.maxBytes = 1024
	if _, err := client.Poll(context.Background()); err == nil {
		t.Fatal("an oversized relay response was accepted")
	}
}

func TestResultPostsTheEnvelopeVerbatim(t *testing.T) {
	var got json.RawMessage
	var path string
	client, _ := newClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		path = r.URL.Path
		got, _ = io.ReadAll(r.Body)
		w.WriteHeader(http.StatusOK)
	}), "t")
	// The session the handshake would have issued. It is set directly rather than by running the
	// handshake because this test is about the result body, and the two tokens being distinct is
	// asserted on its own in TestTheBearerAndTheSessionAreDifferentCredentials.
	client.session = "session-token"
	envelope := json.RawMessage(`{"v":1,"node":"nd_test0001","p":"AAAA"}`)
	if err := client.Result(context.Background(), envelope); err != nil {
		t.Fatal(err)
	}
	if path != PathResult {
		t.Fatalf("posted to %q", path)
	}
	if string(got) != string(envelope) {
		t.Fatalf("relay received %s, want %s", got, envelope)
	}
}

// The device registration posts the node's authoritative paired set under the session, which is
// what lets the relay learn about freshly paired phones. Empty is posted too: "nobody" is a
// revocation claim, not a no-op.
func TestPushDevicesPostsTheFullSetUnderTheSession(t *testing.T) {
	var got map[string]any
	client, _ := newClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != PathDevices {
			t.Fatalf("posted to %q", r.URL.Path)
		}
		if r.Header.Get("X-MSA-Session") != "session-token" {
			t.Errorf("the device push carried session %q", r.Header.Get("X-MSA-Session"))
		}
		if r.Header.Get("Authorization") != "Bearer t" {
			t.Errorf("the device push carried authorization %q", r.Header.Get("Authorization"))
		}
		json.NewDecoder(r.Body).Decode(&got)
		w.WriteHeader(http.StatusOK)
	}), "t")
	client.session = "session-token"
	err := client.PushDevices(context.Background(), []DeviceRegistration{
		{App: "ap_5b1d0000", Token: "tok-a"},
	})
	if err != nil {
		t.Fatal(err)
	}
	if got["node"] != "nd_test0001" {
		t.Fatalf("node is %v", got["node"])
	}
	devices, _ := got["devices"].([]any)
	if len(devices) != 1 {
		t.Fatalf("devices: %v", got["devices"])
	}
	// An empty list is a real claim, not a skipped call.
	client2, _ := newClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}), "t")
	client2.session = "session-token"
	if err := client2.PushDevices(context.Background(), []DeviceRegistration{}); err != nil {
		t.Fatalf("empty push: %v", err)
	}
}

func TestRetryAfterHasASaneDefault(t *testing.T) {
	client, _ := newClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}), "t")
	if got := client.RetryAfter(); got <= 0 || got > 10*time.Second {
		t.Fatalf("RetryAfter() = %v", got)
	}
}

// The relay's Retry-After is an instruction, and the loop honours it; a value the agent cannot
// parse falls back to the default rather than becoming an unbounded wait or an immediate retry.
func TestRetryAfterHonoursTheRelayHeader(t *testing.T) {
	client, _ := newClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Retry-After", "7")
		w.WriteHeader(http.StatusTooManyRequests)
	}), "t")
	if _, err := client.Poll(context.Background()); !errors.Is(err, ErrRateLimited) {
		t.Fatalf("got %v, want ErrRateLimited", err)
	}
	if got := client.RetryAfter(); got != 7*time.Second {
		t.Fatalf("RetryAfter() = %v, want 7s", got)
	}
	for _, junk := range []string{"", "Wed, 21 Oct 2015 07:28:00 GMT", "-3", "999999"} {
		if got := parseRetryAfter(junk); got != 2*time.Second {
			t.Errorf("parseRetryAfter(%q) = %v, want the default", junk, got)
		}
	}
}

// The bearer token and the session token are different credentials with different jobs, and swapping
// them is the failure this test exists to catch. The bearer identifies the node to the relay's
// authentication layer; the session is what the relay issued once the challenge signature proved
// possession of the node's key. Sending the session as the bearer (or vice versa) makes the relay
// answer 401, which the loop reads as "re-authenticate" — so the symptom is an agent that connects
// and immediately reconnects forever, with a log line that blames the token.
func TestTheBearerAndTheSessionAreDifferentCredentials(t *testing.T) {
	var seen []struct {
		path          string
		authorization string
		session       string
	}
	client, _ := newClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		seen = append(seen, struct {
			path          string
			authorization string
			session       string
		}{r.URL.Path, r.Header.Get("Authorization"), r.Header.Get("X-MSA-Session")})
		w.Header().Set("Content-Type", "application/json")
		switch r.URL.Path {
		case PathHello:
			var body HelloRequest
			json.NewDecoder(r.Body).Decode(&body)
			if body.Challenge == "" {
				json.NewEncoder(w).Encode(HelloResponse{Challenge: "issued"})
			} else {
				json.NewEncoder(w).Encode(HelloResponse{Session: Session{Token: "session-abc"}})
			}
		case PathPoll:
			w.WriteHeader(http.StatusNoContent)
		default:
			w.WriteHeader(http.StatusOK)
		}
	}), "bearer-token")

	if _, err := client.Hello(context.Background(), "build", "sha"); err != nil {
		t.Fatalf("Hello: %v", err)
	}
	if _, err := client.Poll(context.Background()); !errors.Is(err, ErrNoWork) {
		t.Fatalf("Poll: %v", err)
	}

	for _, request := range seen {
		if request.authorization != "Bearer bearer-token" {
			t.Errorf("%s sent authorization %q", request.path, request.authorization)
		}
		switch request.path {
		case PathHello:
			// The handshake is authenticated by the bearer and the signature; there is no session yet.
			if request.session != "" {
				t.Errorf("the handshake carried a session %q", request.session)
			}
		case PathPoll:
			if request.session != "session-abc" {
				t.Errorf("poll carried session %q, want the one the handshake issued", request.session)
			}
		}
	}
}

// A pinned CA replaces the system store rather than adding to it, and a certificate that only the
// system store would accept must then fail. That asymmetry is the point: `relayCAFile` is for a
// relay behind a private CA, and an operator who set it expects the relay to be verified against
// exactly that root.
func TestRootCAsReplaceTheSystemStore(t *testing.T) {
	server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	}))
	t.Cleanup(server.Close)

	_, signer, err := ed25519.GenerateKey(nil)
	if err != nil {
		t.Fatal(err)
	}
	// The test server's own certificate is self-signed, so trusting it works...
	pool := x509.NewCertPool()
	pool.AddCert(server.Certificate())
	client, err := New(Config{
		BaseURL: server.URL, NodeID: "nd_test0001", Token: "t", Signer: signer,
		Timeout: 2 * time.Second, RootCAs: pool,
	})
	if err != nil {
		t.Fatal(err)
	}
	// The transport is left as configured, so this exercises the pinned pool rather than the test
	// server's convenience client. The handler answers 200 with no body, so ErrNoWork is the success
	// case here: what is under test is whether the TLS handshake completed, and a verification
	// failure surfaces as a transport error instead.
	if _, err := client.Poll(context.Background()); !errors.Is(err, ErrNoWork) {
		t.Fatalf("a pinned CA did not verify the relay: %v", err)
	}

	// ...and an empty pool verifies nothing, which is the failure mode an operator sees when they
	// point relayCAFile at the wrong file: a TLS error, never a silent fallback.
	empty, err := New(Config{
		BaseURL: server.URL, NodeID: "nd_test0001", Token: "t", Signer: signer,
		Timeout: 2 * time.Second, RootCAs: x509.NewCertPool(),
	})
	if err != nil {
		t.Fatal(err)
	}
	if _, err := empty.Poll(context.Background()); err == nil || errors.Is(err, ErrNoWork) {
		t.Fatalf("an empty root pool still verified the relay: %v", err)
	}
}