reactive-mail-server/backend/src/main/java/com/mailserver/smtp/SmtpServer.java

package com.mailserver.smtp;

import io.netty.handler.codec.LineBasedFrameDecoder;
import io.netty.handler.timeout.ReadTimeoutHandler;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;
import reactor.netty.DisposableServer;
import reactor.netty.tcp.TcpServer;

import java.io.ByteArrayOutputStream;
import java.net.InetSocketAddress;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.time.ZonedDateTime;
import java.time.format.DateTimeFormatter;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;
import java.util.Locale;
import java.util.concurrent.TimeUnit;

/**
 * SMTP server (RFC 5321) on Reactor Netty.
 * <ul>
 *   <li>input is split into lines by {@link LineBasedFrameDecoder} — a network packet may hold half a line or many
 *       pipelined commands, both work;</li>
 *   <li>AUTH PLAIN / LOGIN against our mailboxes; someone else's address is accepted only from a logged-in sender
 *       (no open relay); a logged-in sender may use only his own address in MAIL FROM;</li>
 *   <li>the message is stored or queued before the client gets 250 — on failure 451 and the client retries;</li>
 *   <li>database and DNS work runs on a worker thread, not on the event loop.</li>
 * </ul>
 * No STARTTLS yet: AUTH is offered only when {@code authWithoutTls} is on (a closed network: inside a VM or a docker network).
 */
public class SmtpServer {

    private static final Logger log = LoggerFactory.getLogger(SmtpServer.class);

    // RFC 5321 4.5.3.1.6: a text line is at most 1000 octets with CRLF; we allow more for sloppy clients
    static final int MAX_LINE = 64 * 1024;
    static final int MAX_RECIPIENTS = 100;
    private static final int MAX_BAD_COMMANDS = 20;

    public record Settings(String host, int port, String hostname, long maxMessageBytes, boolean authWithoutTls,
                           Duration idleTimeout) {
    }

    private final Settings settings;
    private final SmtpBackend backend;
    private DisposableServer server;

    public SmtpServer(Settings settings, SmtpBackend backend) {
        this.settings = settings;
        this.backend = backend;
    }

    public void start() {
        server = TcpServer.create()
                .host(settings.host())
                .port(settings.port())
                .doOnConnection(connection -> connection
                        .addHandlerFirst("smtp-idle", new ReadTimeoutHandler(settings.idleTimeout().toSeconds(), TimeUnit.SECONDS))
                        .addHandlerLast("smtp-lines", new LineBasedFrameDecoder(MAX_LINE, true, false)))
                .handle((inbound, outbound) -> {
                    String[] remoteIp = {"unknown"};
                    inbound.withConnection(c -> {
                        if (c.address() instanceof InetSocketAddress address && address.getAddress() != null) {
                            remoteIp[0] = address.getAddress().getHostAddress();
                        }
                    });
                    Session session = new Session(remoteIp[0]);
                    Flux<String> replies = inbound.receive().asByteArray()
                            .concatMap(session::onLine)
                            .onErrorResume(e -> {
                                log.debug("SMTP {}: connection closed: {}", session.remoteIp, e.toString());
                                return Mono.just("421 4.4.2 " + settings.hostname() + " closing connection\r\n");
                            })
                            .takeUntil(reply -> session.closing);
                    return outbound.sendString(Flux.concat(
                            Mono.just("220 " + settings.hostname() + " ESMTP ready\r\n"), replies), StandardCharsets.US_ASCII);
                })
                .bindNow(Duration.ofSeconds(10));
        log.info("SMTP server listening on {}:{} ({})", settings.host(), port(), settings.hostname());
    }

    public int port() {
        return server.port();
    }

    public void stop() {
        if (server != null) {
            server.disposeNow(Duration.ofSeconds(5));
        }
    }

    private enum State { COMMAND, DATA, AUTH_PLAIN, AUTH_LOGIN_USER, AUTH_LOGIN_PASSWORD }

    private final class Session {
        final String remoteIp;
        State state = State.COMMAND;
        String helo;
        boolean extended;
        String authenticatedUser;
        String pendingAuthUser;
        String mailFrom;
        final List<SmtpBackend.Recipient> recipients = new ArrayList<>();
        final ByteArrayOutputStream data = new ByteArrayOutputStream();
        boolean dataTooBig;
        int badCommands;
        volatile boolean closing;

        Session(String remoteIp) {
            this.remoteIp = remoteIp;
        }

        Mono<String> onLine(byte[] raw) {
            byte[] line = stripCr(raw);
            return switch (state) {
                case DATA -> onDataLine(line);
                case AUTH_PLAIN -> authPlain(ascii(line));
                case AUTH_LOGIN_USER -> authLoginUser(ascii(line));
                case AUTH_LOGIN_PASSWORD -> authLoginPassword(ascii(line));
                case COMMAND -> onCommand(ascii(line));
            };
        }

        private Mono<String> onCommand(String line) {
            String[] parts = line.trim().split("\\s+", 2);
            String verb = parts[0].toUpperCase(Locale.ROOT);
            String arg = parts.length > 1 ? parts[1].trim() : "";
            return switch (verb) {
                case "EHLO", "HELO" -> hello(verb.equals("EHLO"), arg);
                case "AUTH" -> auth(arg);
                case "MAIL" -> mail(arg);
                case "RCPT" -> rcpt(arg);
                case "DATA" -> dataStart();
                case "RSET" -> {
                    resetTransaction();
                    yield reply("250 2.0.0 OK");
                }
                case "NOOP" -> reply("250 2.0.0 OK");
                case "VRFY" -> reply("252 2.5.0 Cannot VRFY user, but will accept message");
                case "STARTTLS" -> reply("454 4.7.0 TLS not available");
                case "QUIT" -> {
                    closing = true;
                    yield reply("221 2.0.0 " + settings.hostname() + " Bye");
                }
                case "" -> reply("500 5.5.2 Empty command");
                default -> bad("500 5.5.1 Command not recognized");
            };
        }

        private Mono<String> hello(boolean ehlo, String arg) {
            if (arg.isEmpty()) {
                return reply("501 5.5.4 Hostname required");
            }
            helo = arg;
            extended = ehlo;
            resetTransaction();
            if (!ehlo) {
                return reply("250 " + settings.hostname());
            }
            StringBuilder sb = new StringBuilder();
            sb.append("250-").append(settings.hostname()).append(" Hello ").append(arg).append("\r\n");
            sb.append("250-SIZE ").append(settings.maxMessageBytes()).append("\r\n");
            sb.append("250-8BITMIME\r\n");
            sb.append("250-PIPELINING\r\n");
            if (settings.authWithoutTls()) {
                sb.append("250-AUTH PLAIN LOGIN\r\n");
            }
            sb.append("250 ENHANCEDSTATUSCODES\r\n");
            return Mono.just(sb.toString());
        }

        private Mono<String> auth(String arg) {
            if (!settings.authWithoutTls()) {
                return reply("538 5.7.11 Encryption required for requested authentication mechanism");
            }
            if (helo == null || !extended) {
                return reply("503 5.5.1 EHLO first");
            }
            if (authenticatedUser != null) {
                return reply("503 5.5.1 Already authenticated");
            }
            if (mailFrom != null) {
                return reply("503 5.5.1 AUTH not allowed during a mail transaction");
            }
            String[] parts = arg.split("\\s+", 2);
            String mechanism = parts[0].toUpperCase(Locale.ROOT);
            String initial = parts.length > 1 ? parts[1] : null;
            if (mechanism.equals("PLAIN")) {
                if (initial == null) {
                    state = State.AUTH_PLAIN;
                    return reply("334 ");
                }
                return authPlain(initial);
            }
            if (mechanism.equals("LOGIN")) {
                if (initial != null) {
                    return authLoginUser(initial);
                }
                state = State.AUTH_LOGIN_USER;
                return reply("334 " + b64("Username:"));
            }
            return reply("504 5.5.4 Unrecognized authentication type");
        }

        private Mono<String> authPlain(String response) {
            state = State.COMMAND;
            if (response.equals("*")) {
                return reply("501 5.0.0 Authentication cancelled");
            }
            String decoded = unb64(response);
            // authzid \0 authcid \0 password
            String[] fields = decoded == null ? new String[0] : decoded.split("\0", -1);
            if (fields.length != 3) {
                return reply("501 5.5.2 Cannot decode AUTH PLAIN response");
            }
            return checkPassword(fields[1], fields[2]);
        }

        private Mono<String> authLoginUser(String response) {
            if (response.equals("*")) {
                state = State.COMMAND;
                return reply("501 5.0.0 Authentication cancelled");
            }
            pendingAuthUser = unb64(response);
            if (pendingAuthUser == null) {
                state = State.COMMAND;
                return reply("501 5.5.2 Cannot decode username");
            }
            state = State.AUTH_LOGIN_PASSWORD;
            return reply("334 " + b64("Password:"));
        }

        private Mono<String> authLoginPassword(String response) {
            state = State.COMMAND;
            String user = pendingAuthUser;
            pendingAuthUser = null;
            if (response.equals("*")) {
                return reply("501 5.0.0 Authentication cancelled");
            }
            String password = unb64(response);
            if (password == null) {
                return reply("501 5.5.2 Cannot decode password");
            }
            return checkPassword(user, password);
        }

        private Mono<String> checkPassword(String user, String password) {
            String username = user.trim().toLowerCase(Locale.ROOT);
            return blocking(() -> backend.authenticate(username, password))
                    .map(ok -> {
                        if (ok) {
                            authenticatedUser = username;
                            log.info("SMTP {}: authenticated as {}", remoteIp, username);
                            return "235 2.7.0 Authentication successful\r\n";
                        }
                        log.warn("SMTP {}: authentication failed for {}", remoteIp, username);
                        badCommands++;
                        return "535 5.7.8 Authentication credentials invalid\r\n";
                    });
        }

        private Mono<String> mail(String arg) {
            if (helo == null) {
                return reply("503 5.5.1 EHLO/HELO first");
            }
            if (mailFrom != null) {
                return reply("503 5.5.1 Sender already specified");
            }
            if (!arg.regionMatches(true, 0, "FROM:", 0, 5)) {
                return reply("501 5.5.4 Syntax: MAIL FROM:<address>");
            }
            String rest = arg.substring(5).trim();
            String address = path(rest);
            if (address == null) {
                return reply("501 5.1.7 Bad sender address syntax");
            }
            Long size = sizeParam(rest);
            if (size != null && size > settings.maxMessageBytes()) {
                return reply("552 5.3.4 Message size exceeds fixed limit");
            }
            if (authenticatedUser != null && !address.isEmpty() && !address.equals(authenticatedUser)) {
                return reply("553 5.7.1 Sender address not owned by authenticated user");
            }
            mailFrom = address;
            return reply("250 2.1.0 Sender OK");
        }

        private Mono<String> rcpt(String arg) {
            if (mailFrom == null) {
                return reply("503 5.5.1 MAIL first");
            }
            if (!arg.regionMatches(true, 0, "TO:", 0, 3)) {
                return reply("501 5.5.4 Syntax: RCPT TO:<address>");
            }
            String address = path(arg.substring(3).trim());
            if (address == null || address.isEmpty() || address.indexOf('@') < 1) {
                return reply("501 5.1.3 Bad recipient address syntax");
            }
            if (recipients.size() >= MAX_RECIPIENTS) {
                return reply("452 4.5.3 Too many recipients");
            }
            String user = authenticatedUser;
            return blocking(() -> backend.route(address, user))
                    .map(route -> switch (route) {
                        case LOCAL, OUTBOUND -> {
                            recipients.add(new SmtpBackend.Recipient(address, route));
                            yield "250 2.1.5 Recipient OK\r\n";
                        }
                        case UNKNOWN_USER -> "550 5.1.1 No such user here\r\n";
                        case RELAY_DENIED -> {
                            log.warn("SMTP {}: relay denied {} -> {}", remoteIp, mailFrom, address);
                            yield "550 5.7.1 Relaying denied: authentication required\r\n";
                        }
                        case OUTBOUND_DISABLED -> "550 5.7.1 Outbound delivery is disabled on this server\r\n";
                    });
        }

        private Mono<String> dataStart() {
            if (mailFrom == null) {
                return reply("503 5.5.1 MAIL first");
            }
            if (recipients.isEmpty()) {
                return reply("554 5.5.1 No valid recipients");
            }
            state = State.DATA;
            data.reset();
            dataTooBig = false;
            return reply("354 End data with <CR><LF>.<CR><LF>");
        }

        private Mono<String> onDataLine(byte[] line) {
            if (line.length == 1 && line[0] == '.') {
                state = State.COMMAND;
                if (dataTooBig) {
                    resetTransaction();
                    return reply("552 5.3.4 Message size exceeds fixed limit");
                }
                SmtpBackend.Envelope envelope = new SmtpBackend.Envelope(
                        remoteIp, helo, mailFrom, authenticatedUser, List.copyOf(recipients));
                byte[] message = withReceivedHeader(data.toByteArray());
                resetTransaction();
                return blocking(() -> {
                    backend.accept(envelope, message);
                    return true;
                })
                        .map(ok -> "250 2.0.0 Message accepted\r\n")
                        .onErrorResume(e -> {
                            log.error("SMTP {}: message from {} not stored: {}", remoteIp, envelope.mailFrom(), e.toString());
                            return Mono.just("451 4.3.0 Message not stored, try again later\r\n");
                        });
            }
            if (!dataTooBig) {
                int offset = line.length > 0 && line[0] == '.' ? 1 : 0; // dot-unstuffing (RFC 5321 4.5.2)
                data.write(line, offset, line.length - offset);
                data.write('\r');
                data.write('\n');
                if (data.size() > settings.maxMessageBytes()) {
                    dataTooBig = true;
                    data.reset();
                }
            }
            return Mono.empty();
        }

        private byte[] withReceivedHeader(byte[] message) {
            String received = "Received: from " + helo + " ([" + remoteIp + "])\r\n"
                    + "\tby " + settings.hostname() + " with " + (authenticatedUser != null ? "ESMTPA" : "ESMTP") + ";\r\n"
                    + "\t" + DateTimeFormatter.RFC_1123_DATE_TIME.format(ZonedDateTime.now()) + "\r\n";
            byte[] header = received.getBytes(StandardCharsets.US_ASCII);
            byte[] all = new byte[header.length + message.length];
            System.arraycopy(header, 0, all, 0, header.length);
            System.arraycopy(message, 0, all, header.length, message.length);
            return all;
        }

        private void resetTransaction() {
            mailFrom = null;
            recipients.clear();
            data.reset();
            dataTooBig = false;
        }

        private Mono<String> bad(String text) {
            if (++badCommands > MAX_BAD_COMMANDS) {
                closing = true;
                return reply("421 4.7.0 Too many errors, closing connection");
            }
            return reply(text);
        }
    }

    private static Mono<String> reply(String line) {
        return Mono.just(line + "\r\n");
    }

    private static <T> Mono<T> blocking(java.util.concurrent.Callable<T> work) {
        return Mono.fromCallable(work).subscribeOn(Schedulers.boundedElastic());
    }

    /** "<a@b> SIZE=1" → "a@b"; "<>" → "" (null sender of bounces); malformed → null. */
    static String path(String text) {
        int start = text.indexOf('<');
        int end = text.indexOf('>', start + 1);
        if (start != 0 || end < 0) {
            return null;
        }
        String address = text.substring(1, end).trim();
        int colon = address.indexOf(':'); // source route "@a,@b:user@host" (RFC 5321 4.1.2) — ignored
        if (address.startsWith("@") && colon > 0) {
            address = address.substring(colon + 1);
        }
        if (address.contains(" ") || address.contains("<")) {
            return null;
        }
        return address.toLowerCase(Locale.ROOT);
    }

    private static Long sizeParam(String text) {
        for (String param : text.substring(text.indexOf('>') + 1).trim().split("\\s+")) {
            if (param.regionMatches(true, 0, "SIZE=", 0, 5)) {
                try {
                    return Long.parseLong(param.substring(5));
                } catch (NumberFormatException e) {
                    return null;
                }
            }
        }
        return null;
    }

    private static byte[] stripCr(byte[] line) {
        if (line.length > 0 && line[line.length - 1] == '\r') {
            byte[] copy = new byte[line.length - 1];
            System.arraycopy(line, 0, copy, 0, copy.length);
            return copy;
        }
        return line;
    }

    private static String ascii(byte[] line) {
        return new String(line, StandardCharsets.UTF_8);
    }

    private static String b64(String text) {
        return Base64.getEncoder().encodeToString(text.getBytes(StandardCharsets.UTF_8));
    }

    private static String unb64(String text) {
        try {
            return new String(Base64.getDecoder().decode(text.trim()), StandardCharsets.UTF_8);
        } catch (IllegalArgumentException e) {
            return null;
        }
    }
}