bukubukuchagama/main.zig

//! Reference implementation of compliance with RFC 2324
//! / RFC 7168 section 2.3.2 — HTCPCP "I'm a teapot"
//!
//! Build:  zig build -Doptimize=ReleaseFast
//! Run:    ./zig-out/bin/bukubukuchagama        (listens on 127.0.0.1:8418)
//!
//! Caddyfile:
//!   @garbage {
//!       path *.php *.asp *.aspx *.cgi /wp-login.php /wp-admin/* /.env /.git/* /phpmyadmin/* /xmlrpc.php
//!   }
//!   reverse_proxy @garbage 127.0.0.1:8418
//!

const std = @import("std");
const log = std.log.scoped(.teapot);

const TEAPOT_DIR = "./teapots";
const LISTEN_ADDR = "127.0.0.1";
const LISTEN_PORT = 8418;

/// Populated once at startup, read-only for the lifetime of the process —
/// safe to read from concurrent connection tasks without synchronization.
var images: std.ArrayList([]const u8) = undefined;

fn loadImageList(allocator: std.mem.Allocator, io: std.Io) !void {
    images = .empty;

    var dir = try std.Io.Dir.cwd().openDir(io, TEAPOT_DIR, .{ .iterate = true });
    defer dir.close(io);

    var it = try dir.walk(allocator);
    while (try it.next(io)) |entry| {
        if (entry.kind != .file) continue;
        if (!std.mem.endsWith(u8, entry.basename, ".webp")) continue;
        try images.append(allocator, try allocator.dupe(u8, entry.basename));
    }

    if (images.items.len == 0) return error.NoTeapotImagesFound;
}

fn serveTeapot(allocator: std.mem.Allocator, io: std.Io, rand: std.Random, req: *std.http.Server.Request) !void {
    const idx = rand.intRangeLessThan(u64, 0, images.items.len);
    const name = images.items[idx];

    const path = try std.fmt.allocPrint(allocator, "{s}/{s}", .{ TEAPOT_DIR, name });
    defer allocator.free(path);

    log.info("{s} {s} -> 418 ({s})", .{ @tagName(req.head.method), req.head.target, name });

    const body = std.Io.Dir.cwd().readFileAlloc(io, path, allocator, .unlimited) catch |err| {
        log.err("failed to read {s}: {t}", .{ path, err });
        try req.respond("I'm a teapot.\n", .{ .status = .teapot });
        return;
    };
    defer allocator.free(body);

    try req.respond(body, .{
        .status = .teapot,
        .reason = "I'm a teapot",
        .extra_headers = &.{
            .{ .name = "content-type", .value = "image/webp" },
            .{ .name = "x-htcpcp", .value = "RFC 2324 / RFC 7168 section 2.3.2" },
        },
    });
}

fn handleStream(io: std.Io, stream: std.Io.net.Stream, rand: std.Random) void {
    defer stream.close(io);

    // Per-connection arena: cheap, and guarantees cleanup even if something
    // in the request path errors out partway through.
    var arena_state = std.heap.ArenaAllocator.init(std.heap.page_allocator);
    defer arena_state.deinit();
    const allocator = arena_state.allocator();

    var read_buffer: [4096]u8 = undefined;
    var write_buffer: [4096]u8 = undefined;
    var reader = stream.reader(io, &read_buffer);
    var writer = stream.writer(io, &write_buffer);

    var http_server = std.http.Server.init(&reader.interface, &writer.interface);

    while (true) {
        var req = http_server.receiveHead() catch |err| switch (err) {
            error.HttpConnectionClosing => return,
            else => return,
        };

        serveTeapot(allocator, io, rand, &req) catch |err| {
            log.err("failed to respond: {t}", .{err});
            return;
        };

        _ = arena_state.reset(.retain_capacity);
    }
}

pub fn main(init: std.process.Init) !void {
    const io = init.io;
    const gpa = init.gpa;
    const ioSource: std.Random.IoSource = .{ .io = io };
    const rand = ioSource.interface();

    try loadImageList(gpa, init.io);
    log.info("loaded {d} teapot image(s) from {s}/", .{ images.items.len, TEAPOT_DIR });

    const addr = std.Io.net.IpAddress.parseIp4(LISTEN_ADDR, LISTEN_PORT) catch unreachable;
    var server = try addr.listen(io, .{ .reuse_address = true });
    defer server.deinit(io);

    log.info("teapot responder listening on http://{s}:{d}", .{ LISTEN_ADDR, LISTEN_PORT });

    var group: std.Io.Group = .init;
    defer group.cancel(io);

    while (true) {
        const stream = try server.accept(io);
        group.async(io, handleStream, .{ io, stream, rand });
    }
}