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

package com.mailserver.controller;

import com.mailserver.config.MailAdmins;
import com.mailserver.config.MailDomains;
import com.mailserver.domain.User;
import com.mailserver.dto.LoginRequest;
import com.mailserver.dto.LoginResponse;
import com.mailserver.dto.PasswordChangeRequest;
import com.mailserver.dto.ProfileResponse;
import com.mailserver.dto.ProfileUpdateRequest;
import com.mailserver.dto.RegisterRequest;
import com.mailserver.service.MailService;
import com.mailserver.service.UserService;
import com.mailserver.security.LoginGuard;
import com.mailserver.util.JwtUtils;
import jakarta.validation.Valid;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpStatus;
import org.springframework.http.ResponseEntity;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.security.core.Authentication;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.PutMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.server.ResponseStatusException;
import reactor.core.publisher.Mono;

import java.net.InetSocketAddress;
import java.time.Duration;
import java.util.Locale;
import java.util.Map;
import java.util.Set;

/** Login, registration, the logged-in user's profile and password; what the login page needs to know. */
@RestController
@RequestMapping("/api")
public class AuthenticationController {

    private static final Logger log = LoggerFactory.getLogger(AuthenticationController.class);
    /** Marker for "no such mailbox or wrong password": both look the same from outside. */
    private static final User NO_USER = new User();

    /** Service addresses nobody can register (RFC 2142 and ours). */
    static final Set<String> RESERVED = Set.of("postmaster", "abuse", "mailer-daemon", "hostmaster", "webmaster",
            "root", "admin", "noreply", "no-reply", "security", "support");

    private final UserService userService;
    private final MailService mailService;
    private final JwtUtils jwtUtils;
    private final MailDomains domains;
    private final boolean registrationEnabled;
    private final LoginGuard guard;
    private final MailAdmins admins;

    public AuthenticationController(UserService userService, MailService mailService, JwtUtils jwtUtils, MailDomains domains,
                                    @Value("${mailserver.registration.enabled:false}") boolean registrationEnabled,
                                    LoginGuard guard, MailAdmins admins) {
        this.guard = guard;
        this.admins = admins;
        this.userService = userService;
        this.mailService = mailService;
        this.jwtUtils = jwtUtils;
        this.domains = domains;
        this.registrationEnabled = registrationEnabled;
    }

    /** Public: is registration open, and in which domains mailboxes can be created. */
    @GetMapping("/config")
    public Map<String, Object> config() {
        return Map.of("registrationEnabled", registrationEnabled, "domains", domains.domains().stream().sorted().toList());
    }

    @PostMapping("/login")
    public Mono<ResponseEntity<?>> login(@Valid @RequestBody Mono<LoginRequest> requestMono, ServerHttpRequest http) {
        String ip = clientIp(http);
        return requestMono.flatMap(request -> {
            String email = request.getEmail().trim().toLowerCase(Locale.ROOT);
            Duration wait = guard.retryAfter(ip, email);
            if (!wait.isZero()) {
                log.warn("WEB {}: login as {} refused, blocked for {}", ip, email, LoginGuard.human(wait));
                return Mono.just(tooMany(wait));
            }
            return userService.findByEmail(email)
                .map(user -> userService.checkPassword(user, request.getPassword()) ? user : NO_USER)
                .onErrorResume(IllegalArgumentException.class, e -> Mono.just(NO_USER))
                .defaultIfEmpty(NO_USER)
                .flatMap(user -> {
                    if (user == NO_USER) {
                        guard.failure(ip, email);
                        log.warn("WEB {}: authentication failed for {}", ip, email);
                        Mono<ResponseEntity<?>> answer = Mono.just(ResponseEntity.badRequest().body(Map.of("error", "Invalid credentials")));
                        Duration delay = guard.failureDelay();
                        return delay.isZero() ? answer : answer.delayElement(delay);
                    }
                    guard.success(ip, email);
                    if (!user.isActive()) {
                        return Mono.just(ResponseEntity.badRequest().body(Map.of("error", "User is deactivated")));
                    }
                    return Mono.just(ResponseEntity.ok(token(user)));
                });
        });
    }

    /** New mailbox in one of our domains; logs in right away (the answer carries the token). */
    @PostMapping("/register")
    public Mono<ResponseEntity<?>> register(@Valid @RequestBody Mono<RegisterRequest> requestMono) {
        if (!registrationEnabled) {
            return Mono.just(ResponseEntity.status(HttpStatus.FORBIDDEN).body(Map.of("error", "Registration is closed")));
        }
        return requestMono.flatMap(request -> {
            String email = request.getEmail().trim().toLowerCase(Locale.ROOT);
            String local = email.substring(0, email.indexOf('@'));
            if (!domains.isLocal(email)) {
                return Mono.just(ResponseEntity.badRequest()
                        .body(Map.of("error", "Mailboxes can only be created in: " + String.join(", ", domains.domains()))));
            }
            if (RESERVED.contains(local) || !local.matches("[a-z0-9]([a-z0-9._-]{0,62}[a-z0-9])?")) {
                return Mono.just(ResponseEntity.badRequest()
                        .body(Map.of("error", "This name is not available: letters, digits, dot, dash, underscore")));
            }
            return userService.register(email, request.getPassword(), request.getDisplayName())
                    .map(user -> ResponseEntity.status(HttpStatus.CREATED).body((Object) token(user)))
                    .onErrorResume(IllegalArgumentException.class, e ->
                            Mono.just(ResponseEntity.badRequest().body(Map.of("error", "This mailbox already exists"))));
        });
    }

    @GetMapping("/me")
    public Mono<ProfileResponse> me(Mono<Authentication> auth) {
        return user(auth).flatMap(this::profile);
    }

    @PutMapping("/me")
    public Mono<ProfileResponse> updateMe(Mono<Authentication> auth, @Valid @RequestBody Mono<ProfileUpdateRequest> body) {
        return user(auth).zipWith(body)
                .flatMap(pair -> userService.updateDisplayName(pair.getT1(), pair.getT2().displayName()))
                .flatMap(this::profile);
    }

    @PutMapping("/me/password")
    public Mono<Map<String, String>> changePassword(Mono<Authentication> auth, @Valid @RequestBody Mono<PasswordChangeRequest> body,
                                                    ServerHttpRequest http) {
        String ip = clientIp(http);
        return user(auth).zipWith(body)
                .flatMap(pair -> {
                    User user = pair.getT1();
                    // Someone with a stolen session must not guess the current password either
                    Duration wait = guard.retryAfter(ip, user.getEmail());
                    if (!wait.isZero()) {
                        return Mono.error(new ResponseStatusException(HttpStatus.TOO_MANY_REQUESTS,
                                "Too many failed attempts, try again in " + LoginGuard.human(wait)));
                    }
                    if (!userService.checkPassword(user, pair.getT2().currentPassword())) {
                        guard.failure(ip, user.getEmail());
                    }
                    return userService.changePassword(user, pair.getT2().currentPassword(), pair.getT2().newPassword());
                })
                .thenReturn(Map.of("message", "Password changed"));
    }

    private static ResponseEntity<?> tooMany(Duration wait) {
        return ResponseEntity.status(HttpStatus.TOO_MANY_REQUESTS)
                .header(HttpHeaders.RETRY_AFTER, String.valueOf(Math.max(1, (wait.toMillis() + 999) / 1000)))
                .body(Map.of("error", "Too many failed logins, try again in " + LoginGuard.human(wait)));
    }

    /** Client address; behind Caddy it comes from X-Forwarded-For (server.forward-headers-strategy=framework). */
    static String clientIp(ServerHttpRequest http) {
        InetSocketAddress address = http.getRemoteAddress();
        if (address == null) {
            return "unknown";
        }
        return address.getAddress() != null ? address.getAddress().getHostAddress() : address.getHostString();
    }

    private Mono<ProfileResponse> profile(User user) {
        return mailService.usedBytes(user.getId())
                .map(used -> new ProfileResponse(user.getEmail(), user.getDisplayName(), used,
                        user.getQuotaBytes() == null ? 0 : user.getQuotaBytes(), user.getCreatedAt(),
                        admins.isAdmin(user.getEmail())));
    }

    private Mono<User> user(Mono<Authentication> auth) {
        return auth.switchIfEmpty(Mono.error(new ResponseStatusException(HttpStatus.UNAUTHORIZED)))
                .flatMap(a -> userService.findByEmail(a.getName()))
                .onErrorMap(IllegalArgumentException.class, e -> new ResponseStatusException(HttpStatus.UNAUTHORIZED));
    }

    private LoginResponse token(User user) {
        LoginResponse response = new LoginResponse();
        response.setToken(jwtUtils.generateToken(user.getEmail()));
        response.setEmail(user.getEmail());
        return response;
    }

}