add --image-base support

Based on #6121 by Jay Petacat.

5 files changed, 56 insertions, 19 deletions

diff --git a/src/Compilation.zig b/src/Compilation.zig
index 16a475b2fc..1085e2d779 100644
--- a/src/Compilation.zig
+++ b/src/Compilation.zig
@@ -378,6 +378,7 @@ pub const InitOptions = struct {
     is_compiler_rt_or_libc: bool = false,
     parent_compilation_link_libc: bool = false,
     stack_size_override: ?u64 = null,
+    image_base_override: ?u64 = null,
     self_exe_path: ?[]const u8 = null,
     version: ?std.builtin.Version = null,
     libc_installation: ?*const LibCInstallation = null,
@@ -452,6 +453,7 @@ pub fn create(gpa: *Allocator, options: InitOptions) !*Compilation {
                 options.link_emit_relocs or
                 options.output_mode == .Lib or
                 options.lld_argv.len != 0 or
+                options.image_base_override != null or
                 options.linker_script != null or options.version_script != null)
             {
                 break :blk true;
@@ -772,6 +774,7 @@ pub fn create(gpa: *Allocator, options: InitOptions) !*Compilation {
             .z_nodelete = options.linker_z_nodelete,
             .z_defs = options.linker_z_defs,
             .stack_size_override = options.stack_size_override,
+            .image_base_override = options.image_base_override,
             .linker_script = options.linker_script,
             .version_script = options.version_script,
             .gc_sections = options.linker_gc_sections,
diff --git a/src/link.zig b/src/link.zig
index 3b00684170..139977b3e2 100644
--- a/src/link.zig
+++ b/src/link.zig
@@ -45,6 +45,7 @@ pub const Options = struct {
     program_code_size_hint: u64 = 256 * 1024,
     entry_addr: ?u64 = null,
     stack_size_override: ?u64,
+    image_base_override: ?u64,
     /// Set to `true` to omit debug info.
     strip: bool,
     /// If this is true then this link code is responsible for outputting an object
diff --git a/src/link/Coff.zig b/src/link/Coff.zig
index 3f462582e7..24c3833e0a 100644
--- a/src/link/Coff.zig
+++ b/src/link/Coff.zig
@@ -22,10 +22,10 @@ const minimum_text_block_size = 64 * allocation_padding;
 
 const section_alignment = 4096;
 const file_alignment = 512;
-const image_base = 0x400_000;
+const default_image_base = 0x400_000;
 const section_table_size = 2 * 40;
 comptime {
-    assert(mem.isAligned(image_base, section_alignment));
+    assert(mem.isAligned(default_image_base, section_alignment));
 }
 
 pub const base_tag: link.File.Tag = .coff;
@@ -55,7 +55,7 @@ offset_table: std.ArrayListUnmanaged(u64) = .{},
 /// Free list of offset table indices
 offset_table_free_list: std.ArrayListUnmanaged(u32) = .{},
 
-/// Virtual address of the entry point procedure relative to `image_base`
+/// Virtual address of the entry point procedure relative to image base.
 entry_addr: ?u32 = null,
 
 /// Absolute virtual address of the text section when the executable is loaded in memory.
@@ -183,14 +183,14 @@ pub fn openPath(allocator: *Allocator, sub_path: []const u8, options: link.Optio
 
     self.section_data_offset = mem.alignForwardGeneric(u32, self.section_table_offset + section_table_size, file_alignment);
     const section_data_relative_virtual_address = mem.alignForwardGeneric(u32, self.section_table_offset + section_table_size, section_alignment);
-    self.offset_table_virtual_address = image_base + section_data_relative_virtual_address;
+    self.offset_table_virtual_address = default_image_base + section_data_relative_virtual_address;
     self.offset_table_size = default_offset_table_size;
     self.section_table_offset = section_table_offset;
-    self.text_section_virtual_address = image_base + section_data_relative_virtual_address + section_alignment;
+    self.text_section_virtual_address = default_image_base + section_data_relative_virtual_address + section_alignment;
     self.text_section_size = default_size_of_code;
 
     // Size of file when loaded in memory
-    const size_of_image = mem.alignForwardGeneric(u32, self.text_section_virtual_address - image_base + default_size_of_code, section_alignment);
+    const size_of_image = mem.alignForwardGeneric(u32, self.text_section_virtual_address - default_image_base + default_size_of_code, section_alignment);
 
     mem.writeIntLittle(u16, hdr_data[index..][0..2], optional_header_size);
     index += 2;
@@ -234,11 +234,11 @@ pub fn openPath(allocator: *Allocator, sub_path: []const u8, options: link.Optio
             index += 4;
 
             // Image base address
-            mem.writeIntLittle(u32, hdr_data[index..][0..4], image_base);
+            mem.writeIntLittle(u32, hdr_data[index..][0..4], default_image_base);
             index += 4;
         } else {
             // Image base address
-            mem.writeIntLittle(u64, hdr_data[index..][0..8], image_base);
+            mem.writeIntLittle(u64, hdr_data[index..][0..8], default_image_base);
             index += 8;
         }
 
@@ -328,7 +328,7 @@ pub fn openPath(allocator: *Allocator, sub_path: []const u8, options: link.Optio
         mem.writeIntLittle(u32, hdr_data[index..][0..4], default_offset_table_size);
         index += 4;
         // Virtual address (u32)
-        mem.writeIntLittle(u32, hdr_data[index..][0..4], self.offset_table_virtual_address - image_base);
+        mem.writeIntLittle(u32, hdr_data[index..][0..4], self.offset_table_virtual_address - default_image_base);
         index += 4;
     } else {
         mem.set(u8, hdr_data[index..][0..8], 0);
@@ -354,7 +354,7 @@ pub fn openPath(allocator: *Allocator, sub_path: []const u8, options: link.Optio
         mem.writeIntLittle(u32, hdr_data[index..][0..4], default_size_of_code);
         index += 4;
         // Virtual address (u32)
-        mem.writeIntLittle(u32, hdr_data[index..][0..4], self.text_section_virtual_address - image_base);
+        mem.writeIntLittle(u32, hdr_data[index..][0..4], self.text_section_virtual_address - default_image_base);
         index += 4;
     } else {
         mem.set(u8, hdr_data[index..][0..8], 0);
@@ -601,7 +601,7 @@ fn writeOffsetTableEntry(self: *Coff, index: usize) !void {
 
             // Write .text new virtual address
             self.text_section_virtual_address = self.text_section_virtual_address + va_offset;
-            mem.writeIntLittle(u32, buf[0..4], self.text_section_virtual_address - image_base);
+            mem.writeIntLittle(u32, buf[0..4], self.text_section_virtual_address - default_image_base);
             try self.base.file.?.pwriteAll(buf[0..4], self.section_table_offset + 40 + 12);
 
             // Fix the VAs in the offset table
@@ -716,7 +716,7 @@ pub fn updateDeclExports(self: *Coff, module: *Module, decl: *const Module.Decl,
             }
         }
         if (mem.eql(u8, exp.options.name, "_start")) {
-            self.entry_addr = decl.link.coff.getVAddr(self.*) - image_base;
+            self.entry_addr = decl.link.coff.getVAddr(self.*) - default_image_base;
         } else {
             try module.failed_exports.ensureCapacity(module.gpa, module.failed_exports.items().len + 1);
             module.failed_exports.putAssumeCapacityNoClobber(
@@ -754,7 +754,7 @@ pub fn flushModule(self: *Coff, comp: *Compilation) !void {
     }
 
     if (self.base.options.output_mode == .Exe and self.size_of_image_dirty) {
-        const new_size_of_image = mem.alignForwardGeneric(u32, self.text_section_virtual_address - image_base + self.text_section_size, section_alignment);
+        const new_size_of_image = mem.alignForwardGeneric(u32, self.text_section_virtual_address - default_image_base + self.text_section_size, section_alignment);
         var buf: [4]u8 = undefined;
         mem.writeIntLittle(u32, &buf, new_size_of_image);
         try self.base.file.?.pwriteAll(&buf, self.optional_header_offset + 56);
@@ -832,6 +832,7 @@ fn linkWithLLD(self: *Coff, comp: *Compilation) !void {
         }
         try man.addOptionalFile(module_obj_path);
         man.hash.addOptional(self.base.options.stack_size_override);
+        man.hash.addOptional(self.base.options.image_base_override);
         man.hash.addListOfBytes(self.base.options.extra_lld_args);
         man.hash.addListOfBytes(self.base.options.lib_dirs);
         man.hash.add(self.base.options.is_compiler_rt_or_libc);
@@ -914,6 +915,9 @@ fn linkWithLLD(self: *Coff, comp: *Compilation) !void {
             const stack_size = self.base.options.stack_size_override orelse 16777216;
             try argv.append(try allocPrint(arena, "-STACK:{d}", .{stack_size}));
         }
+        if (self.base.options.image_base_override) |image_base| {
+            try argv.append(try std.fmt.allocPrint(arena, "-BASE:{d}", .{image_base}));
+        }
 
         if (target.cpu.arch == .i386) {
             try argv.append("-MACHINE:X86");
diff --git a/src/link/Elf.zig b/src/link/Elf.zig
index 8e7a02114e..df56bb8bb4 100644
--- a/src/link/Elf.zig
+++ b/src/link/Elf.zig
@@ -1284,6 +1284,7 @@ fn linkWithLLD(self: *Elf, comp: *Compilation) !void {
         // We can skip hashing libc and libc++ components that we are in charge of building from Zig
         // installation sources because they are always a product of the compiler version + target information.
         man.hash.add(stack_size);
+        man.hash.addOptional(self.base.options.image_base_override);
         man.hash.add(gc_sections);
         man.hash.add(self.base.options.eh_frame_hdr);
         man.hash.add(self.base.options.emit_relocs);
@@ -1353,6 +1354,10 @@ fn linkWithLLD(self: *Elf, comp: *Compilation) !void {
         try argv.append(try std.fmt.allocPrint(arena, "stack-size={}", .{stack_size}));
     }
 
+    if (self.base.options.image_base_override) |image_base| {
+        try argv.append(try std.fmt.allocPrint(arena, "--image-base={d}", .{image_base}));
+    }
+
     if (self.base.options.linker_script) |linker_script| {
         try argv.append("-T");
         try argv.append(linker_script);
diff --git a/src/main.zig b/src/main.zig
index 4fddf9ed21..acfe703291 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -279,6 +279,7 @@ const usage_build_generic =
     \\  -Bsymbolic                     Bind global references locally
     \\  --subsystem [subsystem]        (windows) /SUBSYSTEM:<subsystem> to the linker\n"
     \\  --stack [size]                 Override default stack size
+    \\  --image-base [addr]            Set base address for executable image
     \\  -framework [name]              (darwin) link against framework
     \\  -F[dir]                        (darwin) add search path for frameworks
     \\
@@ -435,6 +436,7 @@ fn buildOutputType(
     var linker_z_defs = false;
     var test_evented_io = false;
     var stack_size_override: ?u64 = null;
+    var image_base_override: ?u64 = null;
     var use_llvm: ?bool = null;
     var use_lld: ?bool = null;
     var use_clang: ?bool = null;
@@ -628,9 +630,11 @@ fn buildOutputType(
                     } else if (mem.eql(u8, arg, "--stack")) {
                         if (i + 1 >= args.len) fatal("expected parameter after {}", .{arg});
                         i += 1;
-                        stack_size_override = std.fmt.parseInt(u64, args[i], 10) catch |err| {
-                            fatal("unable to parse '{}': {}", .{ arg, @errorName(err) });
-                        };
+                        stack_size_override = parseAnyBaseInt(args[i]);
+                    } else if (mem.eql(u8, arg, "--image-base")) {
+                        if (i + 1 >= args.len) fatal("expected parameter after {}", .{arg});
+                        i += 1;
+                        image_base_override = parseAnyBaseInt(args[i]);
                     } else if (mem.eql(u8, arg, "--name")) {
                         if (i + 1 >= args.len) fatal("expected parameter after {}", .{arg});
                         i += 1;
@@ -1147,9 +1151,13 @@ fn buildOutputType(
                     if (i >= linker_args.items.len) {
                         fatal("expected linker arg after '{}'", .{arg});
                     }
-                    stack_size_override = std.fmt.parseInt(u64, linker_args.items[i], 10) catch |err| {
-                        fatal("unable to parse '{}': {}", .{ arg, @errorName(err) });
-                    };
+                    stack_size_override = parseAnyBaseInt(linker_args.items[i]);
+                } else if (mem.eql(u8, arg, "--image-base")) {
+                    i += 1;
+                    if (i >= linker_args.items.len) {
+                        fatal("expected linker arg after '{}'", .{arg});
+                    }
+                    image_base_override = parseAnyBaseInt(linker_args.items[i]);
                 } else {
                     warn("unsupported linker arg: {}", .{arg});
                 }
@@ -1595,6 +1603,7 @@ fn buildOutputType(
         .link_eh_frame_hdr = link_eh_frame_hdr,
         .link_emit_relocs = link_emit_relocs,
         .stack_size_override = stack_size_override,
+        .image_base_override = image_base_override,
         .strip = strip,
         .single_threaded = single_threaded,
         .function_sections = function_sections,
@@ -3051,3 +3060,18 @@ pub fn cleanExit() void {
         process.exit(0);
     }
 }
+
+fn parseAnyBaseInt(prefixed_bytes: []const u8) u64 {
+    const base: u8 = if (mem.startsWith(u8, prefixed_bytes, "0x"))
+        16
+    else if (mem.startsWith(u8, prefixed_bytes, "0o"))
+        8
+    else if (mem.startsWith(u8, prefixed_bytes, "0b"))
+        2
+    else
+        @as(u8, 10);
+    const bytes = if (base == 10) prefixed_bytes else prefixed_bytes[2..];
+    return std.fmt.parseInt(u64, bytes, base) catch |err| {
+        fatal("unable to parse '{}': {}", .{ prefixed_bytes, @errorName(err) });
+    };
+}