fediglam/src/http/server/response.zig

const std = @import("std");
const util = @import("util");
const http = @import("../lib.zig");

const Status = http.Status;
const Fields = http.Fields;

const chunk_size = 16 * 1024;
pub fn open(
    alloc: std.mem.Allocator,
    writer: anytype,
    headers: *const Fields,
    status: Status,
) !ResponseStream(@TypeOf(writer)) {
    const buf = try alloc.alloc(u8, chunk_size);
    errdefer alloc.free(buf);

    try writeStatusLine(writer, status);
    try writeFields(writer, headers);

    return ResponseStream(@TypeOf(writer)){
        .allocator = alloc,
        .base_writer = writer,
        .headers = headers,
        .buffer = buf,
    };
}

pub fn writeRequestHeader(writer: anytype, headers: *const Fields, status: Status) !void {
    try writeStatusLine(writer, status);
    try writeFields(writer, headers);
    try writer.writeAll("\r\n");
}

fn writeStatusLine(writer: anytype, status: Status) !void {
    const status_text = status.phrase() orelse "";
    try writer.print("HTTP/1.1 {} {s}\r\n", .{ @enumToInt(status), status_text });
}

fn writeFields(writer: anytype, headers: *const Fields) !void {
    var iter = headers.iterator();
    while (iter.next()) |header| {
        if (std.ascii.eqlIgnoreCase("Set-Cookie", header.key_ptr.*)) continue;

        try writer.print("{s}: {s}\r\n", .{ header.key_ptr.*, percentEncode(header.value_ptr.*) });
    }

    var cookie_iter = headers.getList("Set-Cookie");
    while (cookie_iter.next()) |cookie| {
        try writer.print("Set-Cookie: {s}\r\n", .{percentEncode(cookie)});
    }
}

const PercentEncode = struct {
    str: []const u8,

    pub fn format(v: PercentEncode, comptime _: []const u8, _: std.fmt.FormatOptions, writer: anytype) !void {
        for (v.str) |ch| switch (ch) {
            ' ', '\t', 0x21...0x7e, 0x80...0xff => try writer.writeByte(ch),
            else => try std.fmt.format(writer, "%{x:0>2}", .{ch}),
        };
    }
};
fn percentEncode(str: []const u8) PercentEncode {
    return PercentEncode{ .str = str };
}

fn writeChunk(writer: anytype, contents: []const u8) @TypeOf(writer).Error!void {
    try writer.print("{x}\r\n", .{contents.len});
    try writer.writeAll(contents);
    try writer.writeAll("\r\n");
}

pub fn ResponseStream(comptime BaseWriter: type) type {
    return struct {
        const Self = @This();
        const Error = BaseWriter.Error;
        const Writer = std.io.Writer(*Self, Error, write);

        allocator: std.mem.Allocator,
        base_writer: BaseWriter,
        headers: *const Fields,
        buffer: []u8,
        buffer_pos: usize = 0,
        chunked: bool = false,

        fn writeToBuffer(self: *Self, bytes: []const u8) void {
            std.mem.copy(u8, self.buffer[self.buffer_pos..], bytes);
            self.buffer_pos += bytes.len;
        }

        fn startChunking(self: *Self) Error!void {
            try self.base_writer.writeAll("Transfer-Encoding: chunked\r\n\r\n");
            self.chunked = true;
        }

        fn flushChunk(self: *Self) Error!void {
            try writeChunk(self.base_writer, self.buffer[0..self.buffer_pos]);
            self.buffer_pos = 0;
        }

        fn writeToChunks(self: *Self, bytes: []const u8) Error!void {
            var cursor: usize = 0;
            while (true) {
                const remaining_in_chunk = self.buffer.len - self.buffer_pos;
                const remaining_to_write = bytes.len - cursor;
                if (remaining_to_write <= remaining_in_chunk) {
                    self.writeToBuffer(bytes[cursor..]);
                    return;
                }

                self.writeToBuffer(bytes[cursor .. cursor + remaining_in_chunk]);
                cursor += remaining_in_chunk;
                try self.flushChunk();
            }
        }

        fn write(self: *Self, bytes: []const u8) Error!usize {
            if (!self.chunked and bytes.len > self.buffer.len - self.buffer_pos) {
                try self.startChunking();
            }

            if (self.chunked) {
                try self.writeToChunks(bytes);
            } else {
                self.writeToBuffer(bytes);
            }

            return bytes.len;
        }

        fn flushBodyUnchunked(self: *Self) Error!void {
            try self.base_writer.print("Content-Length: {}\r\n", .{self.buffer_pos});

            try self.base_writer.writeAll("\r\n");

            if (self.buffer_pos != 0) {
                try self.base_writer.writeAll(self.buffer[0..self.buffer_pos]);
            }
            self.buffer_pos = 0;
        }

        pub fn writer(self: *Self) Writer {
            return Writer{ .context = self };
        }

        pub fn finish(self: *Self) Error!void {
            if (!self.chunked) {
                try self.flushBodyUnchunked();
            } else {
                if (self.buffer_pos != 0) {
                    try self.flushChunk();
                }
                try self.base_writer.writeAll("0\r\n\r\n");
            }
        }

        pub fn close(self: *Self) void {
            self.allocator.free(self.buffer);
        }
    };
}

test {
    _ = _tests;
}
const _tests = struct {
    const toCrlf = util.comptimeToCrlf;

    const test_buffer_size = chunk_size * 4;
    test "ResponseStream no headers empty body" {
        var buffer: [test_buffer_size]u8 = undefined;
        var test_stream = std.io.fixedBufferStream(&buffer);
        var headers = Fields.init(std.testing.allocator);
        defer headers.deinit();

        {
            var stream = try open(
                std.testing.allocator,
                test_stream.writer(),
                &headers,
                .ok,
            );
            defer stream.close();

            try stream.finish();
        }

        try std.testing.expectEqualStrings(
            toCrlf(
                \\HTTP/1.1 200 OK
                \\
                \\
            ),
            buffer[0..(try test_stream.getPos())],
        );
    }

    test "ResponseStream empty body" {
        var buffer: [test_buffer_size]u8 = undefined;
        var test_stream = std.io.fixedBufferStream(&buffer);
        var headers = Fields.init(std.testing.allocator);
        defer headers.deinit();

        try headers.put("Content-Type", "text/plain");

        {
            var stream = try open(
                std.testing.allocator,
                test_stream.writer(),
                &headers,
                .ok,
            );
            defer stream.close();

            try stream.finish();
        }

        try std.testing.expectEqualStrings(
            toCrlf(
                \\HTTP/1.1 200 OK
                \\Content-Type: text/plain
                \\
                \\
            ),
            buffer[0..(try test_stream.getPos())],
        );
    }

    test "ResponseStream not 200 OK" {
        var buffer: [test_buffer_size]u8 = undefined;
        var test_stream = std.io.fixedBufferStream(&buffer);
        var headers = Fields.init(std.testing.allocator);
        defer headers.deinit();

        try headers.put("Content-Type", "text/plain");

        {
            var stream = try open(
                std.testing.allocator,
                test_stream.writer(),
                &headers,
                .not_found,
            );
            defer stream.close();

            try stream.finish();
        }

        try std.testing.expectEqualStrings(
            toCrlf(
                \\HTTP/1.1 404 Not Found
                \\Content-Type: text/plain
                \\
                \\
            ),
            buffer[0..(try test_stream.getPos())],
        );
    }
    test "ResponseStream small body" {
        var buffer: [test_buffer_size]u8 = undefined;
        var test_stream = std.io.fixedBufferStream(&buffer);
        var headers = Fields.init(std.testing.allocator);
        defer headers.deinit();

        try headers.put("Content-Type", "text/plain");

        {
            var stream = try open(
                std.testing.allocator,
                test_stream.writer(),
                &headers,
                .ok,
            );
            defer stream.close();

            try stream.writer().writeAll("Index Page");

            try stream.finish();
        }

        try std.testing.expectEqualStrings(
            toCrlf(
                \\HTTP/1.1 200 OK
                \\Content-Type: text/plain
                \\Content-Length: 10
                \\
                \\Index Page
            ),
            buffer[0..(try test_stream.getPos())],
        );
    }

    test "ResponseStream large body" {
        var buffer: [test_buffer_size]u8 = undefined;
        var test_stream = std.io.fixedBufferStream(&buffer);
        var headers = Fields.init(std.testing.allocator);
        defer headers.deinit();

        try headers.put("Content-Type", "text/plain");

        {
            var stream = try open(
                std.testing.allocator,
                test_stream.writer(),
                &headers,
                .ok,
            );
            defer stream.close();

            try stream.writer().writeAll("quuz" ** 6000);

            try stream.finish();
        }

        // zig fmt: off
        try std.testing.expectEqualStrings(
            toCrlf(
                \\HTTP/1.1 200 OK
                \\Content-Type: text/plain
                \\Transfer-Encoding: chunked
                \\
                \\
                ++ "4000\n"
                ++ "quuz" ** 4096 ++ "\n"
                ++ "1dc0\n"
                ++ "quuz" ** (1904) ++ "\n"
                ++ "0\n"
            ),
            buffer[0..(try test_stream.getPos())],
        );
        // zig fmt: on
    }

    test "ResponseStream large body ending on chunk boundary" {
        var buffer: [test_buffer_size]u8 = undefined;
        var test_stream = std.io.fixedBufferStream(&buffer);
        var headers = Fields.init(std.testing.allocator);
        defer headers.deinit();

        try headers.put("Content-Type", "text/plain");

        {
            var stream = try open(
                std.testing.allocator,
                test_stream.writer(),
                &headers,
                .ok,
            );
            defer stream.close();

            try stream.writer().writeAll("quuz" ** (chunk_size / 2));

            try stream.finish();
        }

        // zig fmt: off
        try std.testing.expectEqualStrings(
            toCrlf(
                \\HTTP/1.1 200 OK
                \\Content-Type: text/plain
                \\Transfer-Encoding: chunked
                \\
                \\
                ++ "4000\n"
                ++ "quuz" ** 4096 ++ "\n"
                ++ "4000\n"
                ++ "quuz" ** 4096 ++ "\n"
                ++ "0\n"
            ),
            buffer[0..(try test_stream.getPos())],
        );
        // zig fmt: on
    }
};