fleetd/bridged/src/main/java/dev/ltms/fleet/auth/CallerResolver.java

package dev.ltms.fleet.auth;

import dev.ltms.fleet.mcp.ConnectionIdentity;
import dev.ltms.fleet.peer.MemberRole;

import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.util.Map;
import java.util.function.Function;
import java.util.function.Supplier;

/**
 * Resolves every caller to a {@link Principal}, for both entry paths into the core (CB-501).
 *
 * <p>There are two of them and they are not layered the way the docs suggest: {@code FleetMcp}
 * calls the service layer directly and is mounted as a raw servlet (so it never passes through a
 * Javalin filter), while the REST routes historically resolved no identity at all. Both now
 * delegate here, so the authorization rules are stated once instead of drifting apart.
 *
 * <p><strong>Resolution order</strong> — connection identity first, token second, nothing third:
 * <ol>
 *   <li>A loopback peer PID that maps to a pane named by {@code leaders:}, by the legacy
 *       {@code primary.terminal} pin, or by an operator-labelled lead tab (CB-307, CB-530, CB-531)
 *       ⇒ {@link Role#PRIMARY}, carrying that lead's
 *       name. The pane mapping is as unforgeable as a worker's, and the config explicitly names
 *       that pane as a lead's own — without this rule a lead running <em>inside</em> a herdr pane
 *       is misread as a worker and locked out of orchestration. More than one pane may be named,
 *       so two leads can work as peers rather than one being demoted.</li>
 *   <li>A loopback peer PID that maps to a pane bound to a CB-548 architect slot ⇒
 *       {@link Role#ARCHITECT}, carrying the slot name. Just unforgeable as a worker's, and
 *       resolved from the <em>live</em> terminal→slot binding (never a request argument), before
 *       the generic worker fallback.</li>
 *   <li>A loopback peer PID that maps to any other herdr pane ⇒ {@link Role#WORKER}. This is
 *       unforgeable (the OS reports the PID, herdr owns the PID→pane map) and is honoured
 *       regardless of auth mode, so enabling auth never breaks the fleet.</li>
 *   <li>Otherwise, under {@code token} mode, a valid bearer token ⇒ {@link Role#PRIMARY}.</li>
 *   <li>Otherwise, under {@code loopback-trust}, a loopback caller ⇒ {@link Role#PRIMARY}
 *       (the historical behaviour, now an explicit configured choice).</li>
 *   <li>Otherwise {@link Role#ANONYMOUS}.</li>
 * </ol>
 */
public final class CallerResolver {

    private final ConnectionIdentity identity;
    private final boolean tokenMode;
    private final byte[] expectedToken; // null unless tokenMode
    /**
     * terminal_id → lead name; empty when nothing is pinned. CB-530.
     *
     * <p>A supplier rather than a map because the registry is no longer fixed at startup: CB-531
     * discovers leads by scanning herdr for operator-labelled tabs, so a lead that opens its tab
     * after the daemon booted must still be recognised. Consulted per resolve; the scanner behind
     * it is TTL-cached, so this is a map lookup in the common case.
     */
    private final Supplier<Map<String, String>> leadTerminals;
    /**
     * terminal_id → architect slot name; empty when nothing is configured. CB-548.
     *
     * <p>Like {@link #leadTerminals}, a supplier rather than a fixed map, so a binding injected
     * after startup — when the later spawn lifecycle establishes a live architect session, or an
     * operator pins one — takes effect without a restart. Consulted per resolve; today's wiring
     * in {@code Fleetd} reads a constant from config, which is the degenerate live case.
     */
    private final Supplier<Map<String, String>> architectTerminals;
    private final Function<String, MemberRole> memberSlotRoles;
    private final Function<String, String> memberSlotNames;

    /** Loopback-trust resolver: no token required, historical behaviour. Test-only. */
    CallerResolver(ConnectionIdentity identity) {
        this(identity, false, null, Map.of());
    }

    /** As {@link #CallerResolver(ConnectionIdentity, boolean, String, Map)} with no leads pinned. Test-only. */
    CallerResolver(ConnectionIdentity identity, boolean tokenMode, String token) {
        this(identity, tokenMode, token, Map.of());
    }

    /**
     * Single-pin form for the legacy {@code primary.terminal}-only configuration — one lead, named
     * {@code primary}.
     *
     * <p>A static factory rather than a fourth constructor overload on purpose: {@code String} and
     * {@code Map} overloads are ambiguous for a literal {@code null} argument, which is a compile
     * error at the call site and exactly the shape "unpinned" is written in.
     *
     * @param pinnedPrimaryTerminal the primary's own herdr {@code terminal_id}
     *                              ({@code null}/blank = unpinned)
     */
    static CallerResolver pinnedTo(ConnectionIdentity identity, boolean tokenMode,
                                   String token, String pinnedPrimaryTerminal) {
        return new CallerResolver(identity, tokenMode, token,
                pinnedPrimaryTerminal == null || pinnedPrimaryTerminal.isBlank()
                        ? Map.of() : Map.of(pinnedPrimaryTerminal, "primary"));
    }

    /**
     * @param identity       connection-based worker identification
     * @param tokenMode      when true, a non-worker caller must present a valid bearer token
     * @param token          the expected bearer token; required (non-blank) when {@code tokenMode}
     * @param leadTerminals  herdr {@code terminal_id} → lead name for every configured lead
     *                       (CB-530). A caller resolving to one of these panes is that lead — a
     *                       {@link Role#PRIMARY} — rather than a worker. Empty = nothing pinned,
     *                       so every pane resolves as a worker.
     */
    CallerResolver(ConnectionIdentity identity, boolean tokenMode, String token,
                   Map<String, String> leadTerminals) {
        this(identity, tokenMode, token, fixed(leadTerminals), null);
    }

    /**
     * Live-registry form: {@code leadTerminals} is consulted on every resolve, so leads discovered
     * after startup (CB-531's tab scan) take effect without a restart.
     *
     * <p>A static factory rather than a fourth constructor overload, for the same reason as
     * {@link #pinnedTo}: {@code Map} and {@code Supplier} overloads are ambiguous for a literal
     * {@code null}.
     */
    static CallerResolver withLeads(ConnectionIdentity identity, boolean tokenMode,
                                    String token,
                                    Supplier<Map<String, String>> leadTerminals) {
        return new CallerResolver(identity, tokenMode, token, leadTerminals, null);
    }

    /**
     * Live registry form that can confirm a bound slot is an architect slot.
     *
     * <p>This is the only public construction path. It keeps terminal bindings and slot roles in
     * the same {@link MemberRegistry}, so a configured architect can resolve as an architect.
     */
    public static CallerResolver withLeadsAndMembers(ConnectionIdentity identity,
                                                       boolean tokenMode, String token,
                                                       Supplier<Map<String, String>> leadTerminals,
                                                       MemberRegistry members) {
        return new CallerResolver(identity, tokenMode, token, leadTerminals,
                members == null ? null : members::snapshot,
                members == null ? null : members::roleForSlot,
                members == null ? null : members::nameForSlot);
    }

    private static Supplier<Map<String, String>> fixed(Map<String, String> leadTerminals) {
        Map<String, String> snapshot = leadTerminals == null ? Map.of() : Map.copyOf(leadTerminals);
        return () -> snapshot;
    }

    private CallerResolver(ConnectionIdentity identity, boolean tokenMode, String token,
                           Supplier<Map<String, String>> leadTerminals,
                           Supplier<Map<String, String>> architectTerminals) {
        this(identity, tokenMode, token, leadTerminals, architectTerminals, null);
    }

    private CallerResolver(ConnectionIdentity identity, boolean tokenMode, String token,
                           Supplier<Map<String, String>> leadTerminals,
                           Supplier<Map<String, String>> architectTerminals,
                           Function<String, MemberRole> memberSlotRoles) {
        this(identity, tokenMode, token, leadTerminals, architectTerminals, memberSlotRoles, null);
    }

    private CallerResolver(ConnectionIdentity identity, boolean tokenMode, String token,
                           Supplier<Map<String, String>> leadTerminals,
                           Supplier<Map<String, String>> architectTerminals,
                           Function<String, MemberRole> memberSlotRoles,
                           Function<String, String> memberSlotNames) {
        if (tokenMode && (token == null || token.isBlank())) {
            throw new IllegalArgumentException(
                    "auth.mode=token requires a non-empty token; check that the env var named by "
                            + "auth.tokenEnv is exported to the daemon's environment");
        }
        this.identity = identity;
        this.tokenMode = tokenMode;
        this.expectedToken = tokenMode ? token.getBytes(StandardCharsets.UTF_8) : null;
        this.leadTerminals = leadTerminals == null ? Map::of : leadTerminals;
        this.architectTerminals = architectTerminals == null ? Map::of : architectTerminals;
        this.memberSlotRoles = memberSlotRoles == null ? _ -> null : memberSlotRoles;
        this.memberSlotNames = memberSlotNames == null ? Function.identity() : memberSlotNames;
    }

    /**
     * The currently-recognised leads, {@code terminal_id → name} (CB-535).
     *
     * <p>Deliberately read from the same supplier {@link #resolve} consults, rather than from a
     * second copy handed to the roster: a lead that is <em>listed</em> but would not <em>resolve</em>
     * (or the reverse) is an address a peer cannot actually reach, and the two answers drifting apart
     * is precisely the confusion this exists to end. Live, so a lead discovered by the tab scan after
     * startup appears without a restart.
     */
    public Map<String, String> leads() {
        return leadTerminals.get();
    }

    /**
     * The currently-recognised architect slots, {@code terminal_id → slot name} (CB-548).
     *
     * <p>Read from the same supplier {@link #resolve} consults, so a slot that is <em>listed</em>
     * here but would not <em>resolve</em> (or the reverse) cannot drift apart. Live for the same
     * reason as {@link #leads()}.
     */
    public Map<String, String> members() {
        return architectTerminals.get();
    }

    /**
     * Resolve the caller of a request.
     *
     * @param remoteAddr           the connection's remote address
     * @param remotePort           the connection's remote port (used for the peer-PID lookup)
     * @param authorizationHeader  the raw {@code Authorization} header, or {@code null}
     */
    public Principal resolve(String remoteAddr, int remotePort, String authorizationHeader) {
        ConnectionIdentity.Caller c = identity.resolve(remoteAddr, remotePort);
        if (c.terminal() != null) {
            String lead = leadTerminals.get().get(c.terminal());
            if (lead != null) {
                // The config names this pane as a lead's own. The pane mapping is exactly as
                // unforgeable as a worker's, so it outranks the token path — no credential needed.
                // Checked before the architect registry so a pane named in BOTH is still the lead
                // (CB-548 preserves every existing leader behaviour).
                return Principal.leader(lead, c.terminal(), c.pid());
            }
            String slot = architectTerminals.get().get(c.terminal());
            if (slot != null && memberSlotRoles.apply(slot) == MemberRole.ARCHITECT) {
                // The config/live binding names this pane as an architect slot's own. Same
                // unforgeable pane mapping; the live binding, never a request argument, decides.
                // Check the slot role too: this defence in depth prevents a bad lifecycle bind from
                // escalating a dev or reviewer into an architect. Checked before the worker fallback.
                return Principal.architect(memberSlotNames.apply(slot), c.terminal(), c.pid());
            }
            return Principal.worker(c.terminal(), c.pid()); // unforgeable; never token-gated
        }

        if (tokenMode) {
            return presentedTokenMatches(authorizationHeader)
                    ? Principal.primary(c.pid())
                    : Principal.anonymous();
        }

        // loopback-trust: same-host callers that are not workers are the primary. A non-loopback
        // caller is anonymous even here — and startup refuses that combination anyway
        // (FleetConfig.validateAuthExposure), so this is defence in depth, not the control.
        return isLoopback(remoteAddr) ? Principal.primary(c.pid()) : Principal.anonymous();
    }

    private boolean presentedTokenMatches(String authorizationHeader) {
        String presented = bearerValue(authorizationHeader);
        if (presented == null) {
            return false;
        }
        // Constant-time: MessageDigest.isEqual does not short-circuit on the first differing byte,
        // so a token cannot be recovered a byte at a time by timing the response.
        return MessageDigest.isEqual(presented.getBytes(StandardCharsets.UTF_8), expectedToken);
    }

    /** Extract the credential from {@code Authorization: Bearer <token>}, or {@code null}. */
    private static String bearerValue(String header) {
        if (header == null) {
            return null;
        }
        String h = header.trim();
        if (h.length() < 7 || !h.regionMatches(true, 0, "Bearer ", 0, 7)) {
            return null;
        }
        String token = h.substring(7).trim();
        return token.isEmpty() ? null : token;
    }

    private static boolean isLoopback(String remoteAddr) {
        if (remoteAddr == null) {
            return false;
        }
        return remoteAddr.equals("127.0.0.1") || remoteAddr.equals("::1")
                || remoteAddr.equals("0:0:0:0:0:0:0:1") || remoteAddr.startsWith("127.");
    }
}