msa-agent/internal/transport/transport_test.go

package transport

import (
	"context"
	"crypto/ed25519"
	"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")
	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)
	}
}

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)
	}
}