mesh-controller/internal/licences/refresh_test.go

package licences

import (
	"context"
	"crypto/ecdh"
	"crypto/rand"
	"encoding/base64"
	"strings"
	"testing"

	"golang.org/x/crypto/nacl/box"

	"github.com/novox/mesh-control/internal/licences/adapters"
	"github.com/novox/mesh-control/internal/secrets"
)

// nodeKeyPair is a node's key as the node would hold it: the public half the mesh seals to, and an
// open closure holding the private half the mesh never sees.
func nodeKeyPair(t *testing.T) (public string, open func(string) ([]byte, error)) {
	t.Helper()
	priv, err := ecdh.X25519().GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	var pub, sk [32]byte
	copy(pub[:], priv.PublicKey().Bytes())
	copy(sk[:], priv.Bytes())
	return base64.StdEncoding.EncodeToString(priv.PublicKey().Bytes()),
		func(sealed string) ([]byte, error) {
			blob, err := base64.StdEncoding.DecodeString(sealed)
			if err != nil {
				return nil, err
			}
			out, ok := box.OpenAnonymous(nil, blob, &pub, &sk)
			if !ok {
				return nil, context.Canceled // any error; the test only checks success/failure
			}
			return out, nil
		}
}

// managerPair is like nodeKeyPair but also returns the private key string, because the manager needs
// to OpenAtRest its own refresh token — the one node that reads it back.
func managerPair(t *testing.T) (public, private string) {
	t.Helper()
	priv, err := ecdh.X25519().GenerateKey(rand.Reader)
	if err != nil {
		t.Fatal(err)
	}
	return base64.StdEncoding.EncodeToString(priv.PublicKey().Bytes()),
		base64.StdEncoding.EncodeToString(priv.Bytes())
}

type fakeRefresher struct {
	access    string
	newAtRest *secrets.AtRest
	got       adapters.RefreshInput
}

func (f *fakeRefresher) Refresh(_ context.Context, in adapters.RefreshInput) (adapters.RefreshResult, error) {
	f.got = in
	return adapters.RefreshResult{AccessToken: f.access, NewAtRest: f.newAtRest}, nil
}

// A refreshable-grant licence set up end to end: a manager, a refresh token sealed at rest to it, two
// holders each with a sealing key, and a fake vendor refresher plugged in.
func aRefreshableLicence(t *testing.T, held *Licences, ctx context.Context, fake *fakeRefresher) (
	managerPub, managerPriv string, holders map[string]func(string) ([]byte, error), keys SealingKeys,
) {
	t.Helper()
	adapters.RegisterRefresher("anthropic", fake)
	t.Cleanup(func() { adapters.RegisterRefresher("anthropic", nil) })

	if err := held.Add(ctx, "personal", "anthropic", map[string]any{"model": "a-model"}); err != nil {
		t.Fatal(err)
	}
	if err := held.SetManager(ctx, "personal", "workstation"); err != nil {
		t.Fatal(err)
	}

	managerPub, managerPriv = managerPair(t)
	at, err := secrets.SealAtRest("rt-the-refresh-token", managerPub)
	if err != nil {
		t.Fatal(err)
	}
	if err := held.SetRefreshGrant(ctx, "personal", at); err != nil {
		t.Fatal(err)
	}

	holders = map[string]func(string) ([]byte, error){}
	pub := map[string]string{}
	for _, node := range []string{"workstation", "laptop"} {
		p, open := nodeKeyPair(t)
		pub[node], holders[node] = p, open
		if err := held.Use(ctx, "personal", node, "assistant"); err != nil {
			t.Fatal(err)
		}
	}
	keys = func(node string) (string, error) { return pub[node], nil }
	return managerPub, managerPriv, holders, keys
}

// The point of the phase, in one test: a refresh seals the ACCESS token to every holder, and the
// refresh token is nowhere a holder can reach it.
func TestARefreshDeliversTheAccessTokenAndNeverTheRefreshToken(t *testing.T) {
	held, ctx := fresh(t)
	fake := &fakeRefresher{access: "at-brand-new-access-token"}
	_, _, holders, keys := aRefreshableLicence(t, held, ctx, fake)

	sealed, err := held.Refresh(ctx, "personal", keys)
	if err != nil {
		t.Fatal(err)
	}
	if sealed != 2 {
		t.Fatalf("%d holder(s) were resealed, expected 2", sealed)
	}

	for node, open := range holders {
		blob, err := held.KeyFor(ctx, "personal", node, "assistant")
		if err != nil {
			t.Fatal(err)
		}
		if blob == "" {
			t.Fatalf("%s got no access token", node)
		}
		got, err := open(blob)
		if err != nil {
			t.Fatalf("%s cannot open what it was delivered", node)
		}
		if string(got) != "at-brand-new-access-token" {
			t.Fatalf("%s was delivered %q, not the access token", node, got)
		}
		// The refresh token is not in the holder's delivery, opened or sealed.
		if string(got) == "rt-the-refresh-token" || strings.Contains(blob, "rt-the-refresh-token") {
			t.Fatalf("%s was delivered the refresh token", node)
		}
	}
}

// KeyFor delivers the access token and cannot deliver the refresh token, because the refresh token
// is not in licence_holder at all — the stripping is structural.
func TestKeyForNeverCarriesTheRefreshToken(t *testing.T) {
	held, ctx := fresh(t)
	fake := &fakeRefresher{access: "at-access"}
	_, _, _, keys := aRefreshableLicence(t, held, ctx, fake)
	if _, err := held.Refresh(ctx, "personal", keys); err != nil {
		t.Fatal(err)
	}

	// Every holder row, straight from the store: none of them holds the refresh token in any form.
	rows, err := held.store.Pool().Query(ctx,
		`select coalesce(sealed, '') from licence_holder where licence = 'personal'`)
	if err != nil {
		t.Fatal(err)
	}
	defer rows.Close()
	for rows.Next() {
		var sealed string
		if err := rows.Scan(&sealed); err != nil {
			t.Fatal(err)
		}
		if strings.Contains(sealed, "rt-the-refresh-token") {
			t.Fatal("a holder row carries the refresh token")
		}
	}
}

// The refresh token at rest is not readable from the database alone: the row holds ciphertext and a
// wrapped key, and only the manager node's private half opens it.
func TestTheRefreshTokenAtRestNeedsTheManagersKey(t *testing.T) {
	held, ctx := fresh(t)
	fake := &fakeRefresher{access: "at-access"}
	managerPub, managerPriv, _, _ := aRefreshableLicence(t, held, ctx, fake)

	// What the database holds, read straight from the row.
	var token, wrapped, managerKey string
	if err := held.store.Pool().QueryRow(ctx,
		`select token, wrapped_key, manager_key from refresh_grant where licence = 'personal'`).
		Scan(&token, &wrapped, &managerKey); err != nil {
		t.Fatal(err)
	}
	if strings.Contains(token, "rt-the-refresh-token") || strings.Contains(wrapped, "rt-the-refresh-token") {
		t.Fatal("the refresh token is in the row in the clear")
	}

	// The manager, holding its private key, reads it back.
	got, err := secrets.OpenAtRest(
		secrets.AtRest{Token: token, WrappedKey: wrapped, ManagerKey: managerKey},
		managerPub, managerPriv)
	if err != nil {
		t.Fatal(err)
	}
	if got != "rt-the-refresh-token" {
		t.Fatalf("the manager read back %q", got)
	}

	// Another node cannot, which is the whole of "the manager node only".
	otherPub, otherPriv := managerPair(t)
	if _, err := secrets.OpenAtRest(
		secrets.AtRest{Token: token, WrappedKey: wrapped, ManagerKey: managerKey},
		otherPub, otherPriv); err == nil {
		t.Fatal("a node that is not the manager opened the refresh token")
	}
}

// After a refresh, holders hold a NEW access token, and the refresh token that was not rotated is
// unchanged — never delivered either way.
func TestAfterRefreshHoldersHoldANewAccessTokenAndTheGrantIsUnchanged(t *testing.T) {
	held, ctx := fresh(t)
	fake := &fakeRefresher{access: "at-first"}
	_, _, holders, keys := aRefreshableLicence(t, held, ctx, fake)

	// A prior access token, so we can see it change.
	if _, err := held.Refresh(ctx, "personal", keys); err != nil {
		t.Fatal(err)
	}
	before := map[string]string{}
	for node := range holders {
		blob, err := held.KeyFor(ctx, "personal", node, "assistant")
		if err != nil {
			t.Fatal(err)
		}
		before[node] = blob
	}
	grantBefore := grantRow(t, held, ctx)

	// A second refresh with a different access token and no rotation of the refresh token.
	fake.access = "at-second"
	if _, err := held.Refresh(ctx, "personal", keys); err != nil {
		t.Fatal(err)
	}
	for node, open := range holders {
		blob, err := held.KeyFor(ctx, "personal", node, "assistant")
		if err != nil {
			t.Fatal(err)
		}
		if blob == before[node] {
			t.Fatalf("%s was not given a new sealed access token", node)
		}
		got, err := open(blob)
		if err != nil {
			t.Fatal(err)
		}
		if string(got) != "at-second" {
			t.Fatalf("%s holds %q, not the new access token", node, got)
		}
	}
	if grantRow(t, held, ctx) != grantBefore {
		t.Fatal("the refresh token changed although the vendor did not rotate it")
	}
}

// When the vendor rotates the refresh token too, the stored envelope is replaced with the
// re-encrypted one — and it is still not deliverable to a holder.
func TestARotatedRefreshTokenReplacesTheStoredEnvelope(t *testing.T) {
	held, ctx := fresh(t)
	fake := &fakeRefresher{access: "at-access"}
	managerPub, managerPriv, _, keys := aRefreshableLicence(t, held, ctx, fake)

	grantBefore := grantRow(t, held, ctx)

	rotated, err := secrets.SealAtRest("rt-a-rotated-refresh-token", managerPub)
	if err != nil {
		t.Fatal(err)
	}
	fake.newAtRest = &rotated
	if _, err := held.Refresh(ctx, "personal", keys); err != nil {
		t.Fatal(err)
	}

	if grantRow(t, held, ctx) == grantBefore {
		t.Fatal("the rotated refresh token did not replace the stored envelope")
	}
	at, ok, err := held.RefreshGrant(ctx, "personal")
	if err != nil || !ok {
		t.Fatalf("the rotated grant is not stored: ok=%v err=%v", ok, err)
	}
	got, err := secrets.OpenAtRest(at, managerPub, managerPriv)
	if err != nil {
		t.Fatal(err)
	}
	if got != "rt-a-rotated-refresh-token" {
		t.Fatalf("the stored grant opened to %q, not the rotated token", got)
	}
}

// A static-key licence has no refresh token and the carve-out never fires: it has no manager, and it
// cannot be refreshed.
func TestAStaticKeyLicenceHasNoManagerAndNoRefresh(t *testing.T) {
	held, ctx := fresh(t)
	if err := held.Add(ctx, "plain", "anthropic-api-key", nil); err != nil {
		t.Fatal(err)
	}
	if err := held.SetManager(ctx, "plain", "workstation"); err == nil {
		t.Fatal("a static-key licence was given a manager")
	}
	if _, ok, err := held.RefreshGrant(ctx, "plain"); err != nil || ok {
		t.Fatalf("a static-key licence has a refresh token stored: ok=%v err=%v", ok, err)
	}
	if _, err := held.Refresh(ctx, "plain", func(string) (string, error) { return "", nil }); err == nil {
		t.Fatal("a static-key licence was refreshed")
	}
}

// A refreshable licence with no manager named cannot be refreshed, and says how to name one.
func TestARefreshableLicenceWithoutAManagerIsRefused(t *testing.T) {
	held, ctx := fresh(t)
	if err := held.Add(ctx, "personal", "anthropic", nil); err != nil {
		t.Fatal(err)
	}
	_, err := held.Refresh(ctx, "personal", func(string) (string, error) { return "", nil })
	if err == nil {
		t.Fatal("a licence with no manager was refreshed")
	}
	if !strings.Contains(err.Error(), "manager") {
		t.Fatalf("the refusal does not point at the missing manager: %v", err)
	}
}

// grantRow is the whole at-rest envelope as one string, for asserting it changed or did not.
func grantRow(t *testing.T, held *Licences, ctx context.Context) string {
	t.Helper()
	at, ok, err := held.RefreshGrant(ctx, "personal")
	if err != nil {
		t.Fatal(err)
	}
	if !ok {
		return ""
	}
	return at.Token + "|" + at.WrappedKey + "|" + at.ManagerKey
}