From 35a2b8a006d82574057d46efdabd4a2d1543017c Mon Sep 17 00:00:00 2001 From: Bob Farrell Date: Sat, 1 Jun 2024 15:25:24 +0100 Subject: [PATCH] Refactor to reduce allocs, add tests for cookie flags --- src/jetzig/http/Cookies.zig | 324 ++++++++++++++++++++++++++++-------- src/jetzig/http/Request.zig | 8 +- src/jetzig/http/Session.zig | 2 +- src/jetzig/util.zig | 4 + 4 files changed, 263 insertions(+), 75 deletions(-) diff --git a/src/jetzig/http/Cookies.zig b/src/jetzig/http/Cookies.zig index ae0b3ce..30c5fe5 100644 --- a/src/jetzig/http/Cookies.zig +++ b/src/jetzig/http/Cookies.zig @@ -4,67 +4,80 @@ const jetzig = @import("../../jetzig.zig"); allocator: std.mem.Allocator, cookie_string: []const u8, -buf: std.ArrayList(u8), cookies: std.StringArrayHashMap(*Cookie), modified: bool = false, +arena: std.heap.ArenaAllocator, const Self = @This(); +const SameSite = enum { strict, lax, none }; pub const CookieOptions = struct { domain: []const u8 = "localhost", path: []const u8 = "/", - same_site: ?enum { strict, lax, none } = null, + same_site: ?SameSite = null, secure: bool = false, expires: ?i64 = null, // if used, set to time in seconds to be added to std.time.timestamp() http_only: bool = false, max_age: ?i64 = null, partitioned: bool = false, - - /// Builds a cookie string - pub fn allocPrint(self: CookieOptions, allocator: std.mem.Allocator, cookie_name: []const u8, cookie_value: []const u8) ![]const u8 { - // add cookie name, cookie value, path, and domain regardless of cookie options - const standard_components = try std.fmt.allocPrint(allocator, "{s}={s}; path={s}; domain={s};", .{ - cookie_name, - cookie_value, - self.path, - self.domain, - }); - var cookie_string = std.ArrayList(u8).fromOwnedSlice(allocator, standard_components); - // secure is required if samesite is set to none - var require_secure = false; - var buf: [256]u8 = undefined; - if (self.same_site) |same_site| { - if (same_site == .none) { - require_secure = true; - } - const kv_pair = try std.fmt.bufPrint(&buf, "{s}={s};", .{ "samesite", @tagName(same_site) }); - try cookie_string.appendSlice(kv_pair); - } - if (self.secure or require_secure) { - try cookie_string.appendSlice("secure;"); - } - if (self.expires) |expires| { - const kv_pair = try std.fmt.bufPrint(&buf, "{s}={d};", .{ "expires", std.time.timestamp() + expires }); - try cookie_string.appendSlice(kv_pair); - } - if (self.max_age) |max_age| { - const kv_pair = try std.fmt.bufPrint(&buf, "{s}={d};", .{ "max-age", max_age }); - try cookie_string.appendSlice(kv_pair); - } - if (self.http_only) { - try cookie_string.appendSlice("httponly;"); - } - if (self.partitioned) { - try cookie_string.appendSlice("partitioned;"); - } - return cookie_string.toOwnedSlice(); - } }; const cookie_options = jetzig.config.get(CookieOptions, "cookie_options"); pub const Cookie = struct { + name: []const u8, value: []const u8, + domain: ?[]const u8 = null, + path: ?[]const u8 = null, + same_site: ?SameSite = null, + secure: ?bool = null, + expires: ?i64 = null, // if used, set to time in seconds to be added to std.time.timestamp() + http_only: ?bool = null, + max_age: ?i64 = null, + partitioned: ?bool = null, + + /// Build a cookie string. + pub fn bufPrint(self: Cookie, buf: *[4096]u8) ![]const u8 { + var options = cookie_options; + inline for (std.meta.fields(CookieOptions)) |field| { + @field(options, field.name) = @field(self, field.name) orelse @field(cookie_options, field.name); + } + + // secure is required if samesite is set to none + const require_secure = if (options.same_site) |same_site| same_site == .none else false; + + var stream = std.io.fixedBufferStream(buf); + const writer = stream.writer(); + + try writer.print("{s}={s}; path={s}; domain={s};", .{ + self.name, + self.value, + options.path, + options.domain, + }); + + if (options.same_site) |same_site| try writer.print(" SameSite={s};", .{@tagName(same_site)}); + if (options.secure or require_secure) try writer.writeAll(" Secure;"); + if (options.expires) |expires| try writer.print(" Expires={d};", .{std.time.timestamp() + expires}); + if (options.max_age) |max_age| try writer.print(" Max-Age={d};", .{max_age}); + if (options.http_only) try writer.writeAll(" HttpOnly;"); + if (options.partitioned) try writer.writeAll(" Partitioned;"); + + return stream.getWritten(); + } + + pub fn applyFlag(self: *Cookie, allocator: std.mem.Allocator, flag: Flag) !void { + switch (flag) { + .domain => |domain| self.domain = try allocator.dupe(u8, domain), + .path => |path| self.path = try allocator.dupe(u8, path), + .same_site => |same_site| self.same_site = same_site, + .secure => |secure| self.secure = secure, + .expires => |expires| self.expires = expires, + .http_only => |http_only| self.http_only = http_only, + .max_age => |max_age| self.max_age = max_age, + .partitioned => |partitioned| self.partitioned = partitioned, + } + } }; pub fn init(allocator: std.mem.Allocator, cookie_string: []const u8) Self { @@ -72,7 +85,7 @@ pub fn init(allocator: std.mem.Allocator, cookie_string: []const u8) Self { .allocator = allocator, .cookie_string = cookie_string, .cookies = std.StringArrayHashMap(*Cookie).init(allocator), - .buf = std.ArrayList(u8).init(allocator), + .arena = std.heap.ArenaAllocator.init(allocator), }; } @@ -84,51 +97,50 @@ pub fn deinit(self: *Self) void { self.allocator.destroy(item.value_ptr.*); } self.cookies.deinit(); - self.buf.deinit(); + self.arena.deinit(); } pub fn get(self: *Self, key: []const u8) ?*Cookie { return self.cookies.get(key); } -pub fn put(self: *Self, key: []const u8, value: Cookie) !void { +pub fn put(self: *Self, cookie: Cookie) !void { self.modified = true; - if (self.cookies.fetchSwapRemove(key)) |entry| { + if (self.cookies.fetchSwapRemove(cookie.name)) |entry| { self.allocator.free(entry.key); self.allocator.free(entry.value.value); self.allocator.destroy(entry.value); } + const key = try self.allocator.dupe(u8, cookie.name); const ptr = try self.allocator.create(Cookie); - ptr.* = value; - ptr.*.value = try self.allocator.dupe(u8, value.value); - try self.cookies.put(try self.allocator.dupe(u8, key), ptr); + ptr.* = cookie; + ptr.name = key; + ptr.value = try self.allocator.dupe(u8, cookie.value); + try self.cookies.put(key, ptr); } pub const HeaderIterator = struct { allocator: std.mem.Allocator, cookies_iterator: std.StringArrayHashMap(*Cookie).Iterator, + buf: *[4096]u8, - pub fn init(allocator: std.mem.Allocator, cookies: *Self) HeaderIterator { - return .{ .allocator = allocator, .cookies_iterator = cookies.cookies.iterator() }; + pub fn init(allocator: std.mem.Allocator, cookies: *Self, buf: *[4096]u8) HeaderIterator { + return .{ .allocator = allocator, .cookies_iterator = cookies.cookies.iterator(), .buf = buf }; } pub fn next(self: *HeaderIterator) !?[]const u8 { - if (self.cookies_iterator.next()) |*item| { - return try cookie_options.allocPrint(self.allocator, item.key_ptr.*, item.value_ptr.*.value); + if (self.cookies_iterator.next()) |entry| { + const cookie = entry.value_ptr.*; + return try cookie.bufPrint(self.buf); } else { return null; } } }; -pub fn headerIterator(self: *Self) HeaderIterator { - var buf = std.ArrayList([]const u8).init(self.allocator); - - defer buf.deinit(); - defer for (buf.items) |item| self.allocator.free(item); - - return HeaderIterator.init(self.allocator, self); +pub fn headerIterator(self: *Self, buf: *[4096]u8) HeaderIterator { + return HeaderIterator.init(self.allocator, self, buf); } // https://datatracker.ietf.org/doc/html/rfc6265#section-4.2.1 @@ -139,9 +151,12 @@ pub fn parse(self: *Self) !void { var value_buf = std.ArrayList(u8).init(self.allocator); var key_terminated = false; var value_started = false; + var cookie_buf = std.ArrayList(Cookie).init(self.allocator); defer key_buf.deinit(); defer value_buf.deinit(); + defer cookie_buf.deinit(); + defer self.modified = false; for (self.cookie_string, 0..) |char, index| { if (char == '=') { @@ -151,10 +166,14 @@ pub fn parse(self: *Self) !void { if (char == ';' or index == self.cookie_string.len - 1) { if (char != ';') try value_buf.append(char); - try self.put( - key_buf.items, - Cookie{ .value = value_buf.items }, - ); + if (parseFlag(key_buf.items, value_buf.items)) |flag| { + for (cookie_buf.items) |*cookie| try cookie.applyFlag(self.arena.allocator(), flag); + } else { + try cookie_buf.append(.{ + .name = try self.arena.allocator().dupe(u8, key_buf.items), + .value = try self.arena.allocator().dupe(u8, value_buf.items), + }); + } key_buf.clearAndFree(); value_buf.clearAndFree(); value_started = false; @@ -177,7 +196,55 @@ pub fn parse(self: *Self) !void { continue; } - unreachable; + return error.JetzigInvalidCookieHeader; + } + + for (cookie_buf.items) |cookie| try self.put(cookie); +} + +const Flag = union(enum) { + domain: []const u8, + path: []const u8, + same_site: SameSite, + secure: bool, + expires: i64, + max_age: i64, + http_only: bool, + partitioned: bool, +}; + +fn parseFlag(key: []const u8, value: []const u8) ?Flag { + if (key.len > 64) return null; + if (value.len > 64) return null; + + var key_buf: [64]u8 = undefined; + + const normalized_key = std.ascii.lowerString(&key_buf, jetzig.util.strip(key)); + const normalized_value = jetzig.util.strip(value); + + if (std.mem.eql(u8, normalized_key, "domain")) { + return .{ .domain = normalized_value }; + } else if (std.mem.eql(u8, normalized_key, "path")) { + return .{ .path = normalized_value }; + } else if (std.mem.eql(u8, normalized_key, "samesite")) { + return if (std.mem.eql(u8, normalized_value, "strict")) + .{ .same_site = .strict } + else if (std.mem.eql(u8, normalized_value, "lax")) + .{ .same_site = .lax } + else + .{ .same_site = .none }; + } else if (std.mem.eql(u8, normalized_key, "secure")) { + return .{ .secure = true }; + } else if (std.mem.eql(u8, normalized_key, "httponly")) { + return .{ .http_only = true }; + } else if (std.mem.eql(u8, normalized_key, "partitioned")) { + return .{ .partitioned = true }; + } else if (std.mem.eql(u8, normalized_key, "expires")) { + return .{ .expires = std.fmt.parseInt(i64, normalized_value, 10) catch return null }; + } else if (std.mem.eql(u8, normalized_key, "max-age")) { + return .{ .max_age = std.fmt.parseInt(i64, normalized_value, 10) catch return null }; + } else { + return null; } } @@ -217,16 +284,135 @@ test "headerIterator" { defer cookies.deinit(); try cookies.parse(); - var it = cookies.headerIterator(); + var it_buf: [4096]u8 = undefined; + var it = cookies.headerIterator(&it_buf); while (try it.next()) |*header| { try writer.writeAll(header.*); try writer.writeAll("\n"); - allocator.free(header.*); } try std.testing.expectEqualStrings( - \\foo=bar; path=/; domain=localhost - \\baz=qux; path=/; domain=localhost + \\foo=bar; path=/; domain=localhost; + \\baz=qux; path=/; domain=localhost; \\ , buf.items); } + +test "modified" { + const allocator = std.testing.allocator; + var cookies = Self.init(allocator, "foo=bar; baz=qux;"); + defer cookies.deinit(); + + try cookies.parse(); + try std.testing.expect(cookies.modified == false); + + try cookies.put(.{ .name = "quux", .value = "corge" }); + try std.testing.expect(cookies.modified == true); +} + +test "domain=example.com" { + const allocator = std.testing.allocator; + var cookies = Self.init(allocator, "foo=bar; baz=qux; Domain=example.com;"); + defer cookies.deinit(); + + try cookies.parse(); + const cookie = cookies.get("foo").?; + try std.testing.expectEqualStrings(cookie.domain.?, "example.com"); +} + +test "path=/example_path" { + const allocator = std.testing.allocator; + var cookies = Self.init(allocator, "foo=bar; baz=qux; Path=/example_path;"); + defer cookies.deinit(); + + try cookies.parse(); + const cookie = cookies.get("foo").?; + try std.testing.expectEqualStrings(cookie.path.?, "/example_path"); +} + +test "SameSite=lax" { + const allocator = std.testing.allocator; + var cookies = Self.init(allocator, "foo=bar; baz=qux; SameSite=lax;"); + defer cookies.deinit(); + + try cookies.parse(); + const cookie = cookies.get("foo").?; + try std.testing.expect(cookie.same_site == .lax); +} + +test "SameSite=none" { + const allocator = std.testing.allocator; + var cookies = Self.init(allocator, "foo=bar; baz=qux; SameSite=none;"); + defer cookies.deinit(); + + try cookies.parse(); + const cookie = cookies.get("foo").?; + try std.testing.expect(cookie.same_site == .none); +} + +test "SameSite=strict" { + const allocator = std.testing.allocator; + var cookies = Self.init(allocator, "foo=bar; baz=qux; SameSite=strict;"); + defer cookies.deinit(); + + try cookies.parse(); + const cookie = cookies.get("foo").?; + try std.testing.expect(cookie.same_site == .strict); +} + +test "Secure" { + const allocator = std.testing.allocator; + var cookies = Self.init(allocator, "foo=bar; baz=qux; Secure;"); + defer cookies.deinit(); + + try cookies.parse(); + const cookie = cookies.get("foo").?; + try std.testing.expect(cookie.secure.?); +} + +test "Partitioned" { + const allocator = std.testing.allocator; + var cookies = Self.init(allocator, "foo=bar; baz=qux; Partitioned;"); + defer cookies.deinit(); + + try cookies.parse(); + const cookie = cookies.get("foo").?; + try std.testing.expect(cookie.partitioned.?); +} + +test "Max-Age" { + const allocator = std.testing.allocator; + var cookies = Self.init(allocator, "foo=bar; baz=qux; Max-Age=123123123;"); + defer cookies.deinit(); + + try cookies.parse(); + const cookie = cookies.get("foo").?; + try std.testing.expect(cookie.max_age.? == 123123123); +} + +test "Expires" { + const allocator = std.testing.allocator; + var cookies = Self.init(allocator, "foo=bar; baz=qux; Expires=123123123;"); + defer cookies.deinit(); + + try cookies.parse(); + const cookie = cookies.get("foo").?; + try std.testing.expect(cookie.expires.? == 123123123); +} + +test "default flags" { + const allocator = std.testing.allocator; + var cookies = Self.init(allocator, "foo=bar; baz=qux;"); + defer cookies.deinit(); + + try cookies.parse(); + const cookie = cookies.get("foo").?; + try std.testing.expect(cookie.domain == null); + try std.testing.expect(cookie.path == null); + try std.testing.expect(cookie.same_site == null); + try std.testing.expect(cookie.secure == null); + try std.testing.expect(cookie.expires == null); + try std.testing.expect(cookie.http_only == null); + try std.testing.expect(cookie.max_age == null); + try std.testing.expect(cookie.partitioned == null); +} diff --git a/src/jetzig/http/Request.zig b/src/jetzig/http/Request.zig index 9bf1379..d0f80cc 100644 --- a/src/jetzig/http/Request.zig +++ b/src/jetzig/http/Request.zig @@ -567,11 +567,9 @@ fn setCookieHeaders(self: *Request) !void { const local_cookies = self._cookies orelse return; if (!local_cookies.modified) return; - var cookie_it = local_cookies.headerIterator(); - while (try cookie_it.next()) |header| { - // FIXME: Skip setting cookies that are already present ? - try self.response.headers.append("Set-Cookie", header); - } + var buf: [4096]u8 = undefined; + var header_it = local_cookies.headerIterator(&buf); + while (try header_it.next()) |header| try self.response.headers.append("Set-Cookie", header); } // Determine if a given route matches the current request. diff --git a/src/jetzig/http/Session.zig b/src/jetzig/http/Session.zig index 0d0a36b..6f27228 100644 --- a/src/jetzig/http/Session.zig +++ b/src/jetzig/http/Session.zig @@ -85,7 +85,7 @@ fn save(self: *Self) !void { defer self.allocator.free(encrypted); const encoded = try jetzig.util.base64Encode(self.allocator, encrypted); defer self.allocator.free(encoded); - try self.cookies.put(cookie_name, .{ .value = encoded }); + try self.cookies.put(.{ .name = cookie_name, .value = encoded }); } fn parseSessionCookie(self: *Self, cookie_value: []const u8) !void { diff --git a/src/jetzig/util.zig b/src/jetzig/util.zig index 23a0180..0d6c82b 100644 --- a/src/jetzig/util.zig +++ b/src/jetzig/util.zig @@ -49,3 +49,7 @@ pub fn generateSecret(allocator: std.mem.Allocator, comptime len: u10) ![]const pub fn duration(start_time: i128) i64 { return @intCast(std.time.nanoTimestamp() - start_time); } + +pub inline fn strip(string: []const u8) []const u8 { + return std.mem.trim(u8, string, &std.ascii.whitespace); +}