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 org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.io.OutputStream;
import java.net.Socket;
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 boolean authenticate(String username, String password) {
            return password.equals(passwords.get(username));
        }

        @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 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();
    }

    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 final Socket socket;
        private final BufferedReader in;
        private final 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();
        }

        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();
        }
    }
}