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

package did

import (
	"context"
	"fmt"
	"slices"

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

// AddRotationKeyOptions configures a rotation-key insert operation.
type AddRotationKeyOptions struct {
	// DID is the resolved did:plc identifier of the service.
	DID string

	// PLCDirectoryURL is the PLC directory endpoint (defaults to https://plc.directory if empty).
	PLCDirectoryURL string

	// RotationKey is the currently-authorized rotation key used to sign the update op.
	RotationKey atcrypto.PrivateKey

	// SigningKey is the local k256 verification key — its public part goes into the
	// new op's VerificationMethods so we don't accidentally drop it during the update.
	SigningKey *atcrypto.PrivateKeyK256

	// VerificationKeyName is the fragment under which SigningKey is registered
	// (e.g. "atproto" for a PDS, "atproto_label" for a labeler).
	VerificationKeyName string

	// NewKey is the rotation key to add. If nil, a fresh K-256 key is generated and
	// returned in the result so the caller can print/persist it.
	NewKey atcrypto.PrivateKeyExportable

	// Prepend places the new key at the highest priority position. When false the key
	// is appended at the lowest priority — only set false when the operator explicitly
	// asks for it.
	Prepend bool
}

// AddRotationKeyResult describes the outcome of an AddRotationKey call.
type AddRotationKeyResult struct {
	NewKey         atcrypto.PrivateKeyExportable
	NewKeyDIDKey   string
	Generated      bool
	AlreadyPresent bool
	ExistingAt     int
	InsertedAt     int
	TotalKeys      int
}

// AddRotationKey fetches the current PLC op log, inserts NewKey (generating one if nil),
// signs the update with RotationKey, and submits it. Caller is responsible for printing
// the generated key material — this function returns it on the result so prints can
// happen in the binary's own format.
func AddRotationKey(ctx context.Context, opt AddRotationKeyOptions) (*AddRotationKeyResult, error) {
	if opt.DID == "" {
		return nil, fmt.Errorf("plc: DID is required")
	}
	if opt.RotationKey == nil {
		return nil, fmt.Errorf("plc: rotation key is required to sign updates")
	}
	if opt.SigningKey == nil {
		return nil, fmt.Errorf("plc: signing key is required (becomes verificationMethods.%s)", opt.VerificationKeyName)
	}
	if opt.VerificationKeyName == "" {
		return nil, fmt.Errorf("plc: VerificationKeyName is required")
	}

	directory := opt.PLCDirectoryURL
	if directory == "" {
		directory = "https://plc.directory"
	}
	client := &didplc.Client{DirectoryURL: directory}

	res := &AddRotationKeyResult{NewKey: opt.NewKey}
	if res.NewKey == nil {
		raw, err := atcrypto.GeneratePrivateKeyK256()
		if err != nil {
			return nil, fmt.Errorf("plc: failed to generate rotation key: %w", err)
		}
		res.NewKey = raw
		res.Generated = true
	}
	newPub, err := res.NewKey.PublicKey()
	if err != nil {
		return nil, fmt.Errorf("plc: failed to derive new public key: %w", err)
	}
	res.NewKeyDIDKey = newPub.DIDKey()

	opLog, err := client.OpLog(ctx, opt.DID)
	if err != nil {
		return nil, fmt.Errorf("plc: failed to fetch op log for %s: %w", opt.DID, err)
	}
	if len(opLog) == 0 {
		return nil, fmt.Errorf("plc: empty op log for %s", opt.DID)
	}
	lastEntry := opLog[len(opLog)-1]
	lastOp := lastEntry.Regular
	if lastOp == nil {
		return nil, fmt.Errorf("plc: last operation is not a regular op")
	}

	for i, k := range lastOp.RotationKeys {
		if k == res.NewKeyDIDKey {
			res.AlreadyPresent = true
			res.ExistingAt = i
			res.TotalKeys = len(lastOp.RotationKeys)
			return res, nil
		}
	}

	rotationKeys := make([]string, 0, len(lastOp.RotationKeys)+1)
	if opt.Prepend {
		rotationKeys = append(rotationKeys, res.NewKeyDIDKey)
		rotationKeys = append(rotationKeys, lastOp.RotationKeys...)
		res.InsertedAt = 0
	} else {
		rotationKeys = append(rotationKeys, lastOp.RotationKeys...)
		rotationKeys = append(rotationKeys, res.NewKeyDIDKey)
		res.InsertedAt = len(rotationKeys) - 1
	}
	res.TotalKeys = len(rotationKeys)

	sigPub, err := opt.SigningKey.PublicKey()
	if err != nil {
		return nil, fmt.Errorf("plc: failed to derive signing public key: %w", err)
	}
	prevCID := lastEntry.AsOperation().CID().String()

	op := &didplc.RegularOp{
		Type:         "plc_operation",
		RotationKeys: rotationKeys,
		VerificationMethods: map[string]string{
			opt.VerificationKeyName: sigPub.DIDKey(),
		},
		AlsoKnownAs: lastOp.AlsoKnownAs,
		Services:    lastOp.Services,
		Prev:        &prevCID,
	}
	if err := op.Sign(opt.RotationKey); err != nil {
		return nil, fmt.Errorf("plc: failed to sign update: %w", err)
	}
	if err := client.Submit(ctx, opt.DID, op); err != nil {
		return nil, fmt.Errorf("plc: failed to submit update: %w", err)
	}
	return res, nil
}

// ListRotationKeysOptions configures a list-rotation-keys read.
type ListRotationKeysOptions struct {
	DID              string
	PLCDirectoryURL  string
	LocalRotationKey atcrypto.PrivateKey // optional — used to compute the LOCAL marker
}

// ListRotationKeysResult holds the priority-ordered rotation keys plus the local
// rotation key's did:key form (if provided), so callers can mark and warn appropriately.
type ListRotationKeysResult struct {
	DID          string
	Directory    string
	Keys         []string
	LocalDIDKey  string
	LocalPresent bool
}

// ListRotationKeys fetches the current PLC op and returns its rotation keys in priority order.
func ListRotationKeys(ctx context.Context, opt ListRotationKeysOptions) (*ListRotationKeysResult, error) {
	if opt.DID == "" {
		return nil, fmt.Errorf("plc: DID is required")
	}
	directory := opt.PLCDirectoryURL
	if directory == "" {
		directory = "https://plc.directory"
	}
	client := &didplc.Client{DirectoryURL: directory}

	opLog, err := client.OpLog(ctx, opt.DID)
	if err != nil {
		return nil, fmt.Errorf("plc: failed to fetch op log for %s: %w", opt.DID, err)
	}
	if len(opLog) == 0 {
		return nil, fmt.Errorf("plc: empty op log for %s", opt.DID)
	}
	lastOp := opLog[len(opLog)-1].Regular
	if lastOp == nil {
		return nil, fmt.Errorf("plc: last operation is not a regular op")
	}

	res := &ListRotationKeysResult{
		DID:       opt.DID,
		Directory: directory,
		Keys:      append([]string(nil), lastOp.RotationKeys...),
	}
	if opt.LocalRotationKey != nil {
		pub, err := opt.LocalRotationKey.PublicKey()
		if err != nil {
			return nil, fmt.Errorf("plc: failed to derive local rotation public key: %w", err)
		}
		res.LocalDIDKey = pub.DIDKey()
		if slices.Contains(res.Keys, res.LocalDIDKey) {
			res.LocalPresent = true
		}
	}
	return res, nil
}