at-container-registry/pkg/hold/pds/appview_token_test.go

package pds

import (
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"

	"atcr.io/pkg/atproto"
	"atcr.io/pkg/auth"

	"github.com/bluesky-social/indigo/atproto/atcrypto"
)

// appviewTestEnv stands up an httptest server that serves a did:web DID
// document for a generated keypair, returns the matching appview DID (with
// percent-encoded port for local-dev), and resets the shared jti cache
// between tests.
type appviewTestEnv struct {
	server     *httptest.Server
	priv       atcrypto.PrivateKey
	appviewDID string
	holdDID    string
}

func newAppviewTestEnv(t *testing.T) *appviewTestEnv {
	t.Helper()
	priv, err := atcrypto.GeneratePrivateKeyP256()
	if err != nil {
		t.Fatalf("generate P-256: %v", err)
	}
	return newAppviewTestEnvWithKey(t, priv)
}

func newAppviewTestEnvWithKey(t *testing.T, priv atcrypto.PrivateKey) *appviewTestEnv {
	t.Helper()
	pub, err := priv.PublicKey()
	if err != nil {
		t.Fatalf("public key: %v", err)
	}

	doc := map[string]any{
		"@context": []string{"https://www.w3.org/ns/did/v1"},
		"verificationMethod": []map[string]string{{
			"id":                 "#atproto",
			"type":               "Multikey",
			"publicKeyMultibase": pub.Multibase(),
		}},
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/.well-known/did.json", func(w http.ResponseWriter, r *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(doc)
	})
	server := httptest.NewServer(mux)
	t.Cleanup(server.Close)

	u, err := url.Parse(server.URL)
	if err != nil {
		t.Fatalf("parse server URL: %v", err)
	}
	// did:web requires percent-encoded ":" for ports
	appviewDID := "did:web:" + strings.ReplaceAll(u.Host, ":", "%3A")

	// fetchAppviewPublicKey downgrades to http when test mode is on
	atproto.SetTestMode(true)
	t.Cleanup(func() { atproto.SetTestMode(false) })

	// Reset shared jti cache so other tests don't pollute this one and vice versa.
	*sharedAppviewJTICache = *newJTIReplayCache()
	t.Cleanup(func() { *sharedAppviewJTICache = *newJTIReplayCache() })

	return &appviewTestEnv{
		server:     server,
		priv:       priv,
		appviewDID: appviewDID,
		holdDID:    "did:web:hold01.atcr.io",
	}
}

func (e *appviewTestEnv) newRequestWithToken(t *testing.T, token string) *http.Request {
	t.Helper()
	req := httptest.NewRequest(http.MethodPost, "/test", nil)
	req.Header.Set("Authorization", "Bearer "+token)
	return req
}

func (e *appviewTestEnv) sign(t *testing.T, userDID string) string {
	t.Helper()
	tok, err := auth.CreateAppviewServiceToken(e.priv, e.appviewDID, e.holdDID, userDID)
	if err != nil {
		t.Fatalf("CreateAppviewServiceToken: %v", err)
	}
	return tok
}

func TestValidateAppviewToken_AcceptsK256SignedToken(t *testing.T) {
	priv, err := atcrypto.GeneratePrivateKeyK256()
	if err != nil {
		t.Fatalf("generate K-256: %v", err)
	}
	env := newAppviewTestEnvWithKey(t, priv)
	tok := env.sign(t, "did:plc:pddp4xt5lgnv2qsegbzzs4xg")

	got, err := ValidateAppviewToken(env.newRequestWithToken(t, tok), env.appviewDID, env.holdDID)
	if err != nil {
		t.Fatalf("ValidateAppviewToken (K-256): %v", err)
	}
	if got != "did:plc:pddp4xt5lgnv2qsegbzzs4xg" {
		t.Errorf("returned sub = %q", got)
	}
}

func TestValidateAppviewToken_AcceptsFreshToken(t *testing.T) {
	env := newAppviewTestEnv(t)
	tok := env.sign(t, "did:plc:pddp4xt5lgnv2qsegbzzs4xg")

	got, err := ValidateAppviewToken(env.newRequestWithToken(t, tok), env.appviewDID, env.holdDID)
	if err != nil {
		t.Fatalf("ValidateAppviewToken: %v", err)
	}
	if got != "did:plc:pddp4xt5lgnv2qsegbzzs4xg" {
		t.Errorf("returned sub = %q", got)
	}
}

func TestValidateAppviewToken_RejectsReplayedJTI(t *testing.T) {
	env := newAppviewTestEnv(t)
	tok := env.sign(t, "did:plc:pddp4xt5lgnv2qsegbzzs4xg")

	if _, err := ValidateAppviewToken(env.newRequestWithToken(t, tok), env.appviewDID, env.holdDID); err != nil {
		t.Fatalf("first validate failed: %v", err)
	}
	_, err := ValidateAppviewToken(env.newRequestWithToken(t, tok), env.appviewDID, env.holdDID)
	if err == nil {
		t.Fatal("expected replay rejection on second submission, got nil")
	}
	if !strings.Contains(err.Error(), "replay") && err != ErrTokenReplayed {
		t.Errorf("error %v does not mention replay", err)
	}
}

// TODO: re-enable when jti becomes mandatory. Today the hold verifier
// accepts tokens without a jti claim (with a warning) so a partial appview
// rollout doesn't break tier updates. Once all appview deployments are
// confirmed to emit jti, restore this test and the corresponding
// ErrMissingJTIClaim check in ValidateAppviewToken.
//
// func TestValidateAppviewToken_RejectsMissingJTI(t *testing.T) {
// 	env := newAppviewTestEnv(t)
// 	header := `{"alg":"ES256","typ":"JWT"}`
// 	now := time.Now().Unix()
// 	body := fmt.Sprintf(`{"iss":%q,"aud":[%q],"sub":%q,"exp":%d,"iat":%d}`,
// 		env.appviewDID, env.holdDID, "did:plc:pddp4xt5lgnv2qsegbzzs4xg", now+60, now)
// 	h := base64.RawURLEncoding.EncodeToString([]byte(header))
// 	p := base64.RawURLEncoding.EncodeToString([]byte(body))
// 	signingInput := h + "." + p
// 	sig, err := env.priv.HashAndSign([]byte(signingInput))
// 	if err != nil {
// 		t.Fatalf("sign: %v", err)
// 	}
// 	tok := signingInput + "." + base64.RawURLEncoding.EncodeToString(sig)
//
// 	_, err = ValidateAppviewToken(env.newRequestWithToken(t, tok), env.appviewDID, env.holdDID)
// 	if !errors.Is(err, ErrMissingJTIClaim) {
// 		t.Errorf("expected ErrMissingJTIClaim, got %v", err)
// 	}
// }

func TestValidateAppviewToken_AcceptsDistinctJTIs(t *testing.T) {
	env := newAppviewTestEnv(t)
	tokA := env.sign(t, "did:plc:pddp4xt5lgnv2qsegbzzs4xg")
	tokB := env.sign(t, "did:plc:pddp4xt5lgnv2qsegbzzs4xg")
	if tokA == tokB {
		t.Fatal("two signed tokens are byte-identical")
	}
	if _, err := ValidateAppviewToken(env.newRequestWithToken(t, tokA), env.appviewDID, env.holdDID); err != nil {
		t.Fatalf("validate A: %v", err)
	}
	if _, err := ValidateAppviewToken(env.newRequestWithToken(t, tokB), env.appviewDID, env.holdDID); err != nil {
		t.Errorf("validate B (distinct jti) rejected: %v", err)
	}
}