// SPDX-FileCopyrightText: © 2026 Jeffrey C. Ollie <jeff@ocjtech.us>
// SPDX-License-Identifier: MIT

//! zlib streams (RFC 1950) as git stores loose objects and pack entries in
//! them, inflated with the standard library's DEFLATE decoder (RFC 1951)
//! and checked the way zlib checks them: the two header bytes and the
//! Adler-32 checksum at the end, which the decoder itself reads without
//! checking. What zlib would say of a corrupt stream is kept in its own
//! words, since git passes them on.
//!
//! Nothing here does I/O beyond the reader and writer it is handed.

const std = @import("std");
const Allocator = std.mem.Allocator;
const Io = std.Io;
const flate = std.compress.flate;

/// What zlib would say of the two bytes that begin a zlib stream, in the
/// order it checks them, or null if it would accept them or there are not
/// two bytes to check. The DEFLATE decoder checks less than zlib does.
pub fn headerProblem(in: *Io.Reader) error{ReadFailed}!?[]const u8 {
    const h = in.peek(2) catch |err| switch (err) {
        error.EndOfStream => return null,
        error.ReadFailed => return error.ReadFailed,
    };
    if ((@as(u16, h[0]) << 8 | h[1]) % 31 != 0) return "data stream error (incorrect header check)";
    if (h[0] & 0x0f != 8) return "data stream error (unknown compression method)";
    if (h[0] >> 4 > 7) return "data stream error (invalid window size)";
    if (h[1] & 0x20 != 0) return "needs dictionary (no message)";
    return null;
}

/// zlib's words for what the DEFLATE decoder found wrong. The decoder's
/// distinctions are not zlib's, so for the rarer corruptions this is the
/// nearest of zlib's messages rather than certainly the one it would give.
pub fn message(err: ?flate.Decompress.Error) []const u8 {
    return switch (err orelse return "data error") {
        error.WrongZlibChecksum => "incorrect data check",
        error.BadZlibHeader => "incorrect header check",
        error.InvalidBlockType => "invalid block type",
        error.WrongStoredBlockNlen => "invalid stored block lengths",
        error.InvalidMatch => "invalid distance too far back",
        error.InvalidCode => "invalid literal/length code",
        error.InvalidDynamicBlockHeader => "invalid code lengths set",
        error.OversubscribedHuffmanTree, error.IncompleteHuffmanTree => "invalid literal/lengths set",
        error.MissingEndOfBlockCode => "invalid code -- missing end-of-block",
        else => "data error",
    };
}

/// The DEFLATE decoder, driven a limited number of bytes at a time into a
/// writer that keeps its history. The decoder reads the zlib stream's
/// checksum without checking it, so this checks it.
pub const Inflater = struct {
    decompress: flate.Decompress,
    adler: std.hash.Adler32 = .{},

    /// What one call to `pump` came to.
    pub const Step = enum {
        /// It wrote what was asked, or some of it.
        more,
        /// The stream ended, checksum and all.
        end,
        /// The stored bytes ran out first.
        input_ended,
        /// The stream is corrupt; `decompress.err` says how.
        data_error,
        /// The stored bytes could not be read.
        read_failed,
        /// The writer failed.
        write_failed,
    };

    /// A decoder reading the zlib stream `in` holds.
    pub fn init(in: *Io.Reader) Inflater {
        // No buffer of its own: it writes straight into the writer it is
        // given, and looks back into that writer's buffer for matches.
        return .{ .decompress = .init(in, .zlib, &.{}) };
    }

    /// Inflates up to `limit` bytes more into `w`, whose buffer must keep
    /// the 32 KiB of history the decoder looks back into.
    pub fn pump(inf: *Inflater, w: *Io.Writer, limit: usize) Step {
        std.debug.assert(limit != 0);
        // No more than the history the writer keeps, so that everything
        // this call writes is still in the writer's buffer afterwards to be
        // added to the checksum.
        const n = @min(limit, flate.history_len);
        // The decoder is handed a writer with room for exactly `n` bytes
        // more, after the history it may copy matches from, rather than
        // `w` itself. Its stored blocks are copied by how much the call
        // was asked for in all, not by how much of that is left (see
        // `stored_after_output` below), so a stored block reached after
        // the same call has already written something would otherwise
        // come to more than `n`, and `Io.Reader.stream` asserts it never
        // does. Room for no more than `n` holds the copy to what is left.
        w.rebase(flate.history_len, n) catch return .write_failed;
        var bounded: Io.Writer = .fixed(w.buffer[0 .. w.end + n]);
        bounded.end = w.end;
        const step: Step = if (inf.decompress.reader.stream(&bounded, .limited(n))) |_| .more else |err| switch (err) {
            error.EndOfStream => .end,
            // Only the decoder writing more than it was asked for fills
            // the bounded writer, which is no stream's doing.
            error.WriteFailed => .data_error,
            error.ReadFailed => if (inf.decompress.err) |e| switch (e) {
                error.EndOfStream => .input_ended,
                error.ReadFailed => .read_failed,
                else => .data_error,
            } else .data_error,
        };
        inf.adler.update(bounded.buffered()[w.end..]);
        w.end = bounded.end;
        if (step != .more) return step;
        if (inf.decompress.state == .end) {
            if (inf.adler.adler != inf.decompress.container_metadata.zlib.adler) {
                inf.decompress.err = error.WrongZlibChecksum;
                return .data_error;
            }
            return .end;
        }
        return .more;
    }

    /// zlib's whole message for the corruption a `data_error` step found,
    /// as git prints it after "inflate: ".
    pub fn problem(inf: *const Inflater) []const u8 {
        return dataErrorMessage(inf.decompress.err);
    }
};

/// A writer passing what is written on to another, a window at a time,
/// keeping the last 32 KiB to hand for the decoder to copy matches from.
pub const WindowWriter = struct {
    writer: Io.Writer,
    out: *Io.Writer,
    /// How many of the first bytes written to drop rather than pass on:
    /// the object's header.
    skip: usize = 0,

    /// A writer keeping `window`, at least `flate.max_window_len` bytes, as
    /// its history and passing everything on to `out`.
    pub fn init(window: []u8, out: *Io.Writer) WindowWriter {
        std.debug.assert(window.len >= flate.max_window_len);
        return .{ .writer = .{ .buffer = window, .vtable = &.{ .drain = drain } }, .out = out };
    }

    fn emit(ww: *WindowWriter, bytes: []const u8) Io.Writer.Error!void {
        const dropped = @min(ww.skip, bytes.len);
        ww.skip -= dropped;
        try ww.out.writeAll(bytes[dropped..]);
    }

    fn drain(w: *Io.Writer, data: []const []const u8, splat: usize) Io.Writer.Error!usize {
        const ww: *WindowWriter = @alignCast(@fieldParentPtr("writer", w));
        try ww.emit(w.buffer[0..w.end]);
        w.end = 0;
        var n: usize = 0;
        for (data[0 .. data.len - 1]) |bytes| {
            try ww.emit(bytes);
            n += bytes.len;
        }
        for (0..splat) |_| try ww.emit(data[data.len - 1]);
        return n + data[data.len - 1].len * splat;
    }

    /// Passes on what is still held in the window.
    pub fn finish(ww: *WindowWriter) Io.Writer.Error!void {
        try ww.emit(ww.writer.buffer[0..ww.writer.end]);
        ww.writer.end = 0;
    }
};

/// The most bytes a DEFLATE stream can inflate to for each byte of it.
/// RFC 1951's longest match is 258 bytes, and a length and a distance
/// take at least one bit each when their Huffman codes are one bit long,
/// so two bits make at most 258 bytes: 1032 to the byte. Stored blocks
/// are one to one, and the zlib header and checksum only add bytes.
pub const max_inflate_ratio = 1032;

/// The most bytes the zlib stream at the start of `compressed_len` bytes
/// could inflate to, however it was made. A size claimed beyond it cannot
/// be met by those bytes, so it is refused before anything that large is
/// allocated: a header is a handful of bytes, and must not be able to ask
/// for more memory than the stream behind it could fill.
pub fn maxInflatedLen(compressed_len: usize) u64 {
    return std.math.mul(u64, compressed_len, max_inflate_ratio) catch std.math.maxInt(u64);
}

/// Why a zlib stream that should have inflated to a known number of bytes
/// did not:
///
/// - `InflateFailed`: the stream is corrupt; the detail is zlib's words.
/// - `TruncatedStream`: the bytes ran out before the stream ended.
/// - `StreamTooLong`: it inflates to more bytes than expected.
/// - `StreamTooShort`: it ends having inflated to fewer.
pub const InflateError = error{ InflateFailed, TruncatedStream, StreamTooLong, StreamTooShort };

/// Inflates the zlib stream at the start of `compressed`, which must
/// inflate to exactly `size` bytes, into `buffer[0..size]`. `buffer` must
/// have room for one byte more, which is how a stream going on past `size`
/// is caught. Whatever follows the stream in `compressed` is ignored, as
/// the next entry in a pack follows the one before it. On
/// `InflateFailed`, `detail` is set to zlib's description.
pub fn inflateExact(compressed: []const u8, buffer: []u8, size: usize, detail: *[]const u8) InflateError!void {
    std.debug.assert(buffer.len > size);
    var in: Io.Reader = .fixed(compressed);
    if (headerProblem(&in) catch unreachable) |problem| {
        detail.* = problem;
        return error.InflateFailed;
    }
    var w: Io.Writer = .fixed(buffer[0 .. size + 1]);
    var inflater: Inflater = .init(&in);
    while (w.end <= size) switch (inflater.pump(&w, size + 1 - w.end)) {
        .more => {},
        .end => return if (w.end == size) {} else if (w.end > size) error.StreamTooLong else error.StreamTooShort,
        .input_ended => return error.TruncatedStream,
        // A fixed reader cannot fail, so `read_failed` is not expected
        // either, and the fixed writer fails only when the stream would go
        // on past its end.
        .data_error, .read_failed => {
            detail.* = inflater.problem();
            return error.InflateFailed;
        },
        .write_failed => return error.StreamTooLong,
    };
    return error.StreamTooLong;
}

/// `inflateExact`, returning how many bytes of `compressed` the stream
/// took up, its checksum included: for streams laid end to end with
/// nothing to say where each ends, as a reftable's log blocks are. The
/// decoder takes only the bytes it uses from its input, so what the
/// reader has given up is the stream's length.
pub fn inflateExactConsumed(compressed: []const u8, buffer: []u8, size: usize, detail: *[]const u8) InflateError!usize {
    std.debug.assert(buffer.len > size);
    var in: Io.Reader = .fixed(compressed);
    if (headerProblem(&in) catch unreachable) |problem| {
        detail.* = problem;
        return error.InflateFailed;
    }
    var w: Io.Writer = .fixed(buffer[0 .. size + 1]);
    var inflater: Inflater = .init(&in);
    while (w.end <= size) switch (inflater.pump(&w, size + 1 - w.end)) {
        .more => {},
        .end => return if (w.end == size) in.seek else if (w.end > size) error.StreamTooLong else error.StreamTooShort,
        .input_ended => return error.TruncatedStream,
        .data_error, .read_failed => {
            detail.* = inflater.problem();
            return error.InflateFailed;
        },
        .write_failed => return error.StreamTooLong,
    };
    return error.StreamTooLong;
}

/// Inflates the start of the zlib stream at the start of `compressed` into
/// `buffer`, stopping when it is full or the stream ends, and returns how
/// many bytes that was: for reading what a stream begins with, such as a
/// delta's sizes, without inflating the rest. A stream that runs out
/// before either is `TruncatedStream`; on `InflateFailed`, `detail` is set
/// to zlib's description.
pub fn inflatePrefix(compressed: []const u8, buffer: []u8, detail: *[]const u8) error{ InflateFailed, TruncatedStream }!usize {
    std.debug.assert(buffer.len > 0);
    var in: Io.Reader = .fixed(compressed);
    if (headerProblem(&in) catch unreachable) |problem| {
        detail.* = problem;
        return error.InflateFailed;
    }
    var w: Io.Writer = .fixed(buffer);
    var inflater: Inflater = .init(&in);
    while (w.end < buffer.len) switch (inflater.pump(&w, buffer.len - w.end)) {
        .more => {},
        .end => break,
        .input_ended => return error.TruncatedStream,
        // Neither a fixed reader nor a fixed writer with room for what is
        // asked can fail, so those two are not expected either.
        .data_error, .read_failed, .write_failed => {
            detail.* = inflater.problem();
            return error.InflateFailed;
        },
    };
    return w.end;
}

/// Writes what the zlib stream at the start of `compressed` inflates to
/// to `out` as it goes, in memory bounded whatever its size, and returns
/// how many bytes that was. `gpa` provides the window the decoder looks
/// back into. If the stream fails part way, everything inflated before
/// the failure has been written by the time the error is returned.
pub fn inflateStream(
    gpa: Allocator,
    compressed: []const u8,
    out: *Io.Writer,
    detail: *[]const u8,
) (Allocator.Error || Io.Writer.Error || error{ InflateFailed, TruncatedStream })!u64 {
    var in: Io.Reader = .fixed(compressed);
    if (headerProblem(&in) catch unreachable) |problem| {
        detail.* = problem;
        return error.InflateFailed;
    }
    const window = try gpa.alloc(u8, flate.max_window_len);
    defer gpa.free(window);
    var counter: Counter = .{ .out = out };
    var ww: WindowWriter = .init(window, &counter.writer);
    var inflater: Inflater = .init(&in);
    while (true) switch (inflater.pump(&ww.writer, std.math.maxInt(usize))) {
        .more => {},
        .end => break,
        .write_failed => return error.WriteFailed,
        else => |step| {
            ww.finish() catch return error.WriteFailed;
            if (step == .input_ended) return error.TruncatedStream;
            detail.* = inflater.problem();
            return error.InflateFailed;
        },
    };
    try ww.finish();
    return counter.count;
}

/// A writer that passes everything on to another, counting it.
const Counter = struct {
    out: *Io.Writer,
    count: u64 = 0,
    writer: Io.Writer = .{ .buffer = &.{}, .vtable = &.{ .drain = drain } },

    fn drain(w: *Io.Writer, data: []const []const u8, splat: usize) Io.Writer.Error!usize {
        const c: *Counter = @alignCast(@fieldParentPtr("writer", w));
        var n: usize = 0;
        for (data[0 .. data.len - 1]) |bytes| {
            try c.out.writeAll(bytes);
            n += bytes.len;
        }
        for (0..splat) |_| try c.out.writeAll(data[data.len - 1]);
        n += data[data.len - 1].len * splat;
        c.count += n;
        return n;
    }
};

/// zlib's whole message for a corrupt stream, as git prints it after
/// "inflate: ".
pub fn dataErrorMessage(err: ?flate.Decompress.Error) []const u8 {
    return switch (err orelse return "data stream error (" ++ comptime message(null) ++ ")") {
        inline else => |e| "data stream error (" ++ comptime message(e) ++ ")",
    };
}

const testing = std.testing;

test "a stream inflates to exactly its size, or says how it did not" {
    const gpa = testing.allocator;
    const fixture = @import("testing/fixture.zig");
    const good = try fixture.zlibStored(gpa, "hello");
    defer gpa.free(good);
    var buffer: [16]u8 = undefined;
    var detail: []const u8 = "";
    try inflateExact(good, &buffer, 5, &detail);
    try testing.expectEqualStrings("hello", buffer[0..5]);
    try testing.expectError(error.StreamTooShort, inflateExact(good, &buffer, 6, &detail));
    try testing.expectError(error.StreamTooLong, inflateExact(good, &buffer, 4, &detail));
    try testing.expectError(error.TruncatedStream, inflateExact(good[0 .. good.len - 1], &buffer, 5, &detail));
    const bad = try gpa.dupe(u8, good);
    defer gpa.free(bad);
    bad[bad.len - 1] ^= 1;
    try testing.expectError(error.InflateFailed, inflateExact(bad, &buffer, 5, &detail));
    try testing.expectEqualStrings("data stream error (incorrect data check)", detail);
    try testing.expectEqual(3, try inflatePrefix(good, buffer[0..3], &detail));
    try testing.expectEqual(5, try inflatePrefix(good, &buffer, &detail));
    var aw: Io.Writer.Allocating = .init(gpa);
    defer aw.deinit();
    try testing.expectEqual(5, try inflateStream(gpa, good, &aw.writer, &detail));
    try testing.expectEqualStrings("hello", aw.written());
}

/// "abc" as a zlib stream of two blocks: "a" in a block of fixed codes,
/// then "bc" in a final stored block. Asked for two bytes at a time, the
/// standard library's decoder (Zig 0.17.0) writes the "a", reaches the
/// stored block with one byte of the two left, and copies as much of the
/// stored block as the whole two allow: three bytes for a call asked for
/// two, which `Io.Reader.stream` asserts cannot happen. Every stored block
/// that begins part way through a call does it, as zlib's encoder makes
/// one wherever data does not compress: a PDF in Ghostty's history, of
/// dynamic blocks and stored ones by turns, was the first found.
const stored_after_output = @embedFile("fuzz/zlib-stored-after-output");

test "a stored block after other output in the same call is copied no further than asked" {
    var buffer: [8]u8 = undefined;
    var detail: []const u8 = "";
    try testing.expectEqual(2, try inflatePrefix(stored_after_output, buffer[0..2], &detail));
    try testing.expectEqualStrings("ab", buffer[0..2]);
    try inflateExact(stored_after_output, &buffer, 3, &detail);
    try testing.expectEqualStrings("abc", buffer[0..3]);

    // One step at a time, each a byte short of the stored block.
    var in: Io.Reader = .fixed(stored_after_output);
    try testing.expectEqual(null, try headerProblem(&in));
    var w: Io.Writer = .fixed(&buffer);
    var inflater: Inflater = .init(&in);
    try testing.expectEqual(.more, inflater.pump(&w, 2));
    try testing.expectEqualStrings("ab", w.buffered());
    try testing.expectEqual(.more, inflater.pump(&w, 2));
    try testing.expectEqualStrings("abc", w.buffered());
    // The stored block has ended there; the checksum is next.
    try testing.expectEqual(.end, inflater.pump(&w, 2));
    try testing.expectEqualStrings("abc", w.buffered());
}

test "an object bigger than one step, its stored block begun part way through one, inflates whole" {
    // As the PDF was: more than the 32 KiB a step inflates, so that the
    // step that reaches the stored block has written something already.
    const gpa = testing.allocator;
    const fixture = @import("testing/fixture.zig");
    const raw = try gpa.alloc(u8, 40_001);
    defer gpa.free(raw);
    for (raw, 0..) |*b, i| b.* = @truncate(i *% 2654435761 >> 13);
    const stream = try fixture.zlibLiteralThenStored(gpa, raw);
    defer gpa.free(stream);
    const buffer = try gpa.alloc(u8, raw.len + 1);
    defer gpa.free(buffer);
    var detail: []const u8 = "";
    try inflateExact(stream, buffer, raw.len, &detail);
    try testing.expectEqualSlices(u8, raw, buffer[0..raw.len]);
    try testing.expectEqual(stream.len, try inflateExactConsumed(stream, buffer, raw.len, &detail));
    var aw: Io.Writer.Allocating = .init(gpa);
    defer aw.deinit();
    try testing.expectEqual(raw.len, try inflateStream(gpa, stream, &aw.writer, &detail));
    try testing.expectEqualSlices(u8, raw, aw.written());
}

fn fuzzOne(_: void, smith: *testing.Smith) anyerror!void {
    const gpa = testing.allocator;
    var input: [4096]u8 = undefined;
    const compressed = fuzz_input.bytes(smith, &input);
    // The whole stream, written out as it goes; then what each of the
    // other ways of inflating it makes must agree with that.
    var aw: Io.Writer.Allocating = .init(gpa);
    defer aw.deinit();
    var detail: []const u8 = "";
    const whole: ?[]const u8 = if (inflateStream(gpa, compressed, &aw.writer, &detail)) |_| aw.written() else |_| null;
    var buffer: [1 << 16]u8 = undefined;
    for ([_]usize{ 1, 2, 3, 20, 255 }) |n| {
        const got = inflatePrefix(compressed, buffer[0..n], &detail) catch continue;
        if (whole) |all| try testing.expectEqualSlices(u8, all[0..@min(all.len, n)], buffer[0..got]);
    }
    if (whole) |all| if (all.len < buffer.len) {
        try inflateExact(compressed, &buffer, all.len, &detail);
        try testing.expectEqualSlices(u8, all, buffer[0..all.len]);
    };
}

const fuzz_input = @import("testing/fuzz.zig");

test "fuzz: zlib streams" {
    const gpa = testing.allocator;
    const fixture = @import("testing/fixture.zig");
    const stored = try fixture.zlibStored(gpa, "hello");
    defer gpa.free(stored);
    try fuzz_input.fuzz({}, fuzzOne, .{ .corpus = &.{
        stored_after_output,
        @embedFile("fuzz/loose-blob"),
        @embedFile("fuzz/loose-stored-after-output"),
        stored,
    } });
}