// 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,
};

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 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.
        message: *const fn (context: ?*anyopaque, envelope: Envelope, data: []const u8) Decision,
    };
};

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

    var greeted = 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");
                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 (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;
                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;
            },
            .unknown => try s.reply(500, "5.5.2 Command not recognized"),
        }
    }
}

/// 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>");
    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),
    }
}

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,

    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 = &.{
            .rcptTo = onRcptTo,
            .message = onMessage,
        } };
    }

    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 "complete session" {
    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" ++
            "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",
        &out_buf,
        h.handler(),
        .{ .hostname = "mx.test" },
    );

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