Refactor to reduce allocs, add tests for cookie flags

This commit is contained in:
Bob Farrell 2024-06-01 15:25:24 +01:00
parent 02fd62b415
commit 35a2b8a006
4 changed files with 263 additions and 75 deletions

View File

@ -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);
}

View File

@ -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.

View File

@ -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 {

View File

@ -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);
}