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"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"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"
)

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.
type Client struct {
	baseURL    string
	nodeID     string
	token      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
}

// 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 = 1 << 20
	}
	return &Client{
		baseURL:   strings.TrimRight(cfg.BaseURL, "/"),
		nodeID:    cfg.NodeID,
		token:     cfg.Token,
		signer:    cfg.Signer,
		userAgent: cfg.UserAgent,
		maxBytes:  maxBytes,
		http:      &http.Client{Timeout: timeout},
	}, 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.token = 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.session(), 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.session(), false)
}

func (c *Client) session() string { return c.token }

// 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:
		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")

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