diff --git a/src/main.zig b/src/main.zig index 9087101..5ae3c02 100644 --- a/src/main.zig +++ b/src/main.zig @@ -398,6 +398,28 @@ fn callLoadFileListWithIncludes(toc_path: [*:0]const u8, md5ctx: *[88]u8, error_ ); } +// ============================================================================= +// Hook: GameEngine_MainInitialize (0x46a400) +// __stdcall(void) — prologue: 55 8B EC 83 EC 28 = 6 bytes, no fixups +// +// Called once from InitializeAllSubsystems during the init callback inside +// WinMain. Fires after Lua env, UI frames, event tables, and DB tables are +// set up, but before the event loop starts frame callbacks (i.e. before the +// login screen renders and screenshots become possible). +// +// We install the screenshot hook here to guarantee it runs AFTER all +// DLL_PROCESS_ATTACH hooks (including UnitXP's CTgaFile::Write hook), +// making us the outermost detour in the hook chain. +// ============================================================================= + +var engine_init_hook: hook.Hook = .{}; + +fn engineInitDetour() callconv(sc) void { + const orig = engine_init_hook.getTrampoline(*const fn () callconv(sc) void); + orig(); + screenshot.installHook(); +} + // ============================================================================= // Hook: CGGameUI_Shutdown (0x490BD0) // Prologue: 56 E8 7A 83 FC FF = 6 bytes, fixup at offset 1 @@ -433,8 +455,9 @@ fn install() void { load_addons_hook.activate(@intFromPtr(thunk)); } - // 5. Screenshot hook — async PNG capture replacing TGA write - screenshot.installHook(); + // 5. GameEngine_MainInitialize — install screenshot hook after all DLLs load + // Fires once inside WinMain, before the login screen renders. + _ = engine_init_hook.install(0x46a400, 6, @intFromPtr(&engineInitDetour), &.{}); // 6. CGGameUI_Shutdown — cleanup _ = shutdown_hook.install(0x490BD0, 6, @intFromPtr(&shutdownDetour), &.{1}); @@ -442,6 +465,7 @@ fn install() void { fn uninstall() void { shutdown_hook.remove(); + engine_init_hook.remove(); screenshot.removeHook(); load_addons_hook.remove(); lsf_hook.remove(); diff --git a/src/png.zig b/src/png.zig index 2cfb221..7fead2b 100644 --- a/src/png.zig +++ b/src/png.zig @@ -64,7 +64,7 @@ pub fn encode( if (@intFromEnum(level) == 0) { writeIdatStore(&s, pixels, w, h); } else { - writeIdatFixedHuffman(&s, pixels, w, h); + writeIdatFixedHuffman(&s, pixels, w, h, @intFromEnum(level)); } // IEND @@ -204,11 +204,11 @@ fn writeIdatStore(s: anytype, pixels: [*]const u8, w: u32, h: u32) void { } // ============================================================================= -// Fixed Huffman (levels 1-9) — RFC 1951 §3.2.6 fixed codes, literals only +// Fixed Huffman + LZ77 — RFC 1951 §3.2.6 fixed codes with back-references // -// Each byte is Huffman-coded using the fixed table. No LZ77 matching. -// This gives ~10-20% compression on typical screenshot data with zero -// runtime state beyond a small bit buffer. +// Sub filter makes adjacent pixel differences small (often zero in flat areas). +// LZ77 with hash-table matching finds repeated byte sequences within a 32KB +// window and encodes them as length-distance pairs. // ============================================================================= const BitBuf = struct { @@ -236,13 +236,14 @@ const BitBuf = struct { } }; -/// RFC 1951 fixed Huffman: encode a literal byte (0-255) or end-of-block (256). -fn fixedLiteral(bb: *BitBuf, s: anytype, val: u16) void { - // RFC 1951 §3.2.6 fixed Huffman code table: - // 0-143: 8 bits, codes 00110000-10111111 - // 144-255: 9 bits, codes 110010000-111111111 - // 256-279: 7 bits, codes 0000000-0010111 - // 280-287: 8 bits, codes 11000000-11000111 +fn bitReverse(comptime T: type, val: T, n: u5) u32 { + const full = @bitReverse(val); + const shift: u5 = @intCast(@typeInfo(T).int.bits - @as(u8, n)); + return @as(u32, full) >> shift; +} + +/// RFC 1951 fixed Huffman: encode a literal/length code (0-285). +fn fixedCode(bb: *BitBuf, s: anytype, val: u16) void { if (val <= 143) { const code: u9 = @as(u9, @intCast(val)) + 0x30; bb.write(s, bitReverse(u9, code, 8), 8); @@ -258,50 +259,105 @@ fn fixedLiteral(bb: *BitBuf, s: anytype, val: u16) void { } } -fn bitReverse(comptime T: type, val: T, n: u5) u32 { - const full = @bitReverse(val); - const shift: u5 = @intCast(@typeInfo(T).int.bits - @as(u8, n)); - return @as(u32, full) >> shift; +// RFC 1951 length/distance encoding tables (from stb_image_write.h) +const length_base = [29]u16{ 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 15, 17, 19, 23, 27, 31, 35, 43, 51, 59, 67, 83, 99, 115, 131, 163, 195, 227, 258 }; +const length_extra = [29]u5{ 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 0 }; +const dist_base = [30]u16{ 1, 2, 3, 4, 5, 7, 9, 13, 17, 25, 33, 49, 65, 97, 129, 193, 257, 385, 513, 769, 1025, 1537, 2049, 3073, 4097, 6145, 8193, 12289, 16385, 24577 }; +const dist_extra = [30]u5{ 0, 0, 0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8, 9, 9, 10, 10, 11, 11, 12, 12, 13, 13 }; + +inline fn zhashFn(d: *const [3]u8) u32 { + var h: u32 = @as(u32, d[0]) +% (@as(u32, d[1]) << 8) +% (@as(u32, d[2]) << 16); + h ^= h *% 8; + h +%= h >> 5; + h ^= h *% 16; + h +%= h >> 17; + h ^= h *% (1 << 25); + h +%= h >> 6; + return h; } -fn writeIdatFixedHuffman(s: anytype, pixels: [*]const u8, w: u32, h: u32) void { - // We can't precompute IDAT payload size for Huffman, so we buffer the - // entire deflate stream, then write it as one IDAT chunk. - // For a 1024x768 screenshot, fixed Huffman with only literals produces - // roughly 8-9 bits per byte ≈ same size or slightly larger than raw. - // But the Sub filter makes most bytes small, yielding good compression. +fn countMatch(data: []const u8, a: u32, b: u32) u32 { + const max_len: u32 = @min(@as(u32, @intCast(data.len)) - b, 258); + var i: u32 = 0; + while (i < max_len) : (i += 1) { + if (data[a + i] != data[b + i]) break; + } + return i; +} - // Allocate output buffer: worst case ~9 bits/byte * raw_size / 8 + overhead +/// Emit a deflate length-distance pair using fixed Huffman codes. +fn emitMatch(bb: *BitBuf, s: anytype, length: u32, distance: u32) void { + // Length code + var li: usize = 0; + while (li + 1 < length_base.len and length >= length_base[li + 1]) : (li += 1) {} + fixedCode(bb, s, @intCast(li + 257)); + if (length_extra[li] > 0) bb.write(s, length - length_base[li], length_extra[li]); + // Distance code: 5-bit reversed + extra bits + var di: usize = 0; + while (di + 1 < dist_base.len and distance >= dist_base[di + 1]) : (di += 1) {} + bb.write(s, bitReverse(u5, @as(u5, @intCast(di)), 5), 5); + if (dist_extra[di] > 0) bb.write(s, distance - dist_base[di], dist_extra[di]); +} + +fn writeIdatFixedHuffman(s: anytype, pixels: [*]const u8, w: u32, h: u32, level: u4) void { const row_bytes: u32 = 1 + w * 3; const raw_size: u32 = h * row_bytes; - // Worst case: 9 bits per byte + block headers + zlib overhead - const max_out: u32 = (raw_size / 8) * 9 + raw_size / 8 + 1024; + + // Pre-filter all pixel data (Sub filter) into contiguous buffer + const filtered = std.heap.page_allocator.alloc(u8, raw_size) catch { + writeIdatStore(s, pixels, w, h); + return; + }; + defer std.heap.page_allocator.free(filtered); + { + var pos: u32 = 0; + var y: u32 = 0; + while (y < h) : (y += 1) { + filtered[pos] = 1; // Sub filter type byte + pos += 1; + const row = y * w * 3; + var x: u32 = 0; + while (x < w * 3) : (x += 1) { + const raw = pixels[row + x]; + filtered[pos] = if (x >= 3) raw -% pixels[row + x - 3] else raw; + pos += 1; + } + } + } + + // Adler-32 over entire filtered buffer + var adler_a: u32 = 1; + var adler_b: u32 = 0; + adlerUpdate(&adler_a, &adler_b, filtered[0..raw_size]); + const adler: u32 = (adler_b << 16) | adler_a; + + // Output buffer (worst case: ~9 bits per byte for fixed Huffman literals) + const max_out: u32 = raw_size + raw_size / 4 + 1024; const out_buf = std.heap.page_allocator.alloc(u8, max_out) catch { - // Fall back to store blocks writeIdatStore(s, pixels, w, h); return; }; defer std.heap.page_allocator.free(out_buf); - // Compress into buffer + // Hash chains for LZ77 matching — chain depth = 2 * level (stb approach) + const ZHASH: u32 = 16384; + const chain_depth: u32 = @as(u32, level) * 2; + const chain_mem = std.heap.page_allocator.alloc(u32, ZHASH * chain_depth) catch { + writeIdatStore(s, pixels, w, h); + return; + }; + defer std.heap.page_allocator.free(chain_mem); + const chain_count = std.heap.page_allocator.alloc(u8, ZHASH) catch { + writeIdatStore(s, pixels, w, h); + return; + }; + defer std.heap.page_allocator.free(chain_count); + @memset(chain_count, 0); + var out_pos: u32 = 0; - - // Zlib header - out_buf[0] = 0x78; - out_buf[1] = 0x9C; // default compression - out_pos = 2; - - var adler_a: u32 = 1; - var adler_b: u32 = 0; - - // Single fixed-Huffman block (BFINAL=1, BTYPE=01) - var bb: BitBuf = .{}; - - // Pack bits into out_buf via a mini stream const OutStream = struct { buf: []u8, pos: *u32, - // Dummy fields matching the CrcAdler interface fn writeCrcAdler(self: *@This(), data: []const u8) void { for (data) |byte| { if (self.pos.* < self.buf.len) { @@ -313,38 +369,87 @@ fn writeIdatFixedHuffman(s: anytype, pixels: [*]const u8, w: u32, h: u32) void { }; var out_stream = OutStream{ .buf = out_buf, .pos = &out_pos }; + // Zlib header + out_buf[0] = 0x78; + out_buf[1] = 0x9C; + out_pos = 2; + + var bb: BitBuf = .{}; // BFINAL=1, BTYPE=01 (fixed Huffman) bb.write(&out_stream, 0b011, 3); - // Encode each scanline with Sub filter - var y: u32 = 0; - while (y < h) : (y += 1) { - const row_start = y * w * 3; + // LZ77 + fixed Huffman encoding + const WINDOW: u32 = 32768; + var i: u32 = 0; + while (i + 2 < raw_size) { + const bucket = zhashFn(filtered[i..][0..3]) & (ZHASH - 1); + const base = bucket * chain_depth; - // Filter byte: 1 = Sub - const filter_byte: u8 = 1; - adlerUpdate(&adler_a, &adler_b, &[1]u8{filter_byte}); - fixedLiteral(&bb, &out_stream, filter_byte); - - // First pixel: Sub filter with no left neighbor = raw bytes - var x: u32 = 0; - while (x < w * 3) : (x += 1) { - const raw = pixels[row_start + x]; - const filtered: u8 = if (x >= 3) - raw -% pixels[row_start + x - 3] - else - raw; - adlerUpdate(&adler_a, &adler_b, &[1]u8{filtered}); - fixedLiteral(&bb, &out_stream, filtered); + // Search chain for best match (prefer closest of equal length via >=) + var best_len: u32 = 3; + var best_dist: u32 = 0; + { + var j: u32 = 0; + while (j < chain_count[bucket]) : (j += 1) { + const prev = chain_mem[base + j]; + if (i -% prev <= WINDOW) { + const ml = countMatch(filtered[0..raw_size], prev, i); + if (ml >= best_len) { + best_len = ml; + best_dist = i - prev; + } + } + } } + + // Add current position to chain; prune oldest half when full + { + var cnt = @as(u32, chain_count[bucket]); + if (cnt >= chain_depth) { + const keep = chain_depth / 2; + const src = base + chain_depth - keep; + @memcpy(chain_mem[base..][0..keep], chain_mem[src..][0..keep]); + cnt = keep; + } + chain_mem[base + cnt] = i; + chain_count[bucket] = @intCast(cnt + 1); + } + + if (best_dist > 0) { + // Lazy matching: check if next position beats current match + if (i + 3 < raw_size) { + const bucket2 = zhashFn(filtered[i + 1 ..][0..3]) & (ZHASH - 1); + const base2 = bucket2 * chain_depth; + var j: u32 = 0; + while (j < chain_count[bucket2]) : (j += 1) { + const prev2 = chain_mem[base2 + j]; + if ((i + 1) -% prev2 <= WINDOW) { + if (countMatch(filtered[0..raw_size], prev2, i + 1) > best_len) { + best_dist = 0; // cancel — next position is better + break; + } + } + } + } + if (best_dist > 0) { + emitMatch(&bb, &out_stream, best_len, best_dist); + i += best_len; + continue; + } + } + + fixedCode(&bb, &out_stream, filtered[i]); + i += 1; + } + // Remaining <3 bytes as literals + while (i < raw_size) : (i += 1) { + fixedCode(&bb, &out_stream, filtered[i]); } - // End of block marker (256) - fixedLiteral(&bb, &out_stream, 256); + fixedCode(&bb, &out_stream, 256); // End of block bb.flush(&out_stream); - // Adler-32 (big-endian, NOT bit-packed — appended as raw bytes after deflate) - const adler = (adler_b << 16) | adler_a; + // Append Adler-32 (big-endian) if (out_pos + 4 <= out_buf.len) { out_buf[out_pos] = @truncate(adler >> 24); out_buf[out_pos + 1] = @truncate(adler >> 16); @@ -353,7 +458,12 @@ fn writeIdatFixedHuffman(s: anytype, pixels: [*]const u8, w: u32, h: u32) void { out_pos += 4; } - // Write as single IDAT chunk + // Fallback to store if compressed is larger + if (out_pos >= raw_size) { + writeIdatStore(s, pixels, w, h); + return; + } + s.writeChunk("IDAT", out_buf[0..out_pos]); } diff --git a/src/screenshot.zig b/src/screenshot.zig index 0ddf2a9..131a38c 100644 --- a/src/screenshot.zig +++ b/src/screenshot.zig @@ -52,7 +52,8 @@ var compression_level: i32 = 6; // user-facing 0–9, kept for Lua interface var tga_hook: hook.Hook = .{}; var screenshot_dir: [260]u8 = undefined; var screenshot_dir_len: usize = 0; -var screenshot_counter: u32 = 1; +var screenshot_counter: u8 = 0; +var last_screenshot_time: u64 = 0; // packed YMDHMS — resets counter on new second // ============================================================================= // Ring buffer queue (max 8 pending screenshots) @@ -65,6 +66,7 @@ const PendingScreenshot = struct { width: u16, height: u16, size: u32, + level: png.Level, }; var queue: [MAX_PENDING]PendingScreenshot = undefined; @@ -160,7 +162,7 @@ fn tgaWriteDetour(self: u32, _edx: u32, filename: u32) callconv(.c) i32 { mutex.lock(); defer mutex.unlock(); - if (!enqueue(.{ .buffer = buffer.ptr, .width = width, .height = height, .size = size })) { + if (!enqueue(.{ .buffer = buffer.ptr, .width = width, .height = height, .size = size, .level = png.mapLevel(compression_level) })) { std.heap.page_allocator.free(buffer); return callOriginal(self, filename); } @@ -213,12 +215,24 @@ fn processScreenshot(shot: PendingScreenshot) void { shot.buffer[off + 2] = tmp; } - // Generate filename: {dir}WoWScrnShot_MMDDYY_HHMMSS_N.png + // Generate filename: {dir}WoWScrnShot_MMDDYY_HHMMSS_X.png (X = 0–F hex) var st: SYSTEMTIME = undefined; GetLocalTime(&st); + // Pack timestamp into a single comparable value — reset counter on new second + const now: u64 = @as(u64, st.wYear) << 32 | @as(u64, st.wMonth) << 24 | + @as(u64, st.wDay) << 16 | @as(u64, st.wHour) << 10 | + @as(u64, st.wMinute) << 4 | @as(u64, st.wSecond); + if (now != last_screenshot_time) { + screenshot_counter = 0; + last_screenshot_time = now; + } + + const suffix: u8 = if (screenshot_counter < 16) "0123456789ABCDEF"[screenshot_counter] else return; + screenshot_counter += 1; + var name_buf: [260]u8 = undefined; - const name_slice = std.fmt.bufPrint(&name_buf, "{s}WoWScrnShot_{:0>2}{:0>2}{:0>2}_{:0>2}{:0>2}{:0>2}_{}.png", .{ + const name_slice = std.fmt.bufPrint(&name_buf, "{s}WoWScrnShot_{:0>2}{:0>2}{:0>2}_{:0>2}{:0>2}{:0>2}_{c}.png", .{ screenshot_dir[0..screenshot_dir_len], st.wMonth, st.wDay, @@ -226,18 +240,14 @@ fn processScreenshot(shot: PendingScreenshot) void { st.wHour, st.wMinute, st.wSecond, - screenshot_counter, + suffix, }) catch return; - screenshot_counter += 1; - if (screenshot_counter > 999) screenshot_counter = 1; - // Null-terminate for CreateFileA if (name_slice.len >= name_buf.len) return; name_buf[name_slice.len] = 0; - const level = png.mapLevel(compression_level); - writePng(@ptrCast(name_slice.ptr), shot.buffer, shot.width, shot.height, level); + writePng(@ptrCast(name_slice.ptr), shot.buffer, shot.width, shot.height, shot.level); } // ============================================================================= @@ -333,6 +343,12 @@ pub fn installHook() void { // CTgaFile::Write at 0x5a4810 // __thiscall(self, filename) — prologue: 55 8B EC 83 EC 08 = 6 bytes, no fixups // Thunk: fastcall(ECX=self, EDX, stack: filename) → cdecl(self, edx, filename) + // + // Another DLL (UnitXP_SP3) hooks this same address during DLL_PROCESS_ATTACH, + // replacing the prologue with an E9 JMP. Restore the original prologue first + // so prepare() builds a trampoline to the real function rather than chaining + // through UnitXP's detour. + hook.writeProtected(0x5a4810, &.{ 0x55, 0x8B, 0xEC, 0x83, 0xEC, 0x08 }); if (tga_hook.prepare(0x5a4810, 6, &.{})) { const thunk = tga_hook.mem.? + 32; _ = hook.buildFastcallToCdeclThunk(thunk, @intFromPtr(&tgaWriteDetour), 1);