fediglam/src/http/multipart.zig

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

const max_boundary = 70;
const read_ahead = max_boundary + 4;

pub fn MultipartStream(comptime ReaderType: type) type {
    return struct {
        const Multipart = @This();

        pub const PartReader = std.io.Reader(*Part, ReaderType.Error, Part.read);

        stream: std.io.PeekStream(.{ .Static = read_ahead }, ReaderType),
        boundary: []const u8,

        pub fn next(self: *Multipart, alloc: std.mem.Allocator) !?Part {
            const reader = self.stream.reader();
            while (true) {
                try reader.skipUntilDelimiterOrEof('\r');
                var line_buf: [read_ahead]u8 = undefined;
                const len = try reader.readAll(line_buf[0 .. self.boundary.len + 3]);
                const line = line_buf[0..len];
                if (line.len == 0) return null;
                if (std.mem.startsWith(u8, line, "\n--") and std.mem.endsWith(u8, line, self.boundary)) {
                    // match, check for end thing
                    var more_buf: [2]u8 = undefined;
                    if (try reader.readAll(&more_buf) != 2) return error.EndOfStream;

                    const more = !(more_buf[0] == '-' and more_buf[1] == '-');
                    try self.stream.putBack(&more_buf);
                    try reader.skipUntilDelimiterOrEof('\n');
                    if (more) return try Part.open(self, alloc) else return null;
                }
            }
        }

        pub const Part = struct {
            base: ?*Multipart,
            fields: fields.Fields,

            pub fn open(base: *Multipart, alloc: std.mem.Allocator) !Part {
                var parsed_fields = try @import("./request/parser.zig").parseHeaders(alloc, base.stream.reader());
                return .{ .base = base, .fields = parsed_fields };
            }

            pub fn reader(self: *Part) PartReader {
                return .{ .context = self };
            }

            pub fn close(self: *Part) void {
                self.fields.deinit();
            }

            pub fn read(self: *Part, buf: []u8) ReaderType.Error!usize {
                const base = self.base orelse return 0;

                const r = base.stream.reader();

                var count: usize = 0;
                while (count < buf.len) {
                    const byte = r.readByte() catch |err| switch (err) {
                        error.EndOfStream => {
                            self.base = null;
                            return count;
                        },
                        else => |e| return e,
                    };

                    buf[count] = byte;
                    count += 1;
                    if (byte != '\r') continue;

                    var line_buf: [read_ahead]u8 = undefined;
                    const line = line_buf[0..try r.readAll(line_buf[0 .. base.boundary.len + 3])];
                    if (!std.mem.startsWith(u8, line, "\n--") or !std.mem.endsWith(u8, line, base.boundary)) {
                        base.stream.putBack(line) catch unreachable;
                        continue;
                    } else {
                        base.stream.putBack(line) catch unreachable;
                        base.stream.putBackByte('\r') catch unreachable;
                        self.base = null;
                        return count - 1;
                    }
                }

                return count;
            }
        };
    };
}

pub fn openMultipart(boundary: []const u8, reader: anytype) !MultipartStream(@TypeOf(reader)) {
    if (boundary.len > max_boundary) return error.BoundaryTooLarge;
    var stream = .{
        .stream = std.io.peekStream(read_ahead, reader),
        .boundary = boundary,
    };

    stream.stream.putBack("\r\n") catch unreachable;
    return stream;
}

pub fn parseFormData(comptime T: type, boundary: []const u8, reader: anytype, alloc: std.mem.Allocator) !T {
    var multipart = try openMultipart(boundary, reader);

    var ds = util.Deserializer(T){};
    defer {
        var iter = ds.iterator();
        while (iter.next()) |pair| {
            alloc.free(pair.value);
        }
    }
    while (true) {
        var part = (try multipart.next(alloc)) orelse break;
        defer part.close();

        const disposition = part.fields.get("Content-Disposition") orelse return error.InvalidForm;

        const name = fields.getParam(disposition, "name") orelse return error.InvalidForm;

        const value = try part.reader().readAllAlloc(alloc, 1 << 32);
        errdefer alloc.free(value);
        try ds.setSerializedField(name, value);
    }

    return try ds.finish(alloc);
}

// TODO: Fix these tests
test "MultipartStream" {
    const ExpectedPart = struct {
        disposition: []const u8,
        value: []const u8,
    };
    const testCase = struct {
        fn case(comptime body: []const u8, boundary: []const u8, expected_parts: []const ExpectedPart) !void {
            var src = std.io.StreamSource{
                .const_buffer = std.io.fixedBufferStream(body),
            };

            var stream = try openMultipart(boundary, src.reader());

            for (expected_parts) |expected| {
                var part = try stream.next(std.testing.allocator) orelse return error.TestExpectedEqual;
                defer part.close();

                const dispo = part.fields.get("Content-Disposition") orelse return error.TestExpectedEqual;
                try std.testing.expectEqualStrings(expected.disposition, dispo);

                var buf: [128]u8 = undefined;
                const count = try part.reader().read(&buf);
                try std.testing.expectEqualStrings(expected.value, buf[0..count]);
            }

            try std.testing.expect(try stream.next(std.testing.allocator) == null);
        }
    }.case;

    try testCase("--abc--\r\n", "abc", &.{});
    try testCase(
        util.comptimeToCrlf(
            \\------abcd
            \\Content-Disposition: form-data; name=first; charset=utf8
            \\
            \\content
            \\------abcd
            \\content-Disposition: form-data; name=second
            \\
            \\no content
            \\------abcd
            \\content-disposition: form-data; name=third
            \\
            \\
            \\------abcd--
            \\
        ),
        "----abcd",
        &.{
            .{ .disposition = "form-data; name=first; charset=utf8", .value = "content" },
            .{ .disposition = "form-data; name=second", .value = "no content" },
            .{ .disposition = "form-data; name=third", .value = "" },
        },
    );

    try testCase(
        util.comptimeToCrlf(
            \\--xyz
            \\Content-Disposition: uhh
            \\
            \\xyz
            \\--xyz
            \\Content-disposition: ok
            \\
            \\  --xyz
            \\--xyz--
            \\
        ),
        "xyz",
        &.{
            .{ .disposition = "uhh", .value = "xyz" },
            .{ .disposition = "ok", .value = "  --xyz" },
        },
    );
}

test "parseFormData" {
    const body = util.comptimeToCrlf(
        \\--abcd
        \\Content-Disposition: form-data; name=foo
        \\
        \\content
        \\--abcd--
        \\
    );
    if (true) return error.SkipZigTest; // TODO
    var src = std.io.StreamSource{ .const_buffer = std.io.fixedBufferStream(body) };
    const val = try parseFormData(struct {
        foo: []const u8,
    }, "abcd", src.reader(), std.testing.allocator);
    _ = val;
}