// SPDX-FileCopyrightText: © 2026 Jeffrey C. Ollie <jeff@ocjtech.us>
// SPDX-License-Identifier: MIT

//! A single-connection SMTP server session. Like the client, it runs over
//! any `Io.Reader`/`Io.Writer` pair; accept a TCP connection and hand its
//! stream reader/writer to `run`. Accepting connections, concurrency, and
//! message storage are left to the caller — the session just speaks the
//! protocol and forwards decisions to a `Handler`.
//!
//! Typical use:
//! ```
//! var session: Server = .init(&stream_reader, &stream_writer, handler, .{
//!     .hostname = "mx.example.com",
//! });
//! try session.run(gpa);
//! ```

const Server = @This();

const std = @import("std");
const Io = std.Io;
const tls = @import("tls");
const protocol = @import("protocol.zig");

reader: *Io.Reader,
writer: *Io.Writer,
handler: Handler,
options: Options,
/// True once a STARTTLS handshake has completed for this session.
secured: bool = false,
tls_connection: tls.Connection = undefined,
tls_reader: tls.Connection.Reader = undefined,
tls_writer: tls.Connection.Writer = undefined,
tls_read_buffer: [4096]u8 = undefined,
tls_write_buffer: [4096]u8 = undefined,

pub const Options = struct {
    /// Hostname announced in the greeting and the EHLO response.
    hostname: []const u8 = "localhost",
    /// Advertised via the SIZE extension and enforced during DATA.
    max_message_size: usize = 16 * 1024 * 1024,
    max_recipients: usize = 100,
    /// When set, STARTTLS is advertised and accepted. The underlying stream
    /// reader/writer handed to `init` must then have buffers of at least
    /// `tls.input_buffer_len` and `tls.output_buffer_len` bytes, since the
    /// handshake and TLS records run over them.
    starttls: ?StartTls = null,
    /// Reject MAIL with 530 until the client has authenticated. Requires a
    /// handler with an `authenticate` callback.
    require_auth: bool = false,
};

pub const StartTls = struct {
    io: Io,
    /// Server certificate chain and private key presented to clients.
    auth: *tls.config.CertKeyPair,
};

/// A handler's verdict on an envelope step or a complete message.
pub const Decision = union(enum) {
    accept,
    reject: Rejection,

    pub const Rejection = struct {
        /// Use 4xx for "try again later", 5xx for permanent rejection.
        code: u16 = 550,
        text: []const u8 = "5.7.1 Rejected",
    };
};

pub const Envelope = struct {
    /// Empty for the null reverse-path (`MAIL FROM:<>`).
    from: []const u8,
    recipients: []const []const u8,
};

/// Callbacks invoked during a session. All slices passed to callbacks are
/// only valid for the duration of the call.
pub const Handler = struct {
    context: ?*anyopaque = null,
    vtable: *const VTable,

    pub const VTable = struct {
        /// Called for AUTH with the decoded credentials; return true to
        /// accept. When set, AUTH PLAIN and AUTH LOGIN are advertised and
        /// accepted (RFC 4954).
        authenticate: ?*const fn (context: ?*anyopaque, username: []const u8, password: []const u8) bool = null,
        /// Called for MAIL FROM. Null accepts every sender.
        mailFrom: ?*const fn (context: ?*anyopaque, from: []const u8) Decision = null,
        /// Called for each RCPT TO. Null accepts every recipient.
        rcptTo: ?*const fn (context: ?*anyopaque, to: []const u8) Decision = null,
        /// Called once the complete message has been received. The data has
        /// CRLF line endings and dot-stuffing already removed. Exactly one
        /// of `message` and `messageReader` must be set.
        message: ?*const fn (context: ?*anyopaque, envelope: Envelope, data: []const u8) Decision = null,
        /// Streaming alternative to `message`: called after DATA with a
        /// reader that yields the message content (dot-stuffing removed,
        /// line endings normalized to CRLF) until end of stream. Anything
        /// the callback leaves unread is drained by the session, so
        /// returning early is fine. `Options.max_message_size` is not
        /// enforced in this mode; individual message lines must fit the
        /// session's stream reader buffer.
        messageReader: ?*const fn (context: ?*anyopaque, envelope: Envelope, message: *Io.Reader) Decision = null,
    };
};

pub fn init(reader: *Io.Reader, writer: *Io.Writer, handler: Handler, options: Options) Server {
    return .{ .reader = reader, .writer = writer, .handler = handler, .options = options };
}

pub const RunError = error{ WriteFailed, ReadFailed, OutOfMemory, TlsHandshakeFailed };

/// Serves the session until the client sends QUIT or disconnects. `gpa`
/// backs per-transaction storage (envelope and message data); everything is
/// freed on return.
pub fn run(s: *Server, gpa: std.mem.Allocator) RunError!void {
    var arena_state: std.heap.ArenaAllocator = .init(gpa);
    defer arena_state.deinit();
    const arena = arena_state.allocator();

    std.debug.assert(!s.options.require_auth or s.handler.vtable.authenticate != null);
    std.debug.assert((s.handler.vtable.message == null) != (s.handler.vtable.messageReader == null));

    var greeted = false;
    var authenticated = false;
    var from: ?[]const u8 = null;
    var recipients: std.ArrayList([]const u8) = .empty;

    try s.writer.print("220 {s} ESMTP ready" ++ protocol.crlf, .{s.options.hostname});
    try s.writer.flush();

    while (true) {
        const line = protocol.readLine(s.reader) catch |err| switch (err) {
            error.EndOfStream => return, // Client disconnected.
            error.ReadFailed => return error.ReadFailed,
            error.LineTooLong => {
                try s.discardLine();
                try s.reply(500, "5.5.2 Line too long");
                continue;
            },
        };
        const command = protocol.Command.parse(line) catch {
            try s.reply(501, "5.5.4 Syntax error in parameters");
            continue;
        };
        switch (command) {
            .helo => {
                greeted = true;
                from = null;
                recipients = .empty;
                _ = arena_state.reset(.retain_capacity);
                try s.reply(250, s.options.hostname);
            },
            .ehlo => {
                greeted = true;
                from = null;
                recipients = .empty;
                _ = arena_state.reset(.retain_capacity);
                try s.writer.print("250-{s}\r\n250-PIPELINING\r\n250-8BITMIME\r\n", .{s.options.hostname});
                if (s.options.starttls != null and !s.secured)
                    try s.writer.writeAll("250-STARTTLS\r\n");
                if (s.handler.vtable.authenticate != null and !authenticated)
                    try s.writer.writeAll("250-AUTH PLAIN LOGIN\r\n");
                try s.writer.print("250 SIZE {d}\r\n", .{s.options.max_message_size});
                try s.writer.flush();
            },
            .mail => |args| {
                if (!greeted) {
                    try s.reply(503, "5.5.1 Send EHLO first");
                    continue;
                }
                if (s.options.require_auth and !authenticated) {
                    try s.reply(530, "5.7.0 Authentication required");
                    continue;
                }
                if (from != null) {
                    try s.reply(503, "5.5.1 Nested MAIL command");
                    continue;
                }
                if (s.handler.vtable.mailFrom) |callback| {
                    switch (callback(s.handler.context, args.path)) {
                        .accept => {},
                        .reject => |r| {
                            try s.reply(r.code, r.text);
                            continue;
                        },
                    }
                }
                from = try arena.dupe(u8, args.path);
                try s.reply(250, "2.1.0 Ok");
            },
            .rcpt => |args| {
                if (from == null) {
                    try s.reply(503, "5.5.1 Need MAIL command first");
                    continue;
                }
                if (recipients.items.len >= s.options.max_recipients) {
                    try s.reply(452, "4.5.3 Too many recipients");
                    continue;
                }
                if (s.handler.vtable.rcptTo) |callback| {
                    switch (callback(s.handler.context, args.path)) {
                        .accept => {},
                        .reject => |r| {
                            try s.reply(r.code, r.text);
                            continue;
                        },
                    }
                }
                try recipients.append(arena, try arena.dupe(u8, args.path));
                try s.reply(250, "2.1.5 Ok");
            },
            .data => {
                if (recipients.items.len == 0) {
                    try s.reply(503, "5.5.1 Need RCPT command first");
                    continue;
                }
                try s.receiveData(arena, .{
                    .from = from.?,
                    .recipients = recipients.items,
                });
                from = null;
                recipients = .empty;
                _ = arena_state.reset(.retain_capacity);
            },
            .rset => {
                from = null;
                recipients = .empty;
                _ = arena_state.reset(.retain_capacity);
                try s.reply(250, "2.0.0 Ok");
            },
            .noop => try s.reply(250, "2.0.0 Ok"),
            .vrfy => try s.reply(252, "2.5.2 Cannot VRFY user"),
            .help => try s.reply(214, "2.0.0 See RFC 5321"),
            .starttls => {
                const config = s.options.starttls orelse {
                    try s.reply(502, "5.5.1 STARTTLS not supported");
                    continue;
                };
                if (s.secured) {
                    try s.reply(503, "5.5.1 TLS already active");
                    continue;
                }
                try s.reply(220, "2.0.0 Ready to start TLS");
                var rng_source: std.Random.IoSource = .{ .io = config.io };
                s.tls_connection = tls.server(s.reader, s.writer, .{
                    .auth = config.auth,
                    .rng = rng_source.interface(),
                    .now = Io.Clock.real.now(config.io),
                }) catch return error.TlsHandshakeFailed;
                s.tls_reader = s.tls_connection.reader(&s.tls_read_buffer);
                s.tls_writer = s.tls_connection.writer(&s.tls_write_buffer);
                s.reader = &s.tls_reader.interface;
                s.writer = &s.tls_writer.interface;
                s.secured = true;
                // RFC 3207 §4.2: both sides return to their initial state;
                // the client must EHLO again.
                greeted = false;
                authenticated = false;
                from = null;
                recipients = .empty;
                _ = arena_state.reset(.retain_capacity);
            },
            .quit => {
                try s.reply(221, "2.0.0 Bye");
                if (s.secured) s.tls_connection.close() catch {};
                return;
            },
            .auth => |args| {
                if (s.handler.vtable.authenticate == null) {
                    try s.reply(503, "5.5.1 Authentication not enabled");
                    continue;
                }
                if (!greeted) {
                    try s.reply(503, "5.5.1 Send EHLO first");
                    continue;
                }
                if (authenticated) {
                    try s.reply(503, "5.5.1 Already authenticated");
                    continue;
                }
                if (from != null) {
                    try s.reply(503, "5.5.1 MAIL transaction in progress");
                    continue;
                }
                switch (try s.receiveAuth(args)) {
                    .authenticated => authenticated = true,
                    .rejected => {},
                    .disconnected => return,
                }
            },
            .unknown => try s.reply(500, "5.5.2 Command not recognized"),
        }
    }
}

const AuthOutcome = enum { authenticated, rejected, disconnected };

/// Runs the challenge/response exchange for AUTH PLAIN or AUTH LOGIN
/// (RFC 4954) and consults the handler's `authenticate` callback. Every
/// outcome except `disconnected` has already sent its reply.
fn receiveAuth(s: *Server, args: protocol.Command.AuthArgs) RunError!AuthOutcome {
    const callback = s.handler.vtable.authenticate.?;

    if (std.ascii.eqlIgnoreCase(args.mechanism, "PLAIN")) {
        var decoded_buf: [576]u8 = undefined;
        var response: []const u8 = args.initial;
        if (response.len == 0) {
            try s.reply(334, "");
            response = switch (try s.takeAuthLine()) {
                .line => |line| line,
                .cancelled => return .rejected,
                .disconnected => return .disconnected,
            };
        }
        const decoded = decodeBase64(&decoded_buf, response) orelse {
            try s.reply(501, "5.5.2 Invalid base64");
            return .rejected;
        };
        // authzid NUL authcid NUL password; the authzid is ignored.
        const first_nul = std.mem.indexOfScalar(u8, decoded, 0) orelse {
            try s.reply(501, "5.5.2 Malformed PLAIN response");
            return .rejected;
        };
        const after_authzid = decoded[first_nul + 1 ..];
        const second_nul = std.mem.indexOfScalar(u8, after_authzid, 0) orelse {
            try s.reply(501, "5.5.2 Malformed PLAIN response");
            return .rejected;
        };
        return s.finishAuth(callback, after_authzid[0..second_nul], after_authzid[second_nul + 1 ..]);
    }

    if (std.ascii.eqlIgnoreCase(args.mechanism, "LOGIN")) {
        var user_buf: [192]u8 = undefined;
        var pass_buf: [192]u8 = undefined;

        var username: []const u8 = undefined;
        if (args.initial.len > 0) {
            // Some clients send the username as an initial response.
            username = decodeBase64(&user_buf, args.initial) orelse {
                try s.reply(501, "5.5.2 Invalid base64");
                return .rejected;
            };
        } else {
            try s.reply(334, "VXNlcm5hbWU6"); // base64("Username:")
            const line = switch (try s.takeAuthLine()) {
                .line => |line| line,
                .cancelled => return .rejected,
                .disconnected => return .disconnected,
            };
            username = decodeBase64(&user_buf, line) orelse {
                try s.reply(501, "5.5.2 Invalid base64");
                return .rejected;
            };
        }
        try s.reply(334, "UGFzc3dvcmQ6"); // base64("Password:")
        const line = switch (try s.takeAuthLine()) {
            .line => |line| line,
            .cancelled => return .rejected,
            .disconnected => return .disconnected,
        };
        const password = decodeBase64(&pass_buf, line) orelse {
            try s.reply(501, "5.5.2 Invalid base64");
            return .rejected;
        };
        return s.finishAuth(callback, username, password);
    }

    try s.reply(504, "5.5.4 Unrecognized authentication type");
    return .rejected;
}

fn finishAuth(
    s: *Server,
    callback: *const fn (?*anyopaque, []const u8, []const u8) bool,
    username: []const u8,
    password: []const u8,
) RunError!AuthOutcome {
    if (callback(s.handler.context, username, password)) {
        try s.reply(235, "2.7.0 Authentication successful");
        return .authenticated;
    }
    try s.reply(535, "5.7.8 Authentication credentials invalid");
    return .rejected;
}

const AuthLine = union(enum) { line: []u8, cancelled, disconnected };

/// Reads one continuation line of an AUTH exchange. `cancelled` covers both
/// an explicit "*" and an overlong line; its reply has already been sent.
fn takeAuthLine(s: *Server) RunError!AuthLine {
    const line = protocol.readLine(s.reader) catch |err| switch (err) {
        error.EndOfStream => return .disconnected,
        error.ReadFailed => return error.ReadFailed,
        error.LineTooLong => {
            try s.discardLine();
            try s.reply(501, "5.5.2 Response too long");
            return .cancelled;
        },
    };
    if (std.mem.eql(u8, line, "*")) {
        try s.reply(501, "5.7.0 Authentication cancelled");
        return .cancelled;
    }
    return .{ .line = line };
}

/// Decodes a base64 AUTH argument; "=" denotes an empty response.
fn decodeBase64(out: []u8, encoded: []const u8) ?[]u8 {
    if (std.mem.eql(u8, encoded, "=")) return out[0..0];
    const len = std.base64.standard.Decoder.calcSizeForSlice(encoded) catch return null;
    if (len > out.len) return null;
    std.base64.standard.Decoder.decode(out[0..len], encoded) catch return null;
    return out[0..len];
}

/// Reads message content after DATA up to the terminating ".\r\n",
/// un-stuffing dots, then asks the handler to accept or reject.
fn receiveData(s: *Server, arena: std.mem.Allocator, envelope: Envelope) RunError!void {
    try s.reply(354, "End data with <CR><LF>.<CR><LF>");

    if (s.handler.vtable.messageReader) |callback| {
        var buffer: [1024]u8 = undefined;
        var data_reader: DataReader = .{
            .session_reader = s.reader,
            .interface = .{
                .buffer = &buffer,
                .vtable = &.{ .stream = DataReader.stream },
                .seek = 0,
                .end = 0,
            },
        };
        const decision = callback(s.handler.context, envelope, &data_reader.interface);
        // Consume whatever the callback left unread, up to and including
        // the terminating ".".
        while (!data_reader.finished) {
            const line = protocol.readLine(s.reader) catch |err| switch (err) {
                error.EndOfStream => return, // Client disconnected mid-message.
                error.ReadFailed => return error.ReadFailed,
                error.LineTooLong => {
                    try s.discardLine();
                    continue;
                },
            };
            if (std.mem.eql(u8, line, ".")) break;
        }
        switch (decision) {
            .accept => try s.reply(250, "2.0.0 Ok, message accepted"),
            .reject => |r| try s.reply(r.code, r.text),
        }
        return;
    }

    var data: std.ArrayList(u8) = .empty;
    var oversize = false;
    while (true) {
        const line = protocol.readLine(s.reader) catch |err| switch (err) {
            error.EndOfStream => return, // Client disconnected mid-message.
            error.ReadFailed => return error.ReadFailed,
            error.LineTooLong => {
                // Longer than our reader buffer; RFC 5321 caps text lines at
                // 1000 octets, so treat it as oversize but keep scanning for
                // the terminator.
                try s.discardLine();
                oversize = true;
                continue;
            },
        };
        if (std.mem.eql(u8, line, ".")) break;
        const content = if (line.len > 0 and line[0] == '.') line[1..] else line;
        if (oversize) continue;
        if (data.items.len + content.len + protocol.crlf.len > s.options.max_message_size) {
            oversize = true;
            continue;
        }
        try data.appendSlice(arena, content);
        try data.appendSlice(arena, protocol.crlf);
    }
    if (oversize) {
        try s.reply(552, "5.3.4 Message exceeds maximum size");
        return;
    }
    switch (s.handler.vtable.message.?(s.handler.context, envelope, data.items)) {
        .accept => try s.reply(250, "2.0.0 Ok, message accepted"),
        .reject => |r| try s.reply(r.code, r.text),
    }
}

/// Adapts the session's line-based DATA phase into an `Io.Reader` of the
/// unstuffed message content for `Handler.VTable.messageReader`.
const DataReader = struct {
    session_reader: *Io.Reader,
    interface: Io.Reader,
    /// Unread remainder of the current line (points into the session
    /// reader's buffer, which only this reader touches during DATA).
    line: []const u8 = &.{},
    line_ending: []const u8 = &.{},
    finished: bool = false,

    fn stream(io_r: *Io.Reader, w: *Io.Writer, limit: Io.Limit) Io.Reader.StreamError!usize {
        const dr: *DataReader = @alignCast(@fieldParentPtr("interface", io_r));
        if (dr.line.len == 0 and dr.line_ending.len == 0) {
            if (dr.finished) return error.EndOfStream;
            const raw = protocol.readLine(dr.session_reader) catch return error.ReadFailed;
            if (std.mem.eql(u8, raw, ".")) {
                dr.finished = true;
                return error.EndOfStream;
            }
            dr.line = if (raw.len > 0 and raw[0] == '.') raw[1..] else raw;
            dr.line_ending = protocol.crlf;
        }
        const dest = limit.slice(try w.writableSliceGreedy(1));
        const line_n = @min(dest.len, dr.line.len);
        @memcpy(dest[0..line_n], dr.line[0..line_n]);
        dr.line = dr.line[line_n..];
        var n = line_n;
        if (dr.line.len == 0) {
            const ending_n = @min(dest.len - n, dr.line_ending.len);
            @memcpy(dest[n..][0..ending_n], dr.line_ending[0..ending_n]);
            dr.line_ending = dr.line_ending[ending_n..];
            n += ending_n;
        }
        w.advance(n);
        return n;
    }
};

fn reply(s: *Server, code: u16, text: []const u8) error{WriteFailed}!void {
    try s.writer.print("{d} {s}" ++ protocol.crlf, .{ code, text });
    try s.writer.flush();
}

/// Discards input through the next newline after `error.LineTooLong`, which
/// leaves the reader positioned at the start of the oversized line.
fn discardLine(s: *Server) error{ReadFailed}!void {
    _ = s.reader.discardDelimiterInclusive('\n') catch |err| switch (err) {
        error.EndOfStream => {},
        error.ReadFailed => return error.ReadFailed,
    };
}

const TestHandler = struct {
    from: std.ArrayList(u8) = .empty,
    recipients: std.ArrayList(u8) = .empty,
    data: std.ArrayList(u8) = .empty,
    messages_accepted: usize = 0,
    reject_recipient: ?[]const u8 = null,
    /// When set, enables the authenticate callback accepting user "alice"
    /// with this password.
    password: ?[]const u8 = null,

    fn deinit(h: *TestHandler) void {
        h.from.deinit(std.testing.allocator);
        h.recipients.deinit(std.testing.allocator);
        h.data.deinit(std.testing.allocator);
    }

    fn handler(h: *TestHandler) Handler {
        return .{ .context = h, .vtable = if (h.password != null) &.{
            .authenticate = onAuthenticate,
            .rcptTo = onRcptTo,
            .message = onMessage,
        } else &.{
            .rcptTo = onRcptTo,
            .message = onMessage,
        } };
    }

    fn onAuthenticate(context: ?*anyopaque, username: []const u8, password: []const u8) bool {
        const h: *TestHandler = @ptrCast(@alignCast(context.?));
        return std.mem.eql(u8, username, "alice") and
            std.mem.eql(u8, password, h.password.?);
    }

    fn onRcptTo(context: ?*anyopaque, to: []const u8) Decision {
        const h: *TestHandler = @ptrCast(@alignCast(context.?));
        if (h.reject_recipient) |rejected| {
            if (std.mem.eql(u8, to, rejected)) return .{ .reject = .{
                .code = 550,
                .text = "5.1.1 No such user",
            } };
        }
        return .accept;
    }

    fn onMessage(context: ?*anyopaque, envelope: Envelope, data: []const u8) Decision {
        const h: *TestHandler = @ptrCast(@alignCast(context.?));
        const gpa = std.testing.allocator;
        h.from.appendSlice(gpa, envelope.from) catch return .{ .reject = .{} };
        for (envelope.recipients) |recipient| {
            h.recipients.appendSlice(gpa, recipient) catch return .{ .reject = .{} };
            h.recipients.append(gpa, ';') catch return .{ .reject = .{} };
        }
        h.data.appendSlice(gpa, data) catch return .{ .reject = .{} };
        h.messages_accepted += 1;
        return .accept;
    }
};

fn runScript(input: []const u8, out_buf: []u8, handler: Handler, options: Options) ![]const u8 {
    var reader: Io.Reader = .fixed(input);
    var writer: Io.Writer = .fixed(out_buf);
    var session: Server = .init(&reader, &writer, handler, options);
    try session.run(std.testing.allocator);
    return writer.buffered();
}

test run {
    var h: TestHandler = .{};
    defer h.deinit();

    var reader: Io.Reader = .fixed("EHLO client.example.org\r\n" ++
        "MAIL FROM:<alice@example.com>\r\n" ++
        "RCPT TO:<bob@example.net>\r\n" ++
        "RCPT TO:<carol@example.net>\r\n" ++
        "DATA\r\n" ++
        "Subject: hi\r\n" ++
        "\r\n" ++
        "..stuffed line\r\n" ++
        "body\r\n" ++
        ".\r\n" ++
        "QUIT\r\n");
    var out_buf: [1024]u8 = undefined;
    var writer: Io.Writer = .fixed(&out_buf);

    var session: Server = .init(&reader, &writer, h.handler(), .{ .hostname = "mx.test" });
    try session.run(std.testing.allocator);
    const output = writer.buffered();

    try std.testing.expectEqualStrings("alice@example.com", h.from.items);
    try std.testing.expectEqualStrings("bob@example.net;carol@example.net;", h.recipients.items);
    try std.testing.expectEqualStrings("Subject: hi\r\n\r\n.stuffed line\r\nbody\r\n", h.data.items);
    try std.testing.expectEqual(@as(usize, 1), h.messages_accepted);

    try std.testing.expectEqualStrings(
        "220 mx.test ESMTP ready\r\n" ++
            "250-mx.test\r\n250-PIPELINING\r\n250-8BITMIME\r\n250 SIZE 16777216\r\n" ++
            "250 2.1.0 Ok\r\n" ++
            "250 2.1.5 Ok\r\n" ++
            "250 2.1.5 Ok\r\n" ++
            "354 End data with <CR><LF>.<CR><LF>\r\n" ++
            "250 2.0.0 Ok, message accepted\r\n" ++
            "221 2.0.0 Bye\r\n",
        output,
    );
}

test "command sequencing is enforced" {
    var h: TestHandler = .{};
    defer h.deinit();

    var out_buf: [1024]u8 = undefined;
    const output = try runScript(
        "MAIL FROM:<early@example.com>\r\n" ++
            "EHLO client.example.org\r\n" ++
            "RCPT TO:<bob@example.net>\r\n" ++
            "DATA\r\n" ++
            "QUIT\r\n",
        &out_buf,
        h.handler(),
        .{},
    );

    try std.testing.expectEqual(@as(usize, 0), h.messages_accepted);
    try std.testing.expect(std.mem.indexOf(u8, output, "503 5.5.1 Send EHLO first") != null);
    try std.testing.expect(std.mem.indexOf(u8, output, "503 5.5.1 Need MAIL command first") != null);
    try std.testing.expect(std.mem.indexOf(u8, output, "503 5.5.1 Need RCPT command first") != null);
}

test "handler can reject a recipient" {
    var h: TestHandler = .{ .reject_recipient = "nobody@example.net" };
    defer h.deinit();

    var out_buf: [1024]u8 = undefined;
    const output = try runScript(
        "EHLO client.example.org\r\n" ++
            "MAIL FROM:<alice@example.com>\r\n" ++
            "RCPT TO:<nobody@example.net>\r\n" ++
            "RCPT TO:<bob@example.net>\r\n" ++
            "DATA\r\n" ++
            "hello\r\n" ++
            ".\r\n" ++
            "QUIT\r\n",
        &out_buf,
        h.handler(),
        .{},
    );

    try std.testing.expect(std.mem.indexOf(u8, output, "550 5.1.1 No such user") != null);
    try std.testing.expectEqualStrings("bob@example.net;", h.recipients.items);
    try std.testing.expectEqual(@as(usize, 1), h.messages_accepted);
}

test "AUTH PLAIN with initial response" {
    var h: TestHandler = .{ .password = "secret" };
    defer h.deinit();

    var out_buf: [1024]u8 = undefined;
    // base64("\x00alice\x00secret")
    const output = try runScript(
        "EHLO client.example.org\r\n" ++
            "AUTH PLAIN AGFsaWNlAHNlY3JldA==\r\n" ++
            "MAIL FROM:<alice@example.com>\r\n" ++
            "RCPT TO:<bob@example.net>\r\n" ++
            "DATA\r\nauthed mail\r\n.\r\n" ++
            "QUIT\r\n",
        &out_buf,
        h.handler(),
        .{ .require_auth = true },
    );

    try std.testing.expect(std.mem.indexOf(u8, output, "250-AUTH PLAIN LOGIN\r\n") != null);
    try std.testing.expect(std.mem.indexOf(u8, output, "235 2.7.0") != null);
    try std.testing.expectEqual(@as(usize, 1), h.messages_accepted);
}

test "AUTH LOGIN challenge exchange" {
    var h: TestHandler = .{ .password = "secret" };
    defer h.deinit();

    var out_buf: [1024]u8 = undefined;
    // base64("alice"), base64("secret")
    const output = try runScript(
        "EHLO client.example.org\r\n" ++
            "AUTH LOGIN\r\n" ++
            "YWxpY2U=\r\n" ++
            "c2VjcmV0\r\n" ++
            "QUIT\r\n",
        &out_buf,
        h.handler(),
        .{},
    );

    try std.testing.expect(std.mem.indexOf(u8, output, "334 VXNlcm5hbWU6\r\n") != null);
    try std.testing.expect(std.mem.indexOf(u8, output, "334 UGFzc3dvcmQ6\r\n") != null);
    try std.testing.expect(std.mem.indexOf(u8, output, "235 2.7.0") != null);
}

test "AUTH failures and sequencing" {
    var h: TestHandler = .{ .password = "secret" };
    defer h.deinit();

    var out_buf: [2048]u8 = undefined;
    const output = try runScript(
        "EHLO client.example.org\r\n" ++
            "MAIL FROM:<alice@example.com>\r\n" ++ // before auth: 530
            "AUTH PLAIN AGFsaWNlAHdyb25n\r\n" ++ // wrong password: 535
            "AUTH GSSAPI\r\n" ++ // unsupported: 504
            "AUTH PLAIN not!base64\r\n" ++ // 501
            "AUTH LOGIN\r\n" ++
            "*\r\n" ++ // cancelled: 501
            "AUTH PLAIN AGFsaWNlAHNlY3JldA==\r\n" ++ // correct: 235
            "AUTH PLAIN AGFsaWNlAHNlY3JldA==\r\n" ++ // again: 503
            "QUIT\r\n",
        &out_buf,
        h.handler(),
        .{ .require_auth = true },
    );

    try std.testing.expect(std.mem.indexOf(u8, output, "530 5.7.0") != null);
    try std.testing.expect(std.mem.indexOf(u8, output, "535 5.7.8") != null);
    try std.testing.expect(std.mem.indexOf(u8, output, "504 5.5.4") != null);
    try std.testing.expect(std.mem.indexOf(u8, output, "501 5.5.2 Invalid base64") != null);
    try std.testing.expect(std.mem.indexOf(u8, output, "501 5.7.0 Authentication cancelled") != null);
    try std.testing.expect(std.mem.indexOf(u8, output, "235 2.7.0") != null);
    try std.testing.expect(std.mem.indexOf(u8, output, "503 5.5.1 Already authenticated") != null);
}

test "AUTH without a handler is refused" {
    var h: TestHandler = .{};
    defer h.deinit();

    var out_buf: [1024]u8 = undefined;
    const output = try runScript(
        "EHLO client.example.org\r\nAUTH PLAIN AGEAYg==\r\nQUIT\r\n",
        &out_buf,
        h.handler(),
        .{},
    );

    try std.testing.expect(std.mem.indexOf(u8, output, "250-AUTH") == null);
    try std.testing.expect(std.mem.indexOf(u8, output, "503 5.5.1 Authentication not enabled") != null);
}

test "oversize message is rejected but session continues" {
    var h: TestHandler = .{};
    defer h.deinit();

    var out_buf: [1024]u8 = undefined;
    const output = try runScript(
        "EHLO client.example.org\r\n" ++
            "MAIL FROM:<alice@example.com>\r\n" ++
            "RCPT TO:<bob@example.net>\r\n" ++
            "DATA\r\n" ++
            "0123456789012345678901234567890123456789\r\n" ++
            ".\r\n" ++
            "NOOP\r\n" ++
            "QUIT\r\n",
        &out_buf,
        h.handler(),
        .{ .max_message_size = 16 },
    );

    try std.testing.expectEqual(@as(usize, 0), h.messages_accepted);
    try std.testing.expect(std.mem.indexOf(u8, output, "552 5.3.4") != null);
    try std.testing.expect(std.mem.indexOf(u8, output, "250 2.0.0 Ok\r\n221") != null);
}

const StreamTestHandler = struct {
    collected: std.ArrayList(u8) = .empty,
    take_only: ?usize = null,

    fn handler(h: *StreamTestHandler) Handler {
        return .{ .context = h, .vtable = &.{
            .messageReader = onMessageReader,
        } };
    }

    fn onMessageReader(context: ?*anyopaque, envelope: Envelope, message: *Io.Reader) Decision {
        const h: *StreamTestHandler = @ptrCast(@alignCast(context.?));
        _ = envelope;
        const gpa = std.testing.allocator;
        if (h.take_only) |n| {
            const bytes = message.take(n) catch return .{ .reject = .{} };
            h.collected.appendSlice(gpa, bytes) catch return .{ .reject = .{} };
            return .accept;
        }
        message.appendRemaining(gpa, &h.collected, .unlimited) catch return .{ .reject = .{} };
        return .accept;
    }
};

test "streaming message handler receives unstuffed content" {
    var h: StreamTestHandler = .{};
    defer h.collected.deinit(std.testing.allocator);

    var out_buf: [1024]u8 = undefined;
    const output = try runScript(
        "EHLO client.example.org\r\n" ++
            "MAIL FROM:<alice@example.com>\r\n" ++
            "RCPT TO:<bob@example.net>\r\n" ++
            "DATA\r\n" ++
            "Subject: streamed\r\n" ++
            "\r\n" ++
            "..dot line\r\n" ++
            "body\r\n" ++
            ".\r\n" ++
            "QUIT\r\n",
        &out_buf,
        h.handler(),
        .{},
    );

    try std.testing.expectEqualStrings(
        "Subject: streamed\r\n\r\n.dot line\r\nbody\r\n",
        h.collected.items,
    );
    try std.testing.expect(std.mem.indexOf(u8, output, "250 2.0.0 Ok, message accepted") != null);
}

test "session drains what a streaming handler leaves unread" {
    var h: StreamTestHandler = .{ .take_only = 7 };
    defer h.collected.deinit(std.testing.allocator);

    var out_buf: [1024]u8 = undefined;
    const output = try runScript(
        "EHLO client.example.org\r\n" ++
            "MAIL FROM:<alice@example.com>\r\n" ++
            "RCPT TO:<bob@example.net>\r\n" ++
            "DATA\r\n" ++
            "Subject: mostly unread\r\n" ++
            "lots of body\r\n" ++
            ".\r\n" ++
            "NOOP\r\n" ++
            "QUIT\r\n",
        &out_buf,
        h.handler(),
        .{},
    );

    try std.testing.expectEqualStrings("Subject", h.collected.items);
    try std.testing.expect(std.mem.indexOf(u8, output, "250 2.0.0 Ok, message accepted") != null);
    // The NOOP after DATA proves the terminator was consumed.
    try std.testing.expect(std.mem.indexOf(u8, output, "250 2.0.0 Ok\r\n221") != null);
}

test "fuzz session with arbitrary client input" {
    try std.testing.fuzz({}, fuzzSession, .{});
}

fn fuzzSession(context: void, smith: *std.testing.Smith) !void {
    _ = context;
    var input_buf: [2048]u8 = undefined;
    const input = input_buf[0..smith.value(u11)];
    smith.bytes(input);

    var h: TestHandler = .{ .password = "secret" };
    defer h.deinit();

    var reader: Io.Reader = .fixed(input);
    var discarding: Io.Writer.Discarding = .init(&.{});
    var session: Server = .init(&reader, &discarding.writer, h.handler(), .{
        .max_message_size = 512,
        .max_recipients = 4,
    });
    // Whatever the "client" sends, the session must fail cleanly, never crash.
    session.run(std.testing.allocator) catch {};
}

test "fuzz collecting and streaming DATA agree" {
    try std.testing.fuzz({}, fuzzDataEquivalence, .{});
}

fn fuzzDataEquivalence(context: void, smith: *std.testing.Smith) !void {
    _ = context;
    var body_buf: [1024]u8 = undefined;
    const body = body_buf[0..smith.value(u10)];
    smith.bytes(body);

    var script_buf: [1200]u8 = undefined;
    const script = std.fmt.bufPrint(
        &script_buf,
        "EHLO fuzz.example.org\r\n" ++
            "MAIL FROM:<a@example.com>\r\n" ++
            "RCPT TO:<b@example.net>\r\n" ++
            "DATA\r\n{s}\r\n.\r\nQUIT\r\n",
        .{body},
    ) catch unreachable;

    var collecting: TestHandler = .{};
    defer collecting.deinit();
    var out_buf: [4096]u8 = undefined;
    _ = runScript(script, &out_buf, collecting.handler(), .{}) catch {};

    var streaming: StreamTestHandler = .{};
    defer streaming.collected.deinit(std.testing.allocator);
    _ = runScript(script, &out_buf, streaming.handler(), .{}) catch {};

    try std.testing.expectEqualSlices(u8, collecting.data.items, streaming.collected.items);
}