From 8767cc4b43f2e5fd650cae0925e8072ced29e30f Mon Sep 17 00:00:00 2001
From: Ryan Liptak <squeek502@hotmail.com>
Date: Tue, 5 Nov 2024 03:53:16 -0800
Subject: [PATCH] cvtres: Add support for defining an external symbol
 (/DEFINE:symbol option of cvtres.exe)

---
 src/cvtres.zig        | 41 +++++++++++++++++++++++++++++++++++++++--
 test/fuzzy_cvtres.zig |  8 ++++++++
 test/utils.zig        | 12 ++++++++++++
 3 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/src/cvtres.zig b/src/cvtres.zig
index b24b25f..f912f21 100644
--- a/src/cvtres.zig
+++ b/src/cvtres.zig
@@ -162,6 +162,8 @@ pub const CoffOptions = struct {
     reproducible: bool = true,
     /// If true, the MEM_WRITE flag will not be set in the .rsrc section header
     read_only: bool = false,
+    /// If non-null, a symbol with this name and storage class EXTERNAL will be added to the symbol table.
+    define_external_symbol: ?[]const u8 = null,
 };
 
 pub fn writeCoff(allocator: Allocator, writer: anytype, resources: []const Resource, options: CoffOptions) !void {
@@ -186,12 +188,13 @@ pub fn writeCoff(allocator: Allocator, writer: anytype, resources: []const Resou
     const flags = std.coff.CoffHeaderFlags{
         .@"32BIT_MACHINE" = 1,
     };
+    const number_of_symbols = 5 + @as(u32, @intCast(resources.len)) + @intFromBool(options.define_external_symbol != null);
     const coff_header = std.coff.CoffHeader{
         .machine = machine_type,
         .number_of_sections = 2,
         .time_date_stamp = @as(u32, @truncate(@as(u64, @bitCast(timestamp)))),
         .pointer_to_symbol_table = pointer_to_symbol_table,
-        .number_of_symbols = 5 + @as(u32, @intCast(resources.len)),
+        .number_of_symbols = number_of_symbols,
         .size_of_optional_header = size_of_optional_header,
         .flags = flags,
     };
@@ -304,8 +307,42 @@ pub fn writeCoff(allocator: Allocator, writer: anytype, resources: []const Resou
         try writeSymbol(writer, resource_symbol);
     }
 
-    const string_table_byte_count = 4;
+    var string_table: std.ArrayListUnmanaged(u8) = .empty;
+    defer string_table.deinit(allocator);
+
+    if (options.define_external_symbol) |external_symbol_name| {
+        const name_bytes: [8]u8 = name_bytes: {
+            if (external_symbol_name.len > 8) {
+                // 4 due to the initial string table byte length field
+                const string_table_offset: u32 = @intCast(4 + string_table.items.len);
+                try string_table.appendSlice(allocator, external_symbol_name);
+                try string_table.append(allocator, 0);
+                var bytes = [_]u8{0} ** 8;
+                std.mem.writeInt(u32, bytes[4..8], string_table_offset, .little);
+                break :name_bytes bytes;
+            } else {
+                var symbol_shortname = [_]u8{0} ** 8;
+                @memcpy(symbol_shortname[0..external_symbol_name.len], external_symbol_name);
+                break :name_bytes symbol_shortname;
+            }
+        };
+
+        try writeSymbol(writer, .{
+            .name = name_bytes,
+            .value = 0,
+            .section_number = .ABSOLUTE,
+            .type = .{
+                .base_type = .NULL,
+                .complex_type = .NULL,
+            },
+            .storage_class = .EXTERNAL,
+            .number_of_aux_symbols = 0,
+        });
+    }
+
+    const string_table_byte_count: u32 = @intCast(4 + string_table.items.len);
     try writer.writeInt(u32, string_table_byte_count, .little);
+    try writer.writeAll(string_table.items);
 }
 
 fn writeSymbol(writer: anytype, symbol: std.coff.Symbol) !void {
diff --git a/test/fuzzy_cvtres.zig b/test/fuzzy_cvtres.zig
index a6828f0..72956de 100644
--- a/test/fuzzy_cvtres.zig
+++ b/test/fuzzy_cvtres.zig
@@ -76,11 +76,19 @@ test "cvtres fuzz" {
             else => unreachable,
         };
 
+        const random_symbol_define: ?[]const u8 = switch (rand.uintLessThan(u8, 3)) {
+            0 => null,
+            1 => "short", // fits within 8 bytes
+            2 => "longerthan8", // needs to go in the string table
+            else => unreachable,
+        };
+
         try utils.expectSameCvtResOutput(allocator, res_buffer.items, .{
             .cwd = tmp.dir,
             .cwd_path = tmp_path,
             .target = random_target,
             .read_only = rand.boolean(),
+            .define_external_symbol = random_symbol_define,
         });
     }
 }
diff --git a/test/utils.zig b/test/utils.zig
index dce4c76..883574d 100644
--- a/test/utils.zig
+++ b/test/utils.zig
@@ -269,6 +269,7 @@ pub const GetCvtResResultOptions = struct {
     output_path: ?[]const u8 = null,
     target: std.coff.MachineType = .X64,
     read_only: bool = false,
+    define_external_symbol: ?[]const u8 = null,
 };
 
 pub const ResinatorCvtResResult = struct {
@@ -298,6 +299,7 @@ pub fn getResinatorCvtResResult(allocator: Allocator, res_source: []const u8, op
     resinator.cvtres.writeCoff(allocator, buf.writer(), resources, .{
         .target = options.target,
         .read_only = options.read_only,
+        .define_external_symbol = options.define_external_symbol,
     }) catch |err| {
         buf.deinit();
         return .{
@@ -354,6 +356,16 @@ pub fn getWin32CvtResResultFromFile(allocator: Allocator, input_path: []const u8
     if (options.read_only) {
         try argv.append("/READONLY");
     }
+    const define_arg: ?[]const u8 = define_arg: {
+        if (options.define_external_symbol) |symbol_name| {
+            const define_arg = try std.fmt.allocPrint(allocator, "/DEFINE:{s}", .{symbol_name});
+            errdefer allocator.free(define_arg);
+            try argv.append(define_arg);
+            break :define_arg define_arg;
+        }
+        break :define_arg null;
+    };
+    defer if (define_arg) |v| allocator.free(v);
 
     try argv.append(input_path);