zms.zig

const std = @import("std");

const BIND_HOST: []const u8 = "0.0.0.0";
const BIND_PORT: u16 = 8090;
const NUM_WORKER_THREADS: u32 = 16;
const DOWNLOAD_DIR: []const u8 = "/tmp";
const MAX_FILE_SIZE: usize = 1024 * 1024 * 500; // 500M
// const UPSTREAM_URL: []const u8 = "https://fs.liujiacai.net/zigbuilds";
const UPSTREAM_URL: []const u8 = "https://ziglang.org";

pub fn main() !void {
    std.log.info("Welcome to Zig Mirror Server...", .{});
    const allocator = std.heap.page_allocator;

    const bind_addr = try std.net.Address.parseIp(BIND_HOST, BIND_PORT);
    var server = try bind_addr.listen(.{
        .kernel_backlog = 128,
        .reuse_address = true,
    });
    var pool = try allocator.create(std.Thread.Pool);
    defer pool.deinit();

    try pool.init(.{
        .allocator = allocator,
        .n_jobs = NUM_WORKER_THREADS,
    });
    std.log.info("Tcp proxy listen on {any}", .{bind_addr});
    while (true) {
        const conn = server.accept() catch |err| {
            std.log.err("Accept connection failed, err: {s}", .{@errorName(err)});
            continue;
        };
        std.log.debug("Got new connection, addr:{any}", .{conn.address});
        var ctx = Context{ .arena = std.heap.ArenaAllocator.init(allocator) };
        pool.spawn(accept, .{ &ctx, conn }) catch |err| {
            std.log.err("Spawn worker task failed, err: {s}", .{@errorName(err)});
            continue;
        };
    }
}

fn accept(ctx: *Context, conn: std.net.Server.Connection) void {
    defer conn.stream.close();
    defer ctx.deinit();

    var read_buffer: [0x4000]u8 = undefined;
    var server = std.http.Server.init(conn, &read_buffer);
    while (server.state == .ready) {
        var request = server.receiveHead() catch |err| switch (err) {
            error.HttpConnectionClosing => return,
            else => {
                std.log.err("closing http connection: {s}", .{@errorName(err)});
                return;
            },
        };
        serveRequest(ctx, &request) catch |err| switch (err) {
            // error.AlreadyReported => return,
            else => |e| {
                std.log.err("unable to serve {s}: {s}", .{ request.head.target, @errorName(e) });
                return;
            },
        };
    }
}

const Context = struct {
    arena: std.heap.ArenaAllocator,

    fn deinit(ctx: Context) void {
        ctx.arena.deinit();
    }
};

fn loadFromDisk(arena_allocator: std.mem.Allocator, filepath: []const u8) ![]u8 {
    const tarball_path = try std.fmt.allocPrint(arena_allocator, "{s}/{s}", .{ DOWNLOAD_DIR, filepath });
    std.log.debug("Load {s} from disk {s}", .{ filepath, tarball_path });

    const f = try std.fs.openFileAbsolute(tarball_path, .{});
    defer f.close();

    return try f.readToEndAlloc(arena_allocator, MAX_FILE_SIZE);
}

fn loadFromUpstream(arena_allocator: std.mem.Allocator, filename: []const u8) ![]u8 {
    const sub_dir = if (std.mem.indexOf(u8, filename, "dev")) |_|
        // "https://ziglang.org/builds/zig-linux-armv7a-0.14.0-dev.1651+ffd071f55.tar.xz"
        "builds"
    else blk: {
        // We need to extract the version number from the filename
        // "https://ziglang.org/download/0.13.0/zig-linux-x86_64-0.13.0.tar.xz"
        const a = std.mem.lastIndexOfScalar(u8, filename, '-') orelse unreachable;
        const b = std.mem.indexOfScalarPos(u8, filename, a + 1, '.') orelse unreachable; // major version
        const c = std.mem.indexOfScalarPos(u8, filename, b + 1, '.') orelse unreachable; // min version
        const d = std.mem.indexOfScalarPos(u8, filename, c + 1, '.') orelse unreachable; // patch version

        const version = filename[a + 1 .. d];
        break :blk try std.fmt.allocPrint(arena_allocator, "download/{s}", .{version});
    };
    const tarball_url = try std.fmt.allocPrint(arena_allocator, "{s}/{s}/{s}", .{
        UPSTREAM_URL,
        sub_dir,
        filename,
    });
    std.log.debug("Downloading {s} from upstream", .{tarball_url});

    var client = std.http.Client{ .allocator = arena_allocator };
    var resp_buffer = std.ArrayList(u8).init(arena_allocator);
    const ret = try client.fetch(.{
        .location = .{ .url = tarball_url },
        .response_storage = .{ .dynamic = &resp_buffer },
        .max_append_size = MAX_FILE_SIZE,
    });

    if (ret.status != .ok) {
        std.log.err("Failed to download {s} from upstream: {d}", .{ tarball_url, ret.status });
        return error.UnexpectedHttpStatusCode;
    }
    return resp_buffer.items;
}

fn loadTarball(arena_allocator: std.mem.Allocator, filename: []const u8) ![]u8 {
    return loadFromDisk(arena_allocator, filename) catch |err| switch (err) {
        error.FileNotFound => {
            return loadFromUpstream(arena_allocator, filename);
        },
        else => {
            return err;
        },
    };
}

fn serveRequest(ctx: *Context, request: *std.http.Server.Request) !void {
    const allocator = ctx.arena.allocator();
    const path = request.head.target;
    const i = std.mem.lastIndexOfScalar(u8, path, '/') orelse unreachable;
    const requested_file = path[i + 1 ..];

    const bytes = try loadTarball(allocator, requested_file);

    try request.respond(bytes, .{
        .extra_headers = &.{
            .{ .name = "content-type", .value = "application/octet-stream" },
        },
    });
}