http header object for zig

http_header.zig

// HTTP Header data structure/type
// Based on lua-http's http.header module
//
// Design criteria:
//   - the same header field is allowed more than once
//       - must be able to fetch seperate occurences (important for some headers e.g. Set-Cookie)
//       - optionally available as comma seperated list
//   - http2 adds flag to headers that they should never be indexed
//   - header order should be recoverable
//
// Headers are implemented as an array of entries.
// An index of field name => array indices is kept.

const std = @import("std");
const debug = std.debug;
const assert = debug.assert;
const testing = std.testing;
const mem = std.mem;
const Allocator = mem.Allocator;

fn never_index_default(name: []const u8) bool {
    if (mem.eql(u8, "authorization", name)) return true;
    if (mem.eql(u8, "proxy-authorization", name)) return true;
    if (mem.eql(u8, "cookie", name)) return true;
    if (mem.eql(u8, "set-cookie", name)) return true;
    return false;
}

const HeaderEntry = struct {
    allocator: *Allocator,
    pub name: []const u8,
    pub value: []u8,
    pub never_index: bool,

    const Self = @This();

    fn init(allocator: *Allocator, name: []const u8, value: []const u8, never_index: ?bool) !Self {
        return Self{
            .allocator = allocator,
            .name = name, // takes reference
            .value = try mem.dupe(allocator, u8, value),
            .never_index = never_index orelse never_index_default(name),
        };
    }

    fn deinit(self: Self) void {
        self.allocator.free(self.value);
    }

    pub fn modify(self: *Self, value: []const u8, never_index: ?bool) !void {
        const old_len = self.value.len;
        if (value.len > old_len) {
            // realloc first
            self.value = try self.allocator.realloc(self.value, value.len);
        }
        mem.copy(u8, self.value, value);
        if (value.len < old_len) {
            // shrink after
            self.value = self.allocator.shrink(self.value, value.len);
        }
        self.never_index = never_index orelse never_index_default(self.name);
    }

    fn compare(a: HeaderEntry, b: HeaderEntry) bool {
        if (a.name.ptr != b.name.ptr and a.name.len != b.name.len) {
            // Things beginning with a colon *must* be before others
            const a_is_colon = a.name[0] == ':';
            const b_is_colon = b.name[0] == ':';
            if (a_is_colon and !b_is_colon) {
                return true;
            } else if (!a_is_colon and b_is_colon) {
                return false;
            }

            // Sort lexigraphically on header name
            return mem.compare(u8, a.name, b.name) == mem.Compare.LessThan;
        }

        // Sort lexigraphically on header value
        if (!mem.eql(u8, a.value, b.value)) {
            return mem.compare(u8, a.value, b.value) == mem.Compare.LessThan;
        }

        // Doesn't matter here; need to pick something for sort consistency
        return a.never_index;
    }
};

test "HeaderEntry" {
    var e = try HeaderEntry.init(debug.global_allocator, "foo", "bar", null);
    defer e.deinit();
    testing.expectEqualSlices(u8, "foo", e.name);
    testing.expectEqualSlices(u8, "bar", e.value);
    testing.expectEqual(false, e.never_index);

    try e.modify("longer value", null);
    testing.expectEqualSlices(u8, "longer value", e.value);

    // shorter value
    try e.modify("x", null);
    testing.expectEqualSlices(u8, "x", e.value);
}

const HeaderList = std.ArrayList(HeaderEntry);
const HeaderIndexList = std.ArrayList(usize);

// https://github.com/ziglang/zig/issues/2247
fn ArrayListRemove(comptime T: type) type {
    return struct {
        const Self = std.ArrayList(T);

        fn remove(self: *Self, i: usize) T {
            const newlen = self.len - 1;
            if (newlen == i) return self.pop();

            const old_item = self.at(i);
            for (self.items[i..newlen]) |*b, j| b.* = self.items[i + 1 + j];
            self.items[newlen] = undefined;
            self.len = newlen;
            return old_item;
        }
    };
}

test "ArrayListRemove" {
    var list = std.ArrayList(i32).init(debug.global_allocator);
    defer list.deinit();
    const remove = ArrayListRemove(i32).remove;

    try list.append(1);
    try list.append(2);
    try list.append(3);
    try list.append(4);
    try list.append(5);
    try list.append(6);
    try list.append(7);

    //remove from middle
    testing.expectEqual(i32(4), remove(&list, 3));
    testing.expectEqual(i32(5), list.at(3));
    testing.expectEqual(usize(6), list.len);

    //remove from end
    testing.expectEqual(i32(7), remove(&list, 5));
    testing.expectEqual(usize(5), list.len);

    //remove from front
    testing.expectEqual(i32(1), remove(&list, 0));
    testing.expectEqual(i32(2), list.at(0));
    testing.expectEqual(usize(4), list.len);
}

const HeaderIndex = std.AutoHashMap([]const u8, HeaderIndexList);

pub const Headers = struct {
    /// the owned header field name is stored in the index as part of the key
    allocator: *Allocator,
    data: HeaderList,
    index: HeaderIndex,

    const Self = @This();

    pub fn init(allocator: *Allocator) Self {
        return Self{
            .allocator = allocator,
            .data = HeaderList.init(allocator),
            .index = HeaderIndex.init(allocator),
        };
    }

    pub fn deinit(self: Self) void {
        var it = self.data.iterator();
        while (it.next()) |entry| {
            entry.deinit();
        }
        self.data.deinit();
        self.index.deinit();
    }

    pub fn clone(self: Self, allocator: *Allocator) !Self {
        var other = Headers.init(allocator);
        errdefer other.deinit();
        try other.data.ensureCapacity(self.data.count());
        try other.index.initCapacity(self.index.entries.len);
        var it = self.data.iterator();
        while (it.next()) |entry| {
            try other.append(entry.name, entry.value, entry.never_index);
        }
        return other;
    }

    pub fn count(self: Self) usize {
        return self.data.count();
    }

    pub const Iterator = HeaderList.Iterator;

    pub fn iterator(self: Self) Iterator {
        return self.data.iterator();
    }

    pub fn append(self: *Self, name: []const u8, value: []const u8, never_index: ?bool) !void {
        const n = self.data.count() + 1;
        try self.data.ensureCapacity(n);
        var entry: HeaderEntry = undefined;
        if (self.index.get(name)) |kv| {
            entry = try HeaderEntry.init(self.allocator, kv.key, value, never_index);
            var dex = &kv.value;
            try dex.append(n - 1);
        } else {
            const name_dup = try mem.dupe(self.allocator, u8, name);
            errdefer self.allocator.free(name_dup);
            entry = try HeaderEntry.init(self.allocator, name_dup, value, never_index);
            var dex = HeaderIndexList.init(self.allocator);
            try dex.append(n - 1);
            errdefer dex.deinit();
            _ = try self.index.put(name, dex);
        }
        self.data.appendAssumeCapacity(entry);
    }

    pub fn upsert(self: *Self, name: []const u8, value: []const u8, never_index: ?bool) !void {
        if (self.index.get(name)) |kv| {
            const dex = kv.value;
            if (dex.count() != 1)
                return error.CannotUpsertMultiValuedField;
            var e = &self.data.at(dex.at(0));
            try e.modify(value, never_index);
        } else {
            try self.append(name, value, never_index);
        }
    }

    /// Returns boolean indicating if the field is present
    pub fn contains(self: Self, name: []const u8) bool {
        return self.index.contains(name);
    }

    /// Returns boolean indicating if something was deleted
    pub fn delete(self: *Self, name: []const u8) bool {
        if (self.index.get(name)) |kv| {
            var dex = &kv.value;
            // iterate backwards
            var i = dex.count();
            while (i > 0) {
                i -= 1;
                const data_index = dex.at(i);
                const removed = ArrayListRemove(HeaderEntry).remove(&self.data, data_index);
                assert(mem.eql(u8, removed.name, name));
            }
            dex.deinit();
            self.allocator.free(kv.key);
            _ = self.index.remove(name);
            self.rebuild_index();
            return true;
        } else {
            return false;
        }
    }

    /// Removes the element at the specified index
    /// Moves items down to fill the empty slot
    pub fn remove(self: *Self, i: usize) void {
        const removed = ArrayListRemove(HeaderEntry).remove(&self.data, i);
        const kv = self.index.get(removed.name).?;
        var dex = &kv.value;
        if (dex.count() == 1) {
            // was last item; delete the index
            dex.deinit();
            self.allocator.free(kv.key);
        } else {
            dex.shrink(dex.count() - 1);
        }
        // if it was the last item; no need to rebuild index
        if (i != self.data.count()) {
            self.rebuild_index();
        }
    }

    /// Removes the element at the specified index..
    /// The empty slot is filled from the end of the list.
    pub fn swapRemove(self: *Self, i: usize) void {
        const removed = self.data.swapRemove(i);
        const kv = self.index.get(removed.name).?;
        var dex = &kv.value;
        if (dex.count() == 1) {
            // was last item; delete the index
            self.allocator.free(kv.key);
            dex.deinit();
        } else {
            dex.shrink(dex.count() - 1);
        }
        // if it was the last item; no need to rebuild index
        if (i != self.data.count()) {
            self.rebuild_index();
        }
    }

    pub fn at(self: Self, i: usize) HeaderEntry {
        return self.data.at(i);
    }

    /// Returns a list of indices containing headers with the given name
    /// The returned list should not be modified by the caller
    pub fn getIndices(self: Self, name: []const u8) ?HeaderIndexList {
        if (self.index.get(name)) |kv| {
            return kv.value;
        } else {
            return null;
        }
    }

    /// Returns a slice
    pub fn get(self: Self, allocator: *Allocator, name: []const u8) !?[]const HeaderEntry {
        const dex = self.getIndices(name) orelse return null;

        const buf = try allocator.alloc(HeaderEntry, dex.count());
        var it = dex.iterator();
        var n: usize = 0;
        while (it.next()) |idx| {
            buf[n] = self.data.at(idx);
            n += 1;
        }
        return buf;
    }

    pub fn getCommaSeparated(self: Self, allocator: *Allocator, name: []const u8) !?[]u8 {
        const dex = self.getIndices(name) orelse return null;

        // adapted from mem.join
        const total_len = blk: {
            var sum: usize = dex.count() - 1; // space for separators
            var it = dex.iterator();
            while (it.next()) |idx|
                sum += self.data.at(idx).value.len;
            break :blk sum;
        };

        const buf = try allocator.alloc(u8, total_len);
        errdefer allocator.free(buf);

        const first_value = self.data.at(dex.at(0)).value;
        mem.copy(u8, buf, first_value);
        var buf_index: usize = first_value.len;
        for (dex.toSlice()[1..]) |idx| {
            const value = self.data.at(idx).value;
            buf[buf_index] = ',';
            buf_index += 1;
            mem.copy(u8, buf[buf_index..], value);
            buf_index += value.len;
        }

        // No need for shrink since buf is exactly the correct size.
        return buf;
    }

    fn rebuild_index(self: *Self) void {
        { // clear out the indexes
            var it = self.index.iterator();
            while (it.next()) |kv| {
                var dex = &kv.value;
                dex.len = 0; // keeps capacity available
            }
        }
        { // fill up indexes again; we know capacity is fine from before
            var it = self.data.iterator();
            while (it.next()) |entry| {
                var dex = &self.index.get(entry.name).?.value;
                dex.appendAssumeCapacity(it.count);
            }
        }
    }

    pub fn sort(self: *Self) void {
        std.sort.sort(HeaderEntry, self.data.toSlice(), HeaderEntry.compare);
        rebuild_index(self);
    }

    pub fn format(self: Self, comptime fmt: []const u8, context: var, comptime Errors: type, output: fn (@typeOf(context), []const u8) Errors!void) Errors!void {
        var it = self.iterator();
        while (it.next()) |entry| {
            try output(context, entry.name);
            try output(context, ": ");
            try output(context, entry.value);
            try output(context, "\n");
        }
    }
};

test "Headers.iterator" {
    var h = Headers.init(debug.global_allocator);
    defer h.deinit();
    try h.append("foo", "bar", null);
    try h.append("cookie", "somevalue", null);

    var count: i32 = 0;
    var it = h.iterator();
    while (it.next()) |e| {
        if (count == 0) {
            testing.expectEqualSlices(u8, "foo", e.name);
            testing.expectEqualSlices(u8, "bar", e.value);
            testing.expectEqual(false, e.never_index);
        } else if (count == 1) {
            testing.expectEqualSlices(u8, "cookie", e.name);
            testing.expectEqualSlices(u8, "somevalue", e.value);
            testing.expectEqual(true, e.never_index);
        }
        count += 1;
    }
    testing.expectEqual(i32(2), count);
}

test "Headers.contains" {
    var h = Headers.init(debug.global_allocator);
    defer h.deinit();
    try h.append("foo", "bar", null);
    try h.append("cookie", "somevalue", null);

    testing.expectEqual(true, h.contains("foo"));
    testing.expectEqual(false, h.contains("flooble"));
}

test "Headers.delete" {
    var h = Headers.init(debug.global_allocator);
    defer h.deinit();
    try h.append("foo", "bar", null);
    try h.append("baz", "qux", null);
    try h.append("cookie", "somevalue", null);

    testing.expectEqual(false, h.delete("not-present"));
    testing.expectEqual(usize(3), h.count());

    testing.expectEqual(true, h.delete("foo"));
    testing.expectEqual(usize(2), h.count());
    {
        const e = h.at(0);
        testing.expectEqualSlices(u8, "baz", e.name);
        testing.expectEqualSlices(u8, "qux", e.value);
        testing.expectEqual(false, e.never_index);
    }
    {
        const e = h.at(1);
        testing.expectEqualSlices(u8, "cookie", e.name);
        testing.expectEqualSlices(u8, "somevalue", e.value);
        testing.expectEqual(true, e.never_index);
    }

    testing.expectEqual(false, h.delete("foo"));
}

test "Headers.remove" {
    var h = Headers.init(debug.global_allocator);
    defer h.deinit();
    try h.append("foo", "bar", null);
    try h.append("baz", "qux", null);
    try h.append("cookie", "somevalue", null);

    h.remove(0);
    testing.expectEqual(usize(2), h.count());
    {
        const e = h.at(0);
        testing.expectEqualSlices(u8, "baz", e.name);
        testing.expectEqualSlices(u8, "qux", e.value);
        testing.expectEqual(false, e.never_index);
    }
    {
        const e = h.at(1);
        testing.expectEqualSlices(u8, "cookie", e.name);
        testing.expectEqualSlices(u8, "somevalue", e.value);
        testing.expectEqual(true, e.never_index);
    }
}

test "Headers.swapRemove" {
    var h = Headers.init(debug.global_allocator);
    defer h.deinit();
    try h.append("foo", "bar", null);
    try h.append("baz", "qux", null);
    try h.append("cookie", "somevalue", null);

    h.swapRemove(0);
    testing.expectEqual(usize(2), h.count());
    {
        const e = h.at(0);
        testing.expectEqualSlices(u8, "cookie", e.name);
        testing.expectEqualSlices(u8, "somevalue", e.value);
        testing.expectEqual(true, e.never_index);
    }
    {
        const e = h.at(1);
        testing.expectEqualSlices(u8, "baz", e.name);
        testing.expectEqualSlices(u8, "qux", e.value);
        testing.expectEqual(false, e.never_index);
    }
}

test "Headers.at" {
    var h = Headers.init(debug.global_allocator);
    defer h.deinit();
    try h.append("foo", "bar", null);
    try h.append("cookie", "somevalue", null);

    {
        const e = h.at(0);
        testing.expectEqualSlices(u8, "foo", e.name);
        testing.expectEqualSlices(u8, "bar", e.value);
        testing.expectEqual(false, e.never_index);
    }
    {
        const e = h.at(1);
        testing.expectEqualSlices(u8, "cookie", e.name);
        testing.expectEqualSlices(u8, "somevalue", e.value);
        testing.expectEqual(true, e.never_index);
    }
}

test "Headers.getIndices" {
    var h = Headers.init(debug.global_allocator);
    defer h.deinit();
    try h.append("foo", "bar", null);
    try h.append("set-cookie", "x=1", null);
    try h.append("set-cookie", "y=2", null);

    testing.expect(null == h.getIndices("not-present"));
    testing.expectEqualSlices(usize, []usize{0}, h.getIndices("foo").?.toSliceConst());
    testing.expectEqualSlices(usize, []usize{1,2}, h.getIndices("set-cookie").?.toSliceConst());
}

test "Headers.get" {
    var h = Headers.init(debug.global_allocator);
    defer h.deinit();
    try h.append("foo", "bar", null);
    try h.append("set-cookie", "x=1", null);
    try h.append("set-cookie", "y=2", null);

    {
        const v = try h.get(debug.global_allocator, "not-present");
        testing.expect(null == v);
    }
    {
        const v = (try h.get(debug.global_allocator, "foo")).?;
        defer debug.global_allocator.free(v);
        const e = v[0];
        testing.expectEqualSlices(u8, "foo", e.name);
        testing.expectEqualSlices(u8, "bar", e.value);
        testing.expectEqual(false, e.never_index);
    }
    {
        const v = (try h.get(debug.global_allocator, "set-cookie")).?;
        defer debug.global_allocator.free(v);
        {
            const e = v[0];
            testing.expectEqualSlices(u8, "set-cookie", e.name);
            testing.expectEqualSlices(u8, "x=1", e.value);
            testing.expectEqual(true, e.never_index);
        }
        {
            const e = v[1];
            testing.expectEqualSlices(u8, "set-cookie", e.name);
            testing.expectEqualSlices(u8, "y=2", e.value);
            testing.expectEqual(true, e.never_index);
        }
    }
}

test "Headers.getCommaSeparated" {
    var h = Headers.init(debug.global_allocator);
    defer h.deinit();
    try h.append("foo", "bar", null);
    try h.append("set-cookie", "x=1", null);
    try h.append("set-cookie", "y=2", null);

    {
        const v = try h.getCommaSeparated(debug.global_allocator, "not-present");
        testing.expect(null == v);
    }
    {
        const v = (try h.getCommaSeparated(debug.global_allocator, "foo")).?;
        defer debug.global_allocator.free(v);
        testing.expectEqualSlices(u8, "bar", v);
    }
    {
        const v = (try h.getCommaSeparated(debug.global_allocator, "set-cookie")).?;
        defer debug.global_allocator.free(v);
        testing.expectEqualSlices(u8, "x=1,y=2", v);
    }
}

test "Headers.sort" {
    var h = Headers.init(debug.global_allocator);
    defer h.deinit();
    try h.append("foo", "bar", null);
    try h.append("cookie", "somevalue", null);

    h.sort();
    {
        const e = h.at(0);
        testing.expectEqualSlices(u8, "cookie", e.name);
        testing.expectEqualSlices(u8, "somevalue", e.value);
        testing.expectEqual(true, e.never_index);
    }
    {
        const e = h.at(1);
        testing.expectEqualSlices(u8, "foo", e.name);
        testing.expectEqualSlices(u8, "bar", e.value);
        testing.expectEqual(false, e.never_index);
    }
}

test "Headers.format" {
    var h = Headers.init(debug.global_allocator);
    defer h.deinit();
    try h.append("foo", "bar", null);
    try h.append("cookie", "somevalue", null);

    var buf: [100]u8 = undefined;
    testing.expectEqualSlices(u8,
        \\foo: bar
        \\cookie: somevalue
        \\
    , try std.fmt.bufPrint(buf[0..], "{}", h));
}