at-container-registry/pkg/testpds/directory.go

package testpds

import (
	"context"
	"fmt"
	"sync"

	"github.com/bluesky-social/indigo/atproto/identity"
	"github.com/bluesky-social/indigo/atproto/syntax"
)

// Directory is an in-memory identity.Directory containing both the user
// identities served by the fake PDS and any service DIDs the test registers
// (e.g. the Hold's did:web identity). Install via atproto.SetDirectory().
type Directory struct {
	mu       sync.RWMutex
	byDID    map[string]*identity.Identity
	byHandle map[string]*identity.Identity
}

func newDirectory() *Directory {
	return &Directory{
		byDID:    make(map[string]*identity.Identity),
		byHandle: make(map[string]*identity.Identity),
	}
}

// Register adds an identity to the directory. Callers (Server.AddIdentity or
// the test harness adding the hold) construct the identity.Identity and pass
// it in. Subsequent Lookup/LookupDID/LookupHandle calls return this entry.
func (d *Directory) Register(ident *identity.Identity) {
	d.mu.Lock()
	defer d.mu.Unlock()
	d.byDID[ident.DID.String()] = ident
	if ident.Handle != syntax.HandleInvalid && ident.Handle.String() != "" {
		d.byHandle[ident.Handle.String()] = ident
	}
}

// LookupHandle implements identity.Directory.
func (d *Directory) LookupHandle(_ context.Context, handle syntax.Handle) (*identity.Identity, error) {
	d.mu.RLock()
	defer d.mu.RUnlock()
	if ident, ok := d.byHandle[handle.String()]; ok {
		return ident, nil
	}
	return nil, fmt.Errorf("%w: %s", identity.ErrHandleNotFound, handle)
}

// LookupDID implements identity.Directory.
func (d *Directory) LookupDID(_ context.Context, did syntax.DID) (*identity.Identity, error) {
	d.mu.RLock()
	defer d.mu.RUnlock()
	if ident, ok := d.byDID[did.String()]; ok {
		return ident, nil
	}
	return nil, fmt.Errorf("%w: %s", identity.ErrDIDNotFound, did)
}

// Lookup implements identity.Directory.
func (d *Directory) Lookup(ctx context.Context, atid syntax.AtIdentifier) (*identity.Identity, error) {
	if did, err := atid.AsDID(); err == nil {
		return d.LookupDID(ctx, did)
	}
	if handle, err := atid.AsHandle(); err == nil {
		return d.LookupHandle(ctx, handle)
	}
	return nil, fmt.Errorf("%w: %s", identity.ErrHandleResolutionFailed, atid)
}

// Purge implements identity.Directory. Our directory is authoritative — entries
// don't expire — so this is a no-op.
func (d *Directory) Purge(_ context.Context, _ syntax.AtIdentifier) error {
	return nil
}