reactive-mail-server/backend/src/main/java/com/mailserver/imap/ImapServer.java

package com.mailserver.imap;

import com.mailserver.security.LoginGuard;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelPipeline;
import io.netty.handler.codec.ByteToMessageDecoder;
import io.netty.handler.codec.TooLongFrameException;
import io.netty.handler.ssl.SslContext;
import io.netty.handler.ssl.SslHandler;
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.publisher.Sinks;
import reactor.core.scheduler.Schedulers;
import reactor.netty.Connection;
import reactor.netty.DisposableServer;
import reactor.netty.NettyOutbound;
import reactor.netty.tcp.TcpServer;

import java.io.ByteArrayOutputStream;
import java.net.InetSocketAddress;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.List;
import java.util.Locale;
import java.util.concurrent.TimeUnit;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * IMAP server (RFC 3501) on Reactor Netty.
 * <ul>
 *   <li>{@link CommandDecoder} collects whole commands with their literals ({n} — after "+ Ready", {n+} — at once);</li>
 *   <li>each command runs on a worker thread (database, MIME parsing), never on the event loop; responses are written
 *       in order, one command at a time;</li>
 *   <li>TLS: STARTTLS, or implicit TLS on a separate port (993); LOGIN/AUTHENTICATE only inside TLS unless
 *       {@code authWithoutTls}; plain text sent together with STARTTLS is never executed;</li>
 *   <li>IDLE: changes of the selected mailbox are pushed every {@code idlePoll} until DONE.</li>
 * </ul>
 */
public class ImapServer {

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

    static final int MAX_LINE = 64 * 1024;
    private static final int MAX_QUEUED = 1000;
    private static final Pattern LITERAL = Pattern.compile("\\{(\\d+)(\\+?)}\r?\n$");
    private static final Pattern IDLE = Pattern.compile("^(\\S+) IDLE\\s*$", Pattern.CASE_INSENSITIVE);
    private static final Duration IDLE_MAX = Duration.ofMinutes(29);

    /**
     * @param tls            certificate; null — no TLS
     * @param implicitTls    TLS from the first byte (port 993)
     * @param maxLiteral     the largest literal accepted (APPEND of a big message)
     * @param idlePoll       how often IDLE looks for changes
     */
    public record Settings(String host, int port, String hostname, boolean authWithoutTls, SslContext tls,
                           boolean implicitTls, int maxLiteral, Duration idleTimeout, Duration idlePoll) {
    }

    /** A whole command and whether it arrived over TLS. */
    record Command(byte[] bytes, boolean tls) {
    }

    private final Settings settings;
    private final ImapStore store;
    private final LoginGuard guard;
    private DisposableServer server;

    public ImapServer(Settings settings, ImapStore store) {
        this(settings, store, LoginGuard.disabled());
    }

    public ImapServer(Settings settings, ImapStore store, LoginGuard guard) {
        this.settings = settings;
        this.store = store;
        this.guard = guard;
    }

    public void start() {
        if (settings.implicitTls() && settings.tls() == null) {
            throw new IllegalStateException("Implicit TLS needs a certificate");
        }
        TcpServer tcp = TcpServer.create()
                .host(settings.host())
                .port(settings.port())
                .doOnConnection(connection -> connection
                        .addHandlerFirst("imap-idle", new ReadTimeoutHandler(settings.idleTimeout().toSeconds(), TimeUnit.SECONDS))
                        .addHandlerLast("imap-commands", new CommandDecoder(settings.maxLiteral())));
        if (settings.implicitTls()) {
            tcp = tcp.secure(spec -> spec.sslContext(settings.tls()));
        }
        server = tcp.handle((inbound, outbound) -> {
                    Connection[] connection = {null};
                    String[] remoteIp = {"unknown"};
                    inbound.withConnection(c -> {
                        connection[0] = c;
                        if (c.address() instanceof InetSocketAddress address && address.getAddress() != null) {
                            remoteIp[0] = address.getAddress().getHostAddress();
                        }
                    });
                    ImapSession session = new ImapSession(
                            new ImapSession.Settings(settings.hostname(), settings.tls() != null, settings.authWithoutTls()),
                            store, guard, remoteIp[0], settings.implicitTls());
                    @SuppressWarnings("unchecked")
                    Sinks.Empty<Void>[] idleDone = new Sinks.Empty[1];

                    Flux<byte[]> responses = inbound.receiveObject().cast(Command.class)
                            // DONE ends IDLE while the IDLE command is still running — handled before the queue
                            .filter(command -> {
                                if (session.idling() && isDone(command.bytes())) {
                                    if (idleDone[0] != null) {
                                        idleDone[0].tryEmitEmpty();
                                    }
                                    return false;
                                }
                                return true;
                            })
                            // read commands as they come (so DONE is seen during IDLE), run them one by one
                            .onBackpressureBuffer(MAX_QUEUED, dropped -> log.warn("IMAP {}: too many queued commands", session.remoteIp))
                            .concatMap(command -> {
                                if (session.tls && !command.tls()) {
                                    // plain text sent together with STARTTLS: never executed, the connection is closed
                                    log.warn("IMAP {}: plain-text command after STARTTLS, closing", session.remoteIp);
                                    session.closing = true;
                                    connection[0].dispose();
                                    return Flux.empty();
                                }
                                String text = new String(command.bytes(), StandardCharsets.UTF_8).trim();
                                Matcher idle = IDLE.matcher(text);
                                if (idle.matches()) {
                                    idleDone[0] = Sinks.empty();
                                    return idle(session, idle.group(1), idleDone[0]);
                                }
                                return blocking(() -> session.handle(command.bytes()));
                            })
                            .onErrorResume(e -> {
                                log.debug("IMAP {}: connection closed: {}", session.remoteIp, e.toString());
                                return Mono.just(ImapSession.line("* BYE " + e.getMessage()));
                            })
                            // after the tagged line of the command that closes (LOGOUT, too many failed logins), not after
                            // its first untagged "* BYE" — the tagged answer must still reach the client
                            .takeUntil(response -> session.closing && !untagged(response));

                    return Flux.concat(Mono.just(session.greeting()), responses)
                            .concatMap(response -> response == ImapSession.STARTTLS_READY
                                    ? startTls(connection[0], session)
                                    : send(outbound, response))
                            .then();
                })
                .bindNow(Duration.ofSeconds(10));
        log.info("IMAP server listening on {}:{} ({}{})", settings.host(), port(), settings.hostname(),
                settings.implicitTls() ? ", implicit TLS" : settings.tls() != null ? ", STARTTLS" : ", no TLS");
    }

    private Flux<byte[]> idle(ImapSession session, String tag, Sinks.Empty<Void> done) {
        Flux<byte[]> polls = Flux.interval(settings.idlePoll())
                .onBackpressureDrop()
                .concatMap(tick -> blocking(session::idlePoll))
                .takeUntilOther(done.asMono())
                .take(IDLE_MAX);
        return Flux.concat(
                blocking(() -> session.idleStart(tag)),
                Flux.defer(() -> session.idling() ? polls : Flux.empty()),
                blocking(session::idleEnd));
    }

    private static boolean isDone(byte[] bytes) {
        return new String(bytes, StandardCharsets.US_ASCII).trim().toUpperCase(Locale.ROOT).equals("DONE");
    }

    private static boolean untagged(byte[] response) {
        return response.length > 1 && response[0] == '*' && response[1] == ' ';
    }

    private static Flux<byte[]> blocking(java.util.concurrent.Callable<List<byte[]>> work) {
        return Mono.fromCallable(work).subscribeOn(Schedulers.boundedElastic()).flatMapIterable(list -> list);
    }

    private static Mono<Void> send(NettyOutbound outbound, byte[] response) {
        return outbound.send(Mono.just(Unpooled.wrappedBuffer(response))).then();
    }

    /** After the tagged OK of STARTTLS has left in plain text: TLS in front, a fresh decoder (no plain leftovers). */
    private Mono<Void> startTls(Connection connection, ImapSession session) {
        ChannelPipeline pipeline = connection.channel().pipeline();
        SslHandler ssl = settings.tls().newHandler(connection.channel().alloc());
        pipeline.replace("imap-commands", "imap-commands", new CommandDecoder(settings.maxLiteral()));
        pipeline.addFirst("imap-tls", ssl);
        return Mono.create(sink -> ssl.handshakeFuture().addListener(future -> {
            if (!future.isSuccess()) {
                log.info("IMAP {}: TLS handshake failed: {}", session.remoteIp, future.cause().toString());
                session.closing = true;
                connection.dispose();
            }
            sink.success();
        }));
    }

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

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

    /**
     * Collects one command: lines plus the literals they announce. A synchronizing literal {n} gets "+ Ready" first;
     * {n+} (LITERAL+) is read at once. Marks whether the command was read behind TLS.
     */
    static final class CommandDecoder extends ByteToMessageDecoder {
        private final int maxLiteral;
        private final ByteArrayOutputStream command = new ByteArrayOutputStream();
        private int literalLeft;
        private int lineBytes;

        CommandDecoder(int maxLiteral) {
            this.maxLiteral = maxLiteral;
        }

        @Override
        protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) {
            while (in.isReadable()) {
                if (literalLeft > 0) {
                    int n = Math.min(literalLeft, in.readableBytes());
                    byte[] chunk = new byte[n];
                    in.readBytes(chunk);
                    command.write(chunk, 0, n);
                    literalLeft -= n;
                    continue;
                }
                int newline = in.indexOf(in.readerIndex(), in.writerIndex(), (byte) '\n');
                if (newline < 0) {
                    if (lineBytes + in.readableBytes() > MAX_LINE) {
                        throw new TooLongFrameException("Command line too long");
                    }
                    return;
                }
                int length = newline - in.readerIndex() + 1;
                if (lineBytes + length > MAX_LINE) {
                    throw new TooLongFrameException("Command line too long");
                }
                byte[] line = new byte[length];
                in.readBytes(line);
                command.write(line, 0, line.length);
                lineBytes += length;
                Matcher literal = LITERAL.matcher(new String(line, StandardCharsets.ISO_8859_1));
                if (literal.find()) {
                    long size = Long.parseLong(literal.group(1));
                    if (size > maxLiteral) {
                        ctx.writeAndFlush(Unpooled.copiedBuffer("* BYE Literal too large\r\n", StandardCharsets.US_ASCII));
                        throw new TooLongFrameException("Literal too large: " + size);
                    }
                    literalLeft = (int) size;
                    if (literal.group(2).isEmpty()) {
                        ctx.writeAndFlush(Unpooled.copiedBuffer("+ Ready for literal data\r\n", StandardCharsets.US_ASCII));
                    }
                    continue;
                }
                out.add(new Command(command.toByteArray(), ctx.pipeline().get(SslHandler.class) != null));
                command.reset();
                lineBytes = 0;
            }
        }
    }
}