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-controller/internal/licences/adapters"
	"github.com/novox/mesh-controller/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 returns the private key string too, because the MANAGER opens
// its own refresh token — the one node that reads it back — and the host on that node does so with
// box.OpenAnonymous, exactly as it opens any sealed credential.
func managerPair(t *testing.T) (public, private 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()),
		base64.StdEncoding.EncodeToString(priv.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
			}
			return out, nil
		}
}

type fakeRefresher struct {
	access        string
	newSealed     string
	newManagerKey string
	got           adapters.RefreshInput
}

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

// A refreshable-grant licence set up end to end: a manager node running the manager module (a holder
// delivered the refresh token), the refresh token sealed to it, two CONSUMER holders — one of them on
// the manager node itself, to exercise co-location — and a fake vendor refresher plugged in.
//
// The manager node is "workstation" and its manager module is "manager"; the consuming module is
// "assistant", present on both "workstation" and "laptop".
func aRefreshableLicence(t *testing.T, held *Licences, ctx context.Context, fake *fakeRefresher) (
	managerPub, managerPriv string, managerOpen func(string) ([]byte, error),
	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", "manager"); err != nil {
		t.Fatal(err)
	}

	managerPub, managerPriv, managerOpen = managerPair(t)
	sealedRefresh, err := secrets.Seal(managerPub, []byte("rt-the-refresh-token"))
	if err != nil {
		t.Fatal(err)
	}
	if err := held.SetRefreshGrant(ctx, "personal", sealedRefresh, managerPub); err != nil {
		t.Fatal(err)
	}
	// The manager module is a holder too, on the manager node, so resealAndPublish has it to skip.
	if err := held.Use(ctx, "personal", "workstation", "manager"); err != nil {
		t.Fatal(err)
	}

	holders = map[string]func(string) ([]byte, error){}
	pub := map[string]string{"workstation": managerPub}
	_, holders["workstation"] = "", managerOpen // consumer on the manager node shares its key
	for _, node := range []string{"workstation", "laptop"} {
		if node == "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, managerOpen, holders, keys
}

// The point of the phase, in one test: a refresh seals the ACCESS token to every CONSUMER holder, the
// manager holder is delivered the refresh token, and the refresh token is nowhere a consumer reaches.
func TestARefreshDeliversTheAccessTokenAndNeverTheRefreshToken(t *testing.T) {
	held, ctx := fresh(t)
	fake := &fakeRefresher{access: "at-brand-new-access-token"}
	_, _, managerOpen, 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 consumer 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)
		}
		if string(got) == "rt-the-refresh-token" || strings.Contains(blob, "rt-the-refresh-token") {
			t.Fatalf("%s was delivered the refresh token", node)
		}
	}

	// The manager holder is delivered the refresh token, and opens it with the node's own key.
	mgrBlob, err := held.KeyFor(ctx, "personal", "workstation", "manager")
	if err != nil {
		t.Fatal(err)
	}
	got, err := managerOpen(mgrBlob)
	if err != nil {
		t.Fatal("the manager cannot open the refresh token delivered to it")
	}
	if string(got) != "rt-the-refresh-token" {
		t.Fatalf("the manager was delivered %q, not the refresh token", got)
	}
}

// A CONSUMER holder is never delivered the refresh token, because a consumer's row never holds it and
// KeyFor for a consumer reads licence_holder — the separation is structural.
func TestKeyForNeverCarriesTheRefreshTokenToAConsumer(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 CONSUMER holder row, straight from the store: none holds the refresh token in any form.
	rows, err := held.store.Pool().Query(ctx,
		`select coalesce(sealed, '') from licence_holder where licence = 'personal' and module = 'assistant'`)
	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 consumer holder row carries the refresh token")
		}
	}
}

// The refresh token is not readable from the database alone: the row holds a sealed box, and only the
// manager node's private half opens it — the same guarantee every sealed credential here has.
func TestTheRefreshTokenNeedsTheManagersKey(t *testing.T) {
	held, ctx := fresh(t)
	fake := &fakeRefresher{access: "at-access"}
	managerPub, managerPriv, _, _, _ := aRefreshableLicence(t, held, ctx, fake)

	var sealed, managerKey string
	if err := held.store.Pool().QueryRow(ctx,
		`select sealed, manager_key from refresh_grant where licence = 'personal'`).
		Scan(&sealed, &managerKey); err != nil {
		t.Fatal(err)
	}
	if strings.Contains(sealed, "rt-the-refresh-token") {
		t.Fatal("the refresh token is in the row in the clear")
	}
	if managerKey != managerPub {
		t.Fatal("the stored manager key is not the manager's public key")
	}

	// The manager, holding its private key, reads it back with box.OpenAnonymous (as the host does).
	got := openAnon(t, sealed, managerPub, managerPriv)
	if string(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 _, ok := tryOpenAnon(sealed, otherPub, otherPriv); ok {
		t.Fatal("a node that is not the manager opened the refresh token")
	}
}

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

	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)

	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 sealed box is replaced with the re-sealed
// one — and it is still delivered only to the manager, opening only with the manager's key.
func TestARotatedRefreshTokenReplacesTheStoredBox(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.Seal(managerPub, []byte("rt-a-rotated-refresh-token"))
	if err != nil {
		t.Fatal(err)
	}
	fake.newSealed, fake.newManagerKey = rotated, managerPub
	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 box")
	}
	sealed, key, ok, err := held.RefreshGrant(ctx, "personal")
	if err != nil || !ok {
		t.Fatalf("the rotated grant is not stored: ok=%v err=%v", ok, err)
	}
	if key != managerPub {
		t.Fatal("the rotated grant is not sealed to the manager's key")
	}
	got := openAnon(t, sealed, managerPub, managerPriv)
	if string(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", "manager"); 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 sealed grant as one string, for asserting it changed or did not.
func grantRow(t *testing.T, held *Licences, ctx context.Context) string {
	t.Helper()
	sealed, key, ok, err := held.RefreshGrant(ctx, "personal")
	if err != nil {
		t.Fatal(err)
	}
	if !ok {
		return ""
	}
	return sealed + "|" + key
}

// openAnon opens an anonymous sealed box with a node's key pair — the host's Unseal, inlined for a test.
func openAnon(t *testing.T, sealed, pubB64, privB64 string) []byte {
	t.Helper()
	out, ok := tryOpenAnon(sealed, pubB64, privB64)
	if !ok {
		t.Fatal("box.OpenAnonymous failed for a value that should open")
	}
	return out
}

func tryOpenAnon(sealed, pubB64, privB64 string) ([]byte, bool) {
	blob, err := base64.StdEncoding.DecodeString(sealed)
	if err != nil {
		return nil, false
	}
	pubRaw, _ := base64.StdEncoding.DecodeString(pubB64)
	privRaw, _ := base64.StdEncoding.DecodeString(privB64)
	var pub, priv [32]byte
	copy(pub[:], pubRaw)
	copy(priv[:], privRaw)
	return box.OpenAnonymous(nil, blob, &pub, &priv)
}