reactive-mail-server/backend/src/main/java/com/mailserver/controller/SsoController.java

package com.mailserver.controller;

import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.mailserver.config.JambotronSso;
import com.mailserver.config.MailAdmins;
import com.mailserver.domain.User;
import com.mailserver.service.UserService;
import com.mailserver.util.JwtUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.reactive.function.client.WebClientResponseException;
import reactor.core.publisher.Mono;

import java.net.URI;
import java.net.URLEncoder;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.security.SecureRandom;
import java.time.Duration;
import java.time.Instant;
import java.util.Base64;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * «Sign in with Jambotron» — OAuth 2.0 authorization code + PKCE (S256):
 * <ol>
 *   <li>{@code /api/oauth/start} → browser to {@code JAMBOTRON_URL/oauth/authorize?client_id=mail&…}</li>
 *   <li>Jambotron logs the person in (Google / GitHub / password), only staff get a one-time code</li>
 *   <li>{@code /api/oauth/callback?code&state} → code exchanged server to server at {@code /api/oauth/token}
 *       (Bearer SERVICE_TOKEN + code_verifier) → who it is, roles, their work mailbox</li>
 *   <li>session of that mailbox → {@code /sso#token=…}; ADMIN of Jambotron gets role ADMIN here</li>
 * </ol>
 */
@RestController
@RequestMapping("/api/oauth")
public class SsoController {

    private static final Logger log = LoggerFactory.getLogger(SsoController.class);
    private static final Duration STATE_TTL = Duration.ofMinutes(10);
    private static final SecureRandom RANDOM = new SecureRandom();

    /** What Jambotron says about the person (its /api/oauth/token). */
    @JsonIgnoreProperties(ignoreUnknown = true)
    public record Identity(String sub, String username, String email, List<String> roles, String mailbox) {
    }

    private record Pending(String verifier, Instant until) {
    }

    private final Map<String, Pending> pending = new ConcurrentHashMap<>();
    private final JambotronSso sso;
    private final UserService users;
    private final MailAdmins admins;
    private final JwtUtils jwt;
    private final WebClient http;

    public SsoController(JambotronSso sso, UserService users, MailAdmins admins, JwtUtils jwt, WebClient.Builder http) {
        this.sso = sso;
        this.users = users;
        this.admins = admins;
        this.jwt = jwt;
        this.http = http.build();
    }

    @GetMapping("/start")
    public ResponseEntity<Void> start() {
        if (!sso.enabled()) {
            return back("/login?sso_error=" + enc("Sign in with Jambotron is not set up on this server"));
        }
        Instant now = Instant.now();
        pending.values().removeIf(p -> p.until().isBefore(now));
        String state = random(24);
        String verifier = random(48);
        pending.put(state, new Pending(verifier, now.plus(STATE_TTL)));
        return back(sso.url() + "/oauth/authorize?response_type=code&client_id=mail"
                + "&redirect_uri=" + enc(sso.callback()) + "&state=" + state
                + "&code_challenge=" + challenge(verifier) + "&code_challenge_method=S256");
    }

    @GetMapping("/callback")
    public Mono<ResponseEntity<Void>> callback(@RequestParam(required = false) String code,
                                               @RequestParam(required = false) String state,
                                               @RequestParam(required = false) String error) {
        Pending p = state == null ? null : pending.remove(state);
        if (error != null) {
            return Mono.just(fail(error));
        }
        if (code == null || p == null || p.until().isBefore(Instant.now())) {
            return Mono.just(fail("The sign-in took too long or was not started here — try again"));
        }
        return http.post().uri(sso.url() + "/api/oauth/token")
                .header("Authorization", "Bearer " + sso.token())
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(Map.of("grant_type", "authorization_code", "client_id", "mail", "code", code,
                        "code_verifier", p.verifier(), "redirect_uri", sso.callback()))
                .retrieve()
                .bodyToMono(Identity.class)
                .timeout(Duration.ofSeconds(15))
                .flatMap(this::session)
                .onErrorResume(WebClientResponseException.class, e -> {
                    log.warn("Sign in with Jambotron: token exchange refused {} {}", e.getStatusCode(), e.getResponseBodyAsString());
                    return Mono.just(fail("Jambotron did not confirm the sign-in (" + e.getStatusCode().value() + ")"));
                })
                .onErrorResume(e -> !(e instanceof SsoRefused), e -> {
                    log.warn("Sign in with Jambotron failed: {}", e.toString());
                    return Mono.just(fail("Jambotron is not reachable — try again later"));
                })
                .onErrorResume(SsoRefused.class, e -> Mono.just(fail(e.getMessage())));
    }

    private Mono<ResponseEntity<Void>> session(Identity who) {
        if (who.mailbox() == null || who.mailbox().isBlank()) {
            return Mono.error(new SsoRefused("You have no work mailbox: mail is for Jambotron staff"));
        }
        boolean admin = who.roles() != null && who.roles().contains("ROLE_ADMIN");
        return users.findOptional(who.mailbox())
                .switchIfEmpty(Mono.error(new SsoRefused("Your mailbox " + who.mailbox() + " is not created yet — it appears within 10 minutes")))
                .flatMap(user -> {
                    if (!user.isActive() || user.isShared() || user.isSystem()) {
                        return Mono.error(new SsoRefused("Mailbox " + who.mailbox() + " is switched off"));
                    }
                    return admin && !user.isAdmin() ? users.setAdmin(user, true) : Mono.just(user);
                })
                .map(user -> {
                    if (user.isAdmin()) {
                        admins.grant(user.getEmail(), true);
                    }
                    log.info("WEB: {} signed in with Jambotron as {} ({})", user.getEmail(), who.username(), who.email());
                    return back("/sso#token=" + enc(jwt.generateToken(user.getEmail())) + "&email=" + enc(user.getEmail()));
                });
    }

    private static final class SsoRefused extends RuntimeException {
        SsoRefused(String message) {
            super(message);
        }
    }

    private static ResponseEntity<Void> fail(String message) {
        return back("/login?sso_error=" + enc(message));
    }

    private static ResponseEntity<Void> back(String location) {
        return ResponseEntity.status(HttpStatus.FOUND).location(URI.create(location)).build();
    }

    private static String enc(String value) {
        return URLEncoder.encode(value, StandardCharsets.UTF_8);
    }

    private static String random(int bytes) {
        byte[] b = new byte[bytes];
        RANDOM.nextBytes(b);
        return Base64.getUrlEncoder().withoutPadding().encodeToString(b);
    }

    static String challenge(String verifier) {
        try {
            byte[] hash = MessageDigest.getInstance("SHA-256").digest(verifier.getBytes(StandardCharsets.US_ASCII));
            return Base64.getUrlEncoder().withoutPadding().encodeToString(hash);
        } catch (NoSuchAlgorithmException e) {
            throw new IllegalStateException(e);
        }
    }
}