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 BaseReader = ReaderType;
        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;
}

const MultipartFormField = struct {
    name: []const u8,
    value: []const u8,

    filename: ?[]const u8 = null,
    content_type: ?[]const u8 = null,
};

pub const FormFile = struct {
    data: []const u8,
    filename: []const u8,
    content_type: []const u8,
};

pub fn MultipartForm(comptime ReaderType: type) type {
    return struct {
        stream: MultipartStream(ReaderType),

        pub fn next(self: *@This(), alloc: std.mem.Allocator) !?MultipartFormField {
            var part = (try self.stream.next(alloc)) orelse return null;
            defer part.close();

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

            if (!std.ascii.eqlIgnoreCase(fields.getParam(disposition, null).?, "form-data")) return error.BadDisposition;
            const name = try util.deepClone(alloc, fields.getParam(disposition, "name") orelse return error.BadDisposition);
            errdefer util.deepFree(alloc, name);
            const filename = try util.deepClone(alloc, fields.getParam(disposition, "filename"));
            errdefer util.deepFree(alloc, filename);
            const content_type = try util.deepClone(alloc, part.fields.get("Content-Type"));
            errdefer util.deepFree(alloc, content_type);

            const value = try part.reader().readAllAlloc(alloc, 1 << 32);

            return MultipartFormField{
                .name = name,
                .value = value,

                .filename = filename,
                .content_type = content_type,
            };
        }
    };
}

pub fn openForm(multipart_stream: anytype) MultipartForm(@TypeOf(multipart_stream).BaseReader) {
    return .{ .stream = multipart_stream };
}

fn Deserializer(comptime Result: type) type {
    return util.DeserializerContext(Result, MultipartFormField, struct {
        pub const options = .{ .isScalar = isScalar, .embed_unions = true };

        pub fn isScalar(comptime T: type) bool {
            if (T == FormFile or T == ?FormFile) return true;
            return util.serialize.defaultIsScalar(T);
        }

        pub fn deserializeScalar(_: @This(), alloc: std.mem.Allocator, comptime T: type, val: MultipartFormField) !T {
            if (T == FormFile or T == ?FormFile) return try deserializeFormFile(alloc, val);

            if (val.filename != null) return error.FilenameProvidedForNonFile;
            return try util.serialize.deserializeString(alloc, T, val.value);
        }

        fn deserializeFormFile(alloc: std.mem.Allocator, val: MultipartFormField) !FormFile {
            const data = try util.deepClone(alloc, val.value);
            errdefer util.deepFree(alloc, data);
            const filename = try util.deepClone(alloc, val.filename orelse "(untitled)");
            errdefer util.deepFree(alloc, filename);
            const content_type = try util.deepClone(alloc, val.content_type orelse "application/octet-stream");
            return FormFile{
                .data = data,
                .filename = filename,
                .content_type = content_type,
            };
        }
    });
}

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

    var ds = Deserializer(T){};
    defer {
        var iter = ds.iterator();
        while (iter.next()) |pair| {
            util.deepFree(alloc, pair.value);
        }
    }
    while (true) {
        var part = (try form.next(alloc)) orelse break;
        errdefer util.deepFree(alloc, part);

        try ds.setSerializedField(part.name, part);
    }

    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 "MultipartForm" {
    const testCase = struct {
        fn case(comptime body: []const u8, boundary: []const u8, expected_parts: []const MultipartFormField) !void {
            var src = std.io.StreamSource{
                .const_buffer = std.io.fixedBufferStream(body),
            };

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

            for (expected_parts) |expected| {
                var data = try form.next(std.testing.allocator) orelse return error.TestExpectedEqual;
                defer util.deepFree(std.testing.allocator, data);

                try util.testing.expectDeepEqual(expected, data);
            }

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

    try testCase(
        util.comptimeToCrlf(
            \\--abcd
            \\Content-Disposition: form-data; name=foo
            \\
            \\content
            \\--abcd--
            \\
        ),
        "abcd",
        &.{.{ .name = "foo", .value = "content" }},
    );
    try testCase(
        util.comptimeToCrlf(
            \\--abcd
            \\Content-Disposition: form-data; name=foo
            \\
            \\content
            \\--abcd
            \\Content-Disposition: form-data; name=bar
            \\Content-Type: blah
            \\
            \\abcd
            \\--abcd
            \\Content-Disposition: form-data; name=baz; filename="myfile.txt"
            \\Content-Type: text/plain
            \\
            \\  --abcd
            \\
            \\--abcd--
            \\
        ),
        "abcd",
        &.{
            .{ .name = "foo", .value = "content" },
            .{ .name = "bar", .value = "abcd", .content_type = "blah" },
            .{
                .name = "baz",
                .value = "  --abcd\r\n",
                .content_type = "text/plain",
                .filename = "myfile.txt",
            },
        },
    );
}

test "parseFormData" {
    const body = util.comptimeToCrlf(
        \\--abcd
        \\Content-Disposition: form-data; name=foo
        \\
        \\content
        \\--abcd--
        \\
    );
    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);
    util.deepFree(std.testing.allocator, val);
}