at-container-registry/pkg/atproto/did/plc_test.go

package did

import (
	"context"
	"encoding/json"
	"io"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"testing"

	"github.com/bluesky-social/indigo/atproto/atcrypto"
	didplc "github.com/did-method-plc/go-didplc/didplc"
)

// fakePLC stands up an httptest server that serves a single op log entry and
// captures any submitted update op for inspection.
type fakePLC struct {
	server     *httptest.Server
	did        string
	logEntries []didplc.OpEnum
	submitted  []didplc.RegularOp
}

func (f *fakePLC) URL() string { return f.server.URL }

func (f *fakePLC) Close() { f.server.Close() }

// newFakePLC creates a fake PLC directory pre-loaded with a single signed
// genesis op containing the given rotation keys (in priority order). Returns
// the fake server and the resulting did:plc DID derived from the genesis op.
func newFakePLC(t *testing.T, rotationKeys []*atcrypto.PrivateKeyK256, signer atcrypto.PrivateKey, signingKey *atcrypto.PrivateKeyK256) *fakePLC {
	t.Helper()

	rotationDIDKeys := make([]string, 0, len(rotationKeys))
	for _, k := range rotationKeys {
		pub, err := k.PublicKey()
		if err != nil {
			t.Fatalf("rotation key public: %v", err)
		}
		rotationDIDKeys = append(rotationDIDKeys, pub.DIDKey())
	}

	sigPub, err := signingKey.PublicKey()
	if err != nil {
		t.Fatalf("signing key public: %v", err)
	}

	op := &didplc.RegularOp{
		Type:         "plc_operation",
		RotationKeys: rotationDIDKeys,
		VerificationMethods: map[string]string{
			"atproto": sigPub.DIDKey(),
		},
		AlsoKnownAs: []string{"at://example.test"},
		Services: map[string]didplc.OpService{
			"atproto_pds": {Type: "AtprotoPersonalDataServer", Endpoint: "https://example.test"},
		},
		Prev: nil,
	}
	if err := op.Sign(signer); err != nil {
		t.Fatalf("sign genesis: %v", err)
	}
	did, err := op.DID()
	if err != nil {
		t.Fatalf("compute DID: %v", err)
	}

	f := &fakePLC{
		did:        did,
		logEntries: []didplc.OpEnum{{Regular: op}},
	}

	mux := http.NewServeMux()
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		if r.Method == http.MethodGet && strings.HasSuffix(r.URL.Path, "/log") {
			w.Header().Set("Content-Type", "application/json")
			_ = json.NewEncoder(w).Encode(f.logEntries)
			return
		}
		if r.Method == http.MethodPost && r.URL.Path == "/"+did {
			body, err := io.ReadAll(r.Body)
			if err != nil {
				http.Error(w, err.Error(), http.StatusBadRequest)
				return
			}
			var op didplc.RegularOp
			if err := json.Unmarshal(body, &op); err != nil {
				http.Error(w, err.Error(), http.StatusBadRequest)
				return
			}
			f.submitted = append(f.submitted, op)
			w.WriteHeader(http.StatusOK)
			return
		}
		http.NotFound(w, r)
	})
	f.server = httptest.NewServer(mux)
	return f
}

// generateK256 returns a fresh K-256 keypair, failing the test on error.
func generateK256(t *testing.T) *atcrypto.PrivateKeyK256 {
	t.Helper()
	k, err := atcrypto.GeneratePrivateKeyK256()
	if err != nil {
		t.Fatalf("generate K-256: %v", err)
	}
	return k
}

// writeSigningKey persists a signing key to a temp file and returns its path.
// Used so EnsureCurrent can load it via oauth.GenerateOrLoadPDSKey.
func writeSigningKey(t *testing.T, dir string, key *atcrypto.PrivateKeyK256) string {
	t.Helper()
	path := filepath.Join(dir, "signing.key")
	if err := os.WriteFile(path, key.Bytes(), 0600); err != nil {
		t.Fatalf("write signing key: %v", err)
	}
	return path
}

func TestEnsureCurrent_PreservesRotationKeys(t *testing.T) {
	ctx := context.Background()
	tmp := t.TempDir()

	// Server-side rotation key (the one stored in database.rotation_key) and an
	// "offline" recovery key that lives only in PLC. The genesis op lists offline
	// FIRST (highest priority).
	serverRot := generateK256(t)
	offlineRot := generateK256(t)

	// Original signing key used to build genesis; the local signing key on disk
	// will be different to force EnsureCurrent into the update path.
	originalSigning := generateK256(t)
	localSigning := generateK256(t)
	writeSigningKey(t, tmp, localSigning)

	fake := newFakePLC(t, []*atcrypto.PrivateKeyK256{offlineRot, serverRot}, serverRot, originalSigning)
	defer fake.Close()

	cfg := Config{
		PublicURL:           "https://example.test",
		PLCDirectoryURL:     fake.URL(),
		VerificationKeyName: "atproto",
		Services: map[string]Service{
			"atproto_pds": {Type: "AtprotoPersonalDataServer", Endpoint: "https://example.test"},
		},
	}

	if err := EnsureCurrent(ctx, fake.did, serverRot, localSigning, cfg); err != nil {
		t.Fatalf("EnsureCurrent: %v", err)
	}

	if len(fake.submitted) != 1 {
		t.Fatalf("expected exactly one update op submitted, got %d", len(fake.submitted))
	}
	got := fake.submitted[0]

	offlinePub, _ := offlineRot.PublicKey()
	serverPub, _ := serverRot.PublicKey()
	want := []string{offlinePub.DIDKey(), serverPub.DIDKey()}

	if len(got.RotationKeys) != len(want) {
		t.Fatalf("rotation keys length: got %d want %d", len(got.RotationKeys), len(want))
	}
	for i := range want {
		if got.RotationKeys[i] != want[i] {
			t.Errorf("rotation key [%d]: got %s want %s", i, got.RotationKeys[i], want[i])
		}
	}

	// Verify signing key was actually rotated (sanity check we hit the update path).
	localPub, _ := localSigning.PublicKey()
	if got.VerificationMethods["atproto"] != localPub.DIDKey() {
		t.Errorf("expected signing key to update to local key %s, got %s",
			localPub.DIDKey(), got.VerificationMethods["atproto"])
	}
}

func TestEnsureCurrent_RefusesUpdateWhenLocalKeyMissing(t *testing.T) {
	ctx := context.Background()
	tmp := t.TempDir()

	// Genesis lists only the offline key. The local server has been rotated out.
	offlineRot := generateK256(t)
	localRot := generateK256(t) // not in PLC

	originalSigning := generateK256(t)
	localSigning := generateK256(t)
	writeSigningKey(t, tmp, localSigning)

	fake := newFakePLC(t, []*atcrypto.PrivateKeyK256{offlineRot}, offlineRot, originalSigning)
	defer fake.Close()

	cfg := Config{
		PublicURL:           "https://example.test",
		PLCDirectoryURL:     fake.URL(),
		VerificationKeyName: "atproto",
		Services: map[string]Service{
			"atproto_pds": {Type: "AtprotoPersonalDataServer", Endpoint: "https://example.test"},
		},
	}

	// Local signing key has drifted, which would normally trigger an update.
	if err := EnsureCurrent(ctx, fake.did, localRot, localSigning, cfg); err != nil {
		t.Fatalf("EnsureCurrent returned error: %v", err)
	}

	if len(fake.submitted) != 0 {
		t.Fatalf("expected no update submission when local rotation key isn't in PLC list, got %d", len(fake.submitted))
	}
}

func TestEnsureCurrent_NoOpWhenCurrent(t *testing.T) {
	ctx := context.Background()
	tmp := t.TempDir()

	serverRot := generateK256(t)
	signing := generateK256(t)
	writeSigningKey(t, tmp, signing)

	fake := newFakePLC(t, []*atcrypto.PrivateKeyK256{serverRot}, serverRot, signing)
	defer fake.Close()

	cfg := Config{
		PublicURL:           "https://example.test",
		PLCDirectoryURL:     fake.URL(),
		VerificationKeyName: "atproto",
		Services: map[string]Service{
			"atproto_pds": {Type: "AtprotoPersonalDataServer", Endpoint: "https://example.test"},
		},
	}

	if err := EnsureCurrent(ctx, fake.did, serverRot, signing, cfg); err != nil {
		t.Fatalf("EnsureCurrent: %v", err)
	}
	if len(fake.submitted) != 0 {
		t.Fatalf("expected no update when state is current, got %d submitted", len(fake.submitted))
	}
}