reactive-mail-server/backend/src/test/java/com/mailserver/smtp/SmtpServerTest.java

package com.mailserver.smtp;

import jakarta.mail.Message;
import jakarta.mail.Session;
import jakarta.mail.Transport;
import jakarta.mail.internet.InternetAddress;
import jakarta.mail.internet.MimeMessage;
import com.mailserver.security.LoginGuard;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.io.TempDir;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.Socket;
import java.nio.file.Path;
import java.security.cert.X509Certificate;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLSocket;
import javax.net.ssl.TrustManager;
import javax.net.ssl.X509TrustManager;
import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.ArrayList;
import java.util.Base64;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.CopyOnWriteArrayList;

import static org.assertj.core.api.Assertions.assertThat;

class SmtpServerTest {

    /** Mailboxes of example.test; alice can log in. */
    static class FakeBackend implements SmtpBackend {
        final Map<String, String> passwords = Map.of("alice@example.test", "secret");
        final List<String> mailboxes = List.of("alice@example.test", "bob@example.test");
        final List<Envelope> envelopes = new CopyOnWriteArrayList<>();
        final List<String> messages = new CopyOnWriteArrayList<>();
        volatile boolean failStore;

        @Override
        public String authenticate(String username, String password) {
            return password.equals(passwords.get(username)) ? username : null;
        }

        @Override
        public Route route(String recipient, String authenticatedUser) {
            if (recipient.endsWith("@example.test")) {
                return mailboxes.contains(recipient) ? Route.LOCAL : Route.UNKNOWN_USER;
            }
            return authenticatedUser == null ? Route.RELAY_DENIED : Route.OUTBOUND;
        }

        @Override
        public void accept(Envelope envelope, byte[] message) {
            if (failStore) {
                throw new IllegalStateException("database down");
            }
            envelopes.add(envelope);
            messages.add(new String(message, StandardCharsets.UTF_8));
        }
    }

    private FakeBackend backend;
    private SmtpServer server;

    @BeforeEach
    void start() {
        backend = new FakeBackend();
        server = new SmtpServer(new SmtpServer.Settings("127.0.0.1", 0, "mx.example.test", 2000, true,
                Duration.ofSeconds(30)), backend);
        server.start();
    }

    @AfterEach
    void stop() {
        server.stop();
    }

    @Test
    void storesMailFromARealClient() throws Exception {
        Properties props = new Properties();
        props.put("mail.smtp.host", "127.0.0.1");
        props.put("mail.smtp.port", String.valueOf(server.port()));
        Session session = Session.getInstance(props);
        MimeMessage message = new MimeMessage(session);
        message.setFrom(new InternetAddress("someone@elsewhere.test"));
        message.setRecipient(Message.RecipientType.TO, new InternetAddress("bob@example.test"));
        message.setSubject("Hello");
        message.setText("first line\n.starts with a dot\nlast line");

        Transport.send(message);

        assertThat(backend.envelopes).hasSize(1);
        SmtpBackend.Envelope envelope = backend.envelopes.get(0);
        assertThat(envelope.mailFrom()).isEqualTo("someone@elsewhere.test");
        assertThat(envelope.authenticatedUser()).isNull();
        assertThat(envelope.recipients()).containsExactly(new SmtpBackend.Recipient("bob@example.test", SmtpBackend.Route.LOCAL));
        assertThat(backend.messages.get(0))
                .startsWith("Received: from ")
                .contains("by mx.example.test with ESMTP;")
                .contains("Subject: Hello")
                .contains("\r\n.starts with a dot\r\n");
    }

    @Test
    void handlesPipelinedCommandsInOnePacketAndLinesSplitAcrossPackets() throws Exception {
        try (Client client = new Client(server.port())) {
            client.read();
            client.write("EHLO c.test\r\nMAIL FROM:<x@elsewhere.test>\r\nRCPT TO:<bob@example.test>\r\nDATA\r\n");
            assertThat(client.readReply()).startsWith("250-mx.example.test");
            assertThat(client.read()).startsWith("250 2.1.0");
            assertThat(client.read()).startsWith("250 2.1.5");
            assertThat(client.read()).startsWith("354");
            client.write("Subject: split\r\n\r\nhal");
            client.write("f a line\r\n..dot\r\n.");
            client.write("\r\nQUIT\r\n");
            assertThat(client.read()).startsWith("250 2.0.0");
            assertThat(client.read()).startsWith("221");
        }
        assertThat(backend.messages.get(0)).endsWith("Subject: split\r\n\r\nhalf a line\r\n.dot\r\n");
    }

    @Test
    void refusesToRelayForStrangers() throws Exception {
        try (Client client = new Client(server.port())) {
            client.read();
            client.send("EHLO spammer.test");
            client.readReply();
            client.send("MAIL FROM:<spam@spammer.test>");
            assertThat(client.send("RCPT TO:<victim@gmail.com>")).startsWith("550 5.7.1 Relaying denied");
            assertThat(client.send("RCPT TO:<nobody@example.test>")).startsWith("550 5.1.1");
            assertThat(client.send("DATA")).startsWith("554");
        }
        assertThat(backend.envelopes).isEmpty();
    }

    @Test
    void relaysForALoggedInUserOnlyFromHisOwnAddress() throws Exception {
        try (Client client = new Client(server.port())) {
            client.read();
            client.send("EHLO laptop");
            assertThat(client.readReplyLines()).contains("250-AUTH PLAIN LOGIN");
            assertThat(client.send("AUTH PLAIN " + b64("\0alice@example.test\0wrong"))).startsWith("535");
            assertThat(client.send("AUTH PLAIN " + b64("\0Alice@Example.test\0secret"))).startsWith("235");
            assertThat(client.send("MAIL FROM:<bob@example.test>")).startsWith("553");
            assertThat(client.send("MAIL FROM:<alice@example.test>")).startsWith("250");
            assertThat(client.send("RCPT TO:<friend@gmail.com>")).startsWith("250");
            assertThat(client.send("DATA")).startsWith("354");
            assertThat(client.send("Subject: out\r\n\r\nhi\r\n.")).startsWith("250");
        }
        SmtpBackend.Envelope envelope = backend.envelopes.get(0);
        assertThat(envelope.authenticatedUser()).isEqualTo("alice@example.test");
        assertThat(envelope.recipients()).containsExactly(new SmtpBackend.Recipient("friend@gmail.com", SmtpBackend.Route.OUTBOUND));
        assertThat(backend.messages.get(0)).contains("with ESMTPA;");
    }

    @Test
    void threeWrongPasswordsCloseTheConnectionAndTheGuardBlocksTheIp() throws Exception {
        server.stop();
        LoginGuard guard = new LoginGuard(new LoginGuard.Settings(true, 4, 100, Duration.ofMinutes(15),
                Duration.ofMinutes(15), Duration.ofHours(24), Duration.ofDays(30), java.util.Set.of(), Duration.ZERO));
        server = new SmtpServer(new SmtpServer.Settings("127.0.0.1", 0, "mx.example.test", 2000, true,
                Duration.ofSeconds(30)), backend, guard);
        server.start();
        try (Client client = new Client(server.port())) {
            client.read();
            client.send("EHLO laptop");
            client.readReplyLines();
            assertThat(client.send("AUTH PLAIN " + b64("\0alice@example.test\0w1"))).startsWith("535");
            assertThat(client.send("AUTH PLAIN " + b64("\0alice@example.test\0w2"))).startsWith("535");
            assertThat(client.send("AUTH PLAIN " + b64("\0alice@example.test\0w3"))).startsWith("535").contains("closing");
            assertThat(client.read()).isNull();
        }
        try (Client client = new Client(server.port())) {
            client.read();
            client.send("EHLO laptop");
            client.readReplyLines();
            assertThat(client.send("AUTH PLAIN " + b64("\0bob@example.test\0w4"))).startsWith("535");
            // 4th failure from 127.0.0.1: blocked — even the right password is not checked
            assertThat(client.send("AUTH PLAIN " + b64("\0alice@example.test\0secret"))).startsWith("421 4.7.0");
            assertThat(client.read()).isNull();
        }
    }

    @Test
    void authLoginWorksToo() throws Exception {
        try (Client client = new Client(server.port())) {
            client.read();
            client.send("EHLO laptop");
            client.readReply();
            assertThat(client.send("AUTH LOGIN")).isEqualTo("334 " + b64("Username:"));
            assertThat(client.send(b64("alice@example.test"))).isEqualTo("334 " + b64("Password:"));
            assertThat(client.send(b64("secret"))).startsWith("235");
        }
    }

    @Test
    void rejectsTooBigMessagesAndKeepsTheSession() throws Exception {
        try (Client client = new Client(server.port())) {
            client.read();
            client.send("HELO c");
            client.send("MAIL FROM:<x@elsewhere.test>");
            client.send("RCPT TO:<bob@example.test>");
            client.send("DATA");
            assertThat(client.send("x".repeat(3000) + "\r\n.")).startsWith("552");
            assertThat(client.send("MAIL FROM:<x@elsewhere.test> SIZE=5000")).startsWith("552");
            assertThat(client.send("NOOP")).startsWith("250");
        }
        assertThat(backend.envelopes).isEmpty();
    }

    @Test
    void answers451WhenTheMessageCannotBeStored() throws Exception {
        backend.failStore = true;
        try (Client client = new Client(server.port())) {
            client.read();
            client.send("HELO c");
            client.send("MAIL FROM:<x@elsewhere.test>");
            client.send("RCPT TO:<bob@example.test>");
            client.send("DATA");
            assertThat(client.send("Subject: x\r\n\r\nbody\r\n.")).startsWith("451");
        }
    }

    @Test
    void withoutTlsAuthIsOffWhenNotAllowed() throws Exception {
        server.stop();
        server = new SmtpServer(new SmtpServer.Settings("127.0.0.1", 0, "mx.example.test", 2000, false,
                Duration.ofSeconds(30)), backend);
        server.start();
        try (Client client = new Client(server.port())) {
            client.read();
            client.send("EHLO laptop");
            assertThat(client.readReplyLines()).noneMatch(line -> line.contains("AUTH"));
            assertThat(client.send("AUTH PLAIN " + b64("\0alice@example.test\0secret"))).startsWith("538");
        }
    }

    @Test
    void parsesPaths() {
        assertThat(SmtpServer.path("<A@B.test> SIZE=10")).isEqualTo("a@b.test");
        assertThat(SmtpServer.path("<>")).isEmpty();
        assertThat(SmtpServer.path("<@relay.test:u@b.test>")).isEqualTo("u@b.test");
        assertThat(SmtpServer.path("u@b.test")).isNull();
        assertThat(SmtpServer.path("<a b@c>")).isNull();
    }

    // ---------- TLS (75.1) ----------

    @TempDir
    Path tlsDir;

    private SmtpServer tlsServer(boolean implicit) throws Exception {
        server.stop();
        server = new SmtpServer(new SmtpServer.Settings("127.0.0.1", 0, "mx.example.test", 2000, false,
                Duration.ofSeconds(30), SmtpTls.serverContext(null, null, tlsDir, "mx.example.test"), implicit), backend);
        server.start();
        return server;
    }

    @Test
    void startTlsThenAuthAndTheSessionStartsOver() throws Exception {
        tlsServer(false);
        try (Client client = new Client(server.port())) {
            client.read();
            client.send("EHLO laptop");
            List<String> plain = client.readReplyLines();
            assertThat(plain).contains("250-STARTTLS").noneMatch(line -> line.contains("AUTH"));
            assertThat(client.send("AUTH PLAIN " + b64("\0alice@example.test\0secret"))).startsWith("538");
            assertThat(client.send("MAIL FROM:<x@elsewhere.test>")).startsWith("250");

            assertThat(client.send("STARTTLS")).isEqualTo("220 2.0.0 Ready to start TLS");
            client.startTls();

            // RFC 3207: MAIL FROM before STARTTLS is forgotten, EHLO again
            assertThat(client.send("MAIL FROM:<x@elsewhere.test>")).startsWith("503");
            client.send("EHLO laptop");
            List<String> secure = client.readReplyLines();
            assertThat(secure).contains("250-AUTH PLAIN LOGIN").noneMatch(line -> line.contains("STARTTLS"));
            assertThat(client.send("STARTTLS")).startsWith("503");
            assertThat(client.send("AUTH PLAIN " + b64("\0alice@example.test\0secret"))).startsWith("235");
            client.send("MAIL FROM:<alice@example.test>");
            client.send("RCPT TO:<friend@gmail.com>");
            client.send("DATA");
            assertThat(client.send("Subject: tls\r\n\r\nhi\r\n.")).startsWith("250");
        }
        assertThat(backend.messages.get(0)).contains("with ESMTPSA;");
    }

    // CVE-2011-0411: commands sent in plain text right after STARTTLS must not run inside the TLS session
    @Test
    void plainTextInjectedAfterStartTlsIsNotExecuted() throws Exception {
        tlsServer(false);
        try (Client client = new Client(server.port())) {
            client.read();
            client.send("EHLO laptop");
            client.readReplyLines();
            client.write("STARTTLS\r\nMAIL FROM:<x@elsewhere.test>\r\nRCPT TO:<bob@example.test>\r\n");
            assertThat(client.read()).isEqualTo("220 2.0.0 Ready to start TLS");
            client.startTls();
            String reply = client.read();
            assertThat(reply == null || reply.startsWith("421")).as(reply).isTrue();
        }
        assertThat(backend.envelopes).isEmpty();
    }

    @Test
    void realClientRequiringStartTlsDelivers() throws Exception {
        tlsServer(false);
        Properties props = new Properties();
        props.put("mail.smtp.host", "127.0.0.1");
        props.put("mail.smtp.port", String.valueOf(server.port()));
        props.put("mail.smtp.starttls.enable", "true");
        props.put("mail.smtp.starttls.required", "true");
        props.put("mail.smtp.ssl.trust", "*");
        props.put("mail.smtp.ssl.checkserveridentity", "false"); // self-signed for mx.example.test, connecting to 127.0.0.1
        props.put("mail.smtp.auth", "true");
        Session session = Session.getInstance(props);
        MimeMessage message = new MimeMessage(session);
        message.setFrom(new InternetAddress("alice@example.test"));
        message.setRecipient(Message.RecipientType.TO, new InternetAddress("bob@example.test"));
        message.setSubject("Over STARTTLS");
        message.setText("hello");

        Transport.send(message, "alice@example.test", "secret");

        assertThat(backend.envelopes).hasSize(1);
        assertThat(backend.envelopes.get(0).authenticatedUser()).isEqualTo("alice@example.test");
        assertThat(backend.messages.get(0)).contains("with ESMTPSA;");
    }

    @Test
    void implicitTlsPortWorksWithoutStartTls() throws Exception {
        tlsServer(true);
        Properties props = new Properties();
        props.put("mail.smtp.host", "127.0.0.1");
        props.put("mail.smtp.port", String.valueOf(server.port()));
        props.put("mail.smtp.ssl.enable", "true");
        props.put("mail.smtp.ssl.trust", "*");
        props.put("mail.smtp.ssl.checkserveridentity", "false"); // self-signed for mx.example.test, connecting to 127.0.0.1
        props.put("mail.smtp.auth", "true");
        Session session = Session.getInstance(props);
        MimeMessage message = new MimeMessage(session);
        message.setFrom(new InternetAddress("alice@example.test"));
        message.setRecipient(Message.RecipientType.TO, new InternetAddress("bob@example.test"));
        message.setSubject("Over 465");
        message.setText("hello");

        Transport.send(message, "alice@example.test", "secret");

        assertThat(backend.messages.get(0)).contains("with ESMTPSA;");
    }

    @Test
    void selfSignedCertificateIsCreatedOnceAndReused() throws Exception {
        SmtpTls.serverContext(null, null, tlsDir, "mx.example.test");
        String first = java.nio.file.Files.readString(tlsDir.resolve(SmtpTls.CERT_FILE));
        SmtpTls.serverContext(null, null, tlsDir, "mx.example.test");

        assertThat(java.nio.file.Files.readString(tlsDir.resolve(SmtpTls.CERT_FILE))).isEqualTo(first);
        assertThat(java.nio.file.Files.readString(tlsDir.resolve(SmtpTls.KEY_FILE))).startsWith("-----BEGIN PRIVATE KEY-----");
    }

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

    /** Line-oriented test client. */
    static final class Client implements AutoCloseable {
        private Socket socket;
        private BufferedReader in;
        private OutputStream out;

        Client(int port) throws Exception {
            socket = new Socket("127.0.0.1", port);
            socket.setSoTimeout(5000);
            in = new BufferedReader(new InputStreamReader(socket.getInputStream(), StandardCharsets.UTF_8));
            out = socket.getOutputStream();
        }

        /** Upgrade the connection to TLS (any certificate is trusted — the test server's is self-signed). */
        void startTls() throws Exception {
            SSLContext context = SSLContext.getInstance("TLS");
            context.init(null, new TrustManager[]{new X509TrustManager() {
                public void checkClientTrusted(X509Certificate[] chain, String authType) {
                }

                public void checkServerTrusted(X509Certificate[] chain, String authType) {
                }

                public X509Certificate[] getAcceptedIssuers() {
                    return new X509Certificate[0];
                }
            }}, null);
            SSLSocket ssl = (SSLSocket) context.getSocketFactory().createSocket(socket, "127.0.0.1", socket.getPort(), true);
            ssl.setUseClientMode(true);
            ssl.startHandshake();
            socket = ssl;
            in = new BufferedReader(new InputStreamReader(ssl.getInputStream(), StandardCharsets.UTF_8));
            out = ssl.getOutputStream();
        }

        void write(String text) throws Exception {
            out.write(text.getBytes(StandardCharsets.UTF_8));
            out.flush();
            Thread.sleep(50); // separate TCP packets
        }

        String send(String line) throws Exception {
            out.write((line + "\r\n").getBytes(StandardCharsets.UTF_8));
            out.flush();
            return read();
        }

        String read() throws Exception {
            return in.readLine();
        }

        /** First line of a multi-line reply; the rest is consumed. */
        String readReply() throws Exception {
            return readReplyLines().get(0);
        }

        List<String> readReplyLines() throws Exception {
            List<String> lines = new ArrayList<>();
            String line;
            do {
                line = in.readLine();
                lines.add(line);
            } while (line != null && line.length() > 3 && line.charAt(3) == '-');
            return lines;
        }

        @Override
        public void close() throws Exception {
            socket.close();
        }
    }
}