sherlock/internal/broker/methods.go

package broker

import (
	"context"
	"encoding/json"
	"errors"
	"time"

	"gitea.alexandru.macocian.me/Charlie/sherlock/internal/authn"
	"gitea.alexandru.macocian.me/Charlie/sherlock/internal/keyring"
	"gitea.alexandru.macocian.me/Charlie/sherlock/internal/rpc"
)

// Methods bundles the broker's dependencies and exposes a MethodTable
// for socket dispatch. The fields are exported so cmd/sherlock-broker
// can construct one with real implementations and tests can swap in
// fakes (the *keyring.Store interface is the obvious seam).
type Methods struct {
	Store     keyring.Store
	Flow      *authn.Flow
	Refresher *authn.Refresher
	Cfg       authn.Config

	// ShutdownFn is invoked by the "shutdown" RPC. Wired by the daemon
	// to cancel the root context.
	ShutdownFn func()

	// AgentNames feeds the status response so the CLI can show which
	// profiles are available without re-loading them.
	AgentNames func() []string
}

// MethodTable returns the map suitable for NewService.
func (m *Methods) MethodTable() map[string]MethodHandler {
	return map[string]MethodHandler{
		rpc.MethodStatus:     m.handleStatus,
		rpc.MethodWhoami:     m.handleWhoami,
		rpc.MethodLoginStart: m.handleLoginStart,
		rpc.MethodLoginWait:  m.handleLoginWait,
		rpc.MethodLogout:     m.handleLogout,
		rpc.MethodGetIDToken: m.handleGetIDToken,
		rpc.MethodShutdown:   m.handleShutdown,
	}
}

// loadTokens reads the keyring once and caches the result in memory.
// We re-read on every status to keep things simple — the keyring is
// the source of truth.
func (m *Methods) currentTokens() (keyring.TokenSet, bool, error) {
	ts, err := m.Store.GetAuthentikTokens()
	if errors.Is(err, keyring.ErrNoTokens) {
		return keyring.TokenSet{}, false, nil
	}
	if err != nil {
		return keyring.TokenSet{}, false, err
	}
	return ts, !ts.Empty(), nil
}

// StatusResult is what the status RPC returns.
type StatusResult struct {
	LoggedIn        bool       `json:"logged_in"`
	Subject         string     `json:"sub,omitempty"`
	Name            string     `json:"name,omitempty"`
	Email           string     `json:"email,omitempty"`
	IDExpiresAt     *time.Time `json:"id_expires_at,omitempty"`
	RefreshExpAt    *time.Time `json:"refresh_expires_at,omitempty"`
	Issuer          string     `json:"issuer,omitempty"`
	AuthnConfigured bool       `json:"authn_configured"`
	Agents          []string   `json:"agents"`
}

func ptrTime(t time.Time) *time.Time {
	if t.IsZero() {
		return nil
	}
	tt := t
	return &tt
}

func (m *Methods) handleStatus(_ context.Context, _ json.RawMessage) (any, *rpc.Error) {
	ts, ok, err := m.currentTokens()
	if err != nil {
		return nil, rpc.Errorf(rpc.CodeInternal, "read keyring: %v", err)
	}
	out := StatusResult{
		LoggedIn:        ok,
		AuthnConfigured: m.Cfg.Configured(),
	}
	if m.AgentNames != nil {
		out.Agents = m.AgentNames()
	}
	if ok {
		out.Subject = ts.Subject
		out.Name = ts.Name
		out.Email = ts.Email
		out.Issuer = ts.Issuer
		out.IDExpiresAt = ptrTime(ts.IDExpiresAt)
		out.RefreshExpAt = ptrTime(ts.RefreshExpAt)
	}
	return out, nil
}

// WhoamiResult is what the whoami RPC returns.
type WhoamiResult struct {
	Subject string `json:"sub"`
	Name    string `json:"name,omitempty"`
	Email   string `json:"email,omitempty"`
}

func (m *Methods) handleWhoami(_ context.Context, _ json.RawMessage) (any, *rpc.Error) {
	ts, ok, err := m.currentTokens()
	if err != nil {
		return nil, rpc.Errorf(rpc.CodeInternal, "read keyring: %v", err)
	}
	if !ok {
		return nil, rpc.NewError(rpc.CodeNotLoggedIn, "not logged in")
	}
	return WhoamiResult{Subject: ts.Subject, Name: ts.Name, Email: ts.Email}, nil
}

// LoginStartResult is what the login_start RPC returns.
type LoginStartResult struct {
	AuthURL string `json:"auth_url"`
	State   string `json:"state"`
}

func (m *Methods) handleLoginStart(ctx context.Context, _ json.RawMessage) (any, *rpc.Error) {
	if !m.Cfg.Configured() {
		return nil, rpc.NewError(rpc.CodeNotConfigured,
			"broker is not configured for Authentik yet; see plan.md Open #11 — set SHERLOCK_AUTHENTIK_ISSUER and SHERLOCK_AUTHENTIK_CLIENT_ID")
	}
	res, err := m.Flow.StartFlow(ctx)
	if err != nil {
		if errors.Is(err, authn.ErrNotConfigured) {
			return nil, rpc.NewError(rpc.CodeNotConfigured, err.Error())
		}
		return nil, rpc.Errorf(rpc.CodeInternal, "start flow: %v", err)
	}
	return LoginStartResult{AuthURL: res.AuthURL, State: res.State}, nil
}

// LoginWaitParams are accepted by the login_wait RPC.
type LoginWaitParams struct {
	State          string `json:"state"`
	TimeoutSeconds int    `json:"timeout_seconds"`
}

// LoginWaitResult is what login_wait returns.
type LoginWaitResult struct {
	OK      bool   `json:"ok"`
	Subject string `json:"sub,omitempty"`
}

func (m *Methods) handleLoginWait(ctx context.Context, raw json.RawMessage) (any, *rpc.Error) {
	var p LoginWaitParams
	if len(raw) > 0 {
		if err := json.Unmarshal(raw, &p); err != nil {
			return nil, rpc.Errorf(rpc.CodeBadRequest, "decode params: %v", err)
		}
	}
	if p.State == "" {
		return nil, rpc.NewError(rpc.CodeBadRequest, "state is required")
	}
	timeout := time.Duration(p.TimeoutSeconds) * time.Second
	res, err := m.Flow.AwaitCallback(ctx, p.State, timeout)
	if err != nil {
		if errors.Is(err, authn.ErrFlowExpired) {
			return nil, rpc.NewError(rpc.CodeFlowExpired, err.Error())
		}
		if errors.Is(err, authn.ErrUnknownState) {
			return nil, rpc.NewError(rpc.CodeBadRequest, err.Error())
		}
		return nil, rpc.Errorf(rpc.CodeInternal, "await callback: %v", err)
	}
	if err := m.Store.SetAuthentikTokens(res.ToTokenSet()); err != nil {
		return nil, rpc.Errorf(rpc.CodeInternal, "persist tokens: %v", err)
	}
	return LoginWaitResult{OK: true, Subject: res.Subject}, nil
}

func (m *Methods) handleLogout(_ context.Context, _ json.RawMessage) (any, *rpc.Error) {
	if err := m.Store.ClearAuthentikTokens(); err != nil {
		return nil, rpc.Errorf(rpc.CodeInternal, "clear keyring: %v", err)
	}
	return struct{}{}, nil
}

// GetIDTokenResult is what get_id_token returns.
type GetIDTokenResult struct {
	Token     string    `json:"token"`
	ExpiresAt time.Time `json:"expires_at"`
}

func (m *Methods) handleGetIDToken(ctx context.Context, _ json.RawMessage) (any, *rpc.Error) {
	ts, ok, err := m.currentTokens()
	if err != nil {
		return nil, rpc.Errorf(rpc.CodeInternal, "read keyring: %v", err)
	}
	if !ok {
		return nil, rpc.NewError(rpc.CodeNotLoggedIn, "not logged in")
	}
	if m.Refresher != nil {
		fresh, err := m.Refresher.EnsureFresh(ctx, ts)
		if err != nil {
			if errors.Is(err, authn.ErrNoRefreshToken) {
				return nil, rpc.NewError(rpc.CodeNotLoggedIn, "refresh token missing — re-login required")
			}
			return nil, rpc.Errorf(rpc.CodeInternal, "refresh: %v", err)
		}
		if fresh.IDToken != ts.IDToken {
			if err := m.Store.SetAuthentikTokens(fresh); err != nil {
				return nil, rpc.Errorf(rpc.CodeInternal, "persist refreshed tokens: %v", err)
			}
			ts = fresh
		}
	}
	return GetIDTokenResult{Token: ts.IDToken, ExpiresAt: ts.IDExpiresAt}, nil
}

func (m *Methods) handleShutdown(_ context.Context, _ json.RawMessage) (any, *rpc.Error) {
	if m.ShutdownFn != nil {
		go m.ShutdownFn() // run async so we can return a response first
	}
	return struct{}{}, nil
}