msa-agent/internal/transport/transport.go

// Package transport is the agent's only outbound channel.
//
// It is long-polled HTTPS rather than a WebSocket, which is a deliberate regression in elegance for
// a gain in dependencies: the agent must build with nothing but the standard library so an operator
// can cross-compile it for an ARM node from any machine, and a WebSocket client is a few hundred
// lines of framing, masking, ping/pong and close handling that would then need its own tests.
//
// The direction is the security property. The agent dials out and the relay answers on the
// connection the agent opened, so the node has no listening socket and no routable name: there is
// nothing on the internal network for anyone to find.
package transport

import (
	"bytes"
	"context"
	"crypto/ed25519"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"time"
)

// Paths on the relay. Versioned, because an agent and a relay will not be upgraded on the same day.
const (
	PathHello   = "/v1/agent/hello"
	PathPoll    = "/v1/agent/poll"
	PathResult  = "/v1/agent/result"
	PathDevices = "/v1/agent/devices"
)

// DefaultMaxBytes bounds a relay response body.
//
// It has to admit a full long-poll batch, which is the largest thing the relay can legitimately send:
// one envelope per mailbox slot. The previous 1 MiB was chosen as "generous" without being derived
// from anything, and it is smaller than what the relay is now allowed to send — a mailbox of eight
// envelopes at the policy's ceiling is about 4 MiB, so a full batch would have been refused as
// "relay response exceeds 1048576 bytes", and the agent would have retried a poll the relay answered
// correctly. That is the same class of defect as the relay's own envelope ceiling: a bound picked by
// feel, on the other side of a seam, disagreeing with the party that produces the data.
//
// The derivation is stated here rather than computed from the policy because the agent reads the
// policy for its own limits; the two numbers are pinned against each other by
// test/limits.test.js, which is where a change on either side becomes visible.
const DefaultMaxBytes int64 = 8 << 20

var (
	// ErrSessionExpired means the relay no longer recognises this session. The caller re-runs the
	// handshake and polls again; it is not an error worth crashing over.
	ErrSessionExpired = errors.New("transport: agent session expired")
	// ErrUnauthorized means the relay rejected the bearer token. Retrying immediately is pointless,
	// so the caller backs off.
	ErrUnauthorized = errors.New("transport: relay rejected the agent credential")
	// ErrNoWork is the long poll's empty result. It is not an error and carries no body.
	ErrNoWork = errors.New("transport: no work")
)

// Session is the state the handshake returns and the poll then presents.
type Session struct {
	Token     string `json:"token"`
	ExpiresIn int    `json:"expiresIn"`
}

// HelloRequest is the agent's opening move.
type HelloRequest struct {
	Node string `json:"node"`
	// Challenge is the relay's nonce, signed with K_agent_sign. Possession of the bearer token is
	// not enough to impersonate an agent: the token limits abuse, the signature proves identity.
	Challenge   string `json:"challenge"`
	Signature   []byte `json:"signature"`
	AgentBuild  string `json:"agentBuild,omitempty"`
	PolicySHA   string `json:"policySha256"`
	ProtocolVer int    `json:"protocolVersion"`
}

// HelloResponse is what the relay answers with.
type HelloResponse struct {
	Session    Session `json:"session"`
	Challenge  string  `json:"challenge,omitempty"`
	PollWaitMs int     `json:"pollWaitMs,omitempty"`
}

// Batch is one long-poll result: zero or more request envelopes waiting for this node.
type Batch struct {
	Envelopes []json.RawMessage `json:"envelopes"`
}

// Client talks to one relay.
//
// The two credentials are kept apart on purpose, and keeping them apart is not cosmetic. The bearer
// token identifies the node to the relay's authentication layer; the session token is what the relay
// issued after the challenge signature proved possession of the node's key. The relay checks both on
// every poll, and using one where the other belongs fails as "relay rejected the agent credential" —
// which reads like a bad token rather than like a mixed-up header.
type Client struct {
	baseURL    string
	nodeID     string
	token      string
	session    string
	signer     ed25519.PrivateKey
	http       *http.Client
	userAgent  string
	maxBytes   int64
	retryAfter time.Duration
}

// Config builds a Client.
type Config struct {
	BaseURL   string
	NodeID    string
	Token     string
	Signer    ed25519.PrivateKey
	Timeout   time.Duration
	UserAgent string
	MaxBytes  int64
	// RootCAs replaces the system trust store when non-nil. It exists for one deployment shape: a
	// relay behind a private CA, where installing the CA into the node's system store is either not
	// possible or not desirable. The zero value keeps the system store, which is what an operator
	// using a public certificate wants and what the agent does by default.
	//
	// It is a *CertPool and not an "insecure" flag on purpose: the way to trust a private relay is to
	// pin the CA that signed it, never to stop verifying.
	RootCAs *x509.CertPool
}

// New validates the relay URL before any traffic is sent.
//
// HTTPS is required, and not negotiable. The relay is untrusted, which means the channel's job is to
// stop a third party from substituting a different ciphertext; an unauthenticated channel would make
// the endpoint-to-endpoint encryption above it decorative.
func New(cfg Config) (*Client, error) {
	parsed, err := url.Parse(cfg.BaseURL)
	if err != nil {
		return nil, fmt.Errorf("transport: relay url: %w", err)
	}
	if parsed.Scheme != "https" {
		return nil, fmt.Errorf("transport: relay url must be https, got %q", parsed.Scheme)
	}
	if parsed.Host == "" {
		return nil, errors.New("transport: relay url has no host")
	}
	if parsed.Path != "" && parsed.Path != "/" {
		return nil, fmt.Errorf("transport: relay url must not carry a path, got %q", parsed.Path)
	}
	if len(cfg.Signer) != ed25519.PrivateKeySize {
		return nil, errors.New("transport: agent signing key is the wrong size")
	}
	timeout := cfg.Timeout
	if timeout <= 0 {
		timeout = 90 * time.Second
	}
	maxBytes := cfg.MaxBytes
	if maxBytes <= 0 {
		maxBytes = DefaultMaxBytes
	}
	client := &http.Client{Timeout: timeout}
	if cfg.RootCAs != nil {
		// A fresh transport rather than a mutated DefaultTransport: mutating the shared one would
		// change trust for every other client in the process, which is a surprising thing for a
		// constructor to do.
		client.Transport = &http.Transport{
			TLSClientConfig: &tls.Config{RootCAs: cfg.RootCAs, MinVersion: tls.VersionTLS12},
		}
	}
	return &Client{
		baseURL:   strings.TrimRight(cfg.BaseURL, "/"),
		nodeID:    cfg.NodeID,
		token:     cfg.Token,
		signer:    cfg.Signer,
		userAgent: cfg.UserAgent,
		maxBytes:  maxBytes,
		http:      client,
	}, nil
}

// Hello performs the challenge-response handshake and stores the session token.
//
// It is called again on expiry, so it must be safe to repeat: no state is carried in beyond the
// signing key and the token.
func (c *Client) Hello(ctx context.Context, build, policySHA string) (HelloResponse, error) {
	var out HelloResponse
	// The first call has no challenge to sign, so the relay is asked for one; the second completes
	// it. Two round trips on startup is an acceptable price for not having the relay mint a nonce
	// the agent never saw.
	err := c.post(ctx, PathHello, HelloRequest{Node: c.nodeID, PolicySHA: policySHA, AgentBuild: build}, &out, "", false)
	if err != nil || out.Challenge == "" {
		return out, err
	}
	signed := HelloRequest{
		Node:        c.nodeID,
		Challenge:   out.Challenge,
		Signature:   ed25519.Sign(c.signer, challengeBytes(c.nodeID, out.Challenge)),
		AgentBuild:  build,
		PolicySHA:   policySHA,
		ProtocolVer: 1,
	}
	var completed HelloResponse
	if err := c.post(ctx, PathHello, signed, &completed, "", false); err != nil {
		return completed, err
	}
	c.session = completed.Session.Token
	return completed, nil
}

// challengeBytes is what the signature covers. Binding the node id in means a signature collected
// from one agent's handshake cannot be replayed by another.
func challengeBytes(nodeID, challenge string) []byte {
	return []byte("msa-relay-v1-hello\x00" + nodeID + "\x00" + challenge)
}

// Poll waits for work. A nil error with an empty batch means the long poll timed out with nothing,
// which is the normal case and must not be treated as a failure.
func (c *Client) Poll(ctx context.Context) (Batch, error) {
	var batch Batch
	// A long poll that times out answers with an empty body, and that is the normal case rather than
	// a failure, so an empty body means ErrNoWork here and only here.
	err := c.post(ctx, PathPoll, map[string]string{"node": c.nodeID}, &batch, c.sessionToken(), true)
	if err != nil {
		return Batch{}, err
	}
	return batch, nil
}

// Result posts one response envelope back, correlated by request id.
//
// An empty answer here is an acknowledgement, not "no work". The endpoints share a transport and
// must not share that reading of an empty body, or a confirmed delivery would be treated as a
// reason to back off.
func (c *Client) Result(ctx context.Context, envelope json.RawMessage) error {
	return c.post(ctx, PathResult, envelope, &json.RawMessage{}, c.sessionToken(), false)
}

// session returns the token the relay issued for this session. It is empty before the handshake, and
// an empty value makes the relay answer 401, which the loop treats as "re-authenticate".
func (c *Client) sessionToken() string { return c.session }

// DeviceRegistration is one paired phone as the relay needs to know it: the bearer token the phone
// presents, minted by the agent at pairing time and never seen by the relay before that moment.
type DeviceRegistration struct {
	App   string `json:"app"`
	Token string `json:"token"`
}

// PushDevices reports the node's complete paired-device set after the handshake.
//
// Full set, not a delta: the relay replaces everything it holds for this node, which is what makes
// revocation work without a second channel. Destroying the device list (or re-pairing) and letting
// the next hello carry the smaller set removes the revoked bearer from the relay's routing, and an
// agent that never reconnects cannot revoke anything but is also unable to command itself.
func (c *Client) PushDevices(ctx context.Context, devices []DeviceRegistration) error {
	return c.post(ctx, PathDevices, map[string]any{
		"node": c.nodeID, "devices": devices,
	}, &json.RawMessage{}, c.sessionToken(), false)
}

// post is the single place a relay request is made, so authentication, size limits, status mapping
// and body draining cannot drift between the three endpoints.
//
// emptyIsNoWork is per-endpoint on purpose: folding "no body" into one meaning across all three
// endpoints would make an acknowledgement indistinguishable from an idle poll.
func (c *Client) post(ctx context.Context, path string, body any, out any, session string, emptyIsNoWork bool) error {
	encoded, err := json.Marshal(body)
	if err != nil {
		return fmt.Errorf("transport: encode: %w", err)
	}
	request, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+path, bytes.NewReader(encoded))
	if err != nil {
		return fmt.Errorf("transport: request: %w", err)
	}
	request.Header.Set("Content-Type", "application/json")
	// The bearer token grants the right to be rate limited, nothing more: it authorises no command,
	// because commands are authorised by the phone's signature inside the envelope.
	request.Header.Set("Authorization", "Bearer "+c.token)
	if session != "" {
		request.Header.Set("X-MSA-Session", session)
	}
	if c.userAgent != "" {
		request.Header.Set("User-Agent", c.userAgent)
	}
	response, err := c.http.Do(request)
	if err != nil {
		return fmt.Errorf("transport: %s: %w", path, err)
	}
	defer func() {
		io.Copy(io.Discard, response.Body)
		response.Body.Close()
	}()

	switch response.StatusCode {
	case http.StatusOK:
	case http.StatusUnauthorized:
		return ErrUnauthorized
	case http.StatusNoContent, http.StatusRequestTimeout:
		// 204 and 408 are how the relay says "nothing yet", so the agent loops straight back.
		return ErrNoWork
	case http.StatusTooManyRequests:
		retry := parseRetryAfter(response.Header.Get("Retry-After"))
		c.retryAfter = retry
		return fmt.Errorf("%w: retry after %s", ErrRateLimited, response.Header.Get("Retry-After"))
	default:
		return fmt.Errorf("transport: %s: relay returned %d", path, response.StatusCode)
	}
	// The response is bounded before it is decoded: an unbounded read is a remote memory decision.
	limited, err := io.ReadAll(io.LimitReader(response.Body, c.maxBytes+1))
	if err != nil {
		return fmt.Errorf("transport: read: %w", err)
	}
	if int64(len(limited)) > c.maxBytes {
		return fmt.Errorf("transport: relay response exceeds %d bytes", c.maxBytes)
	}
	if len(bytes.TrimSpace(limited)) == 0 {
		if emptyIsNoWork {
			return ErrNoWork
		}
		return nil
	}
	if err := json.Unmarshal(limited, out); err != nil {
		return fmt.Errorf("transport: decode: %w", err)
	}
	return nil
}

// ErrRateLimited is returned when the relay asks the agent to slow down.
var ErrRateLimited = errors.New("transport: rate limited by the relay")

// parseRetryAfter reads an RFC 9110 delay-seconds value. Dates are not supported and fall back to
// the default: a clock-skewed node should not be waiting on an arithmetic date subtraction.
func parseRetryAfter(value string) time.Duration {
	seconds, err := strconv.Atoi(strings.TrimSpace(value))
	if err != nil || seconds < 0 || seconds > 3600 {
		return 2 * time.Second
	}
	return time.Duration(seconds) * time.Second
}

// RetryAfter reports how long the agent should wait before its next poll, honouring the relay's
// Retry-After when it gave one.
func (c *Client) RetryAfter() time.Duration {
	if c.retryAfter <= 0 {
		return 2 * time.Second
	}
	return c.retryAfter
}