Commit 9fa55ae777

Andrew Kelley <andrew@ziglang.org>
2022-01-11 19:02:28
Merge pull request #10566 from fifty-six/master
std.os.uefi improvements/fixes
1 parent 53e4168
lib/std/fs/path.zig
@@ -14,11 +14,17 @@ const native_os = builtin.target.os.tag;
 
 pub const sep_windows = '\\';
 pub const sep_posix = '/';
-pub const sep = if (native_os == .windows) sep_windows else sep_posix;
+pub const sep = switch (native_os) {
+    .windows, .uefi => sep_windows,
+    else => sep_posix,
+};
 
 pub const sep_str_windows = "\\";
 pub const sep_str_posix = "/";
-pub const sep_str = if (native_os == .windows) sep_str_windows else sep_str_posix;
+pub const sep_str = switch (native_os) {
+    .windows, .uefi => sep_str_windows,
+    else => sep_str_posix,
+};
 
 pub const delimiter_windows = ';';
 pub const delimiter_posix = ':';
@@ -26,11 +32,11 @@ pub const delimiter = if (native_os == .windows) delimiter_windows else delimite
 
 /// Returns if the given byte is a valid path separator
 pub fn isSep(byte: u8) bool {
-    if (native_os == .windows) {
-        return byte == '/' or byte == '\\';
-    } else {
-        return byte == '/';
-    }
+    return switch (native_os) {
+        .windows => byte == '/' or byte == '\\',
+        .uefi => byte == '\\',
+        else => byte == '/',
+    };
 }
 
 /// This is different from mem.join in that the separator will not be repeated if
@@ -110,6 +116,17 @@ pub fn joinZ(allocator: Allocator, paths: []const []const u8) ![:0]u8 {
     return out[0 .. out.len - 1 :0];
 }
 
+fn testJoinMaybeZUefi(paths: []const []const u8, expected: []const u8, zero: bool) !void {
+    const uefiIsSep = struct {
+        fn isSep(byte: u8) bool {
+            return byte == '\\';
+        }
+    }.isSep;
+    const actual = try joinSepMaybeZ(testing.allocator, sep_windows, uefiIsSep, paths, zero);
+    defer testing.allocator.free(actual);
+    try testing.expectEqualSlices(u8, expected, if (zero) actual[0 .. actual.len - 1 :0] else actual);
+}
+
 fn testJoinMaybeZWindows(paths: []const []const u8, expected: []const u8, zero: bool) !void {
     const windowsIsSep = struct {
         fn isSep(byte: u8) bool {
@@ -158,6 +175,11 @@ test "join" {
             zero,
         );
 
+        try testJoinMaybeZUefi(&[_][]const u8{ "EFI", "Boot", "bootx64.efi" }, "EFI\\Boot\\bootx64.efi", zero);
+        try testJoinMaybeZUefi(&[_][]const u8{ "EFI\\Boot", "bootx64.efi" }, "EFI\\Boot\\bootx64.efi", zero);
+        try testJoinMaybeZUefi(&[_][]const u8{ "EFI\\", "\\Boot", "bootx64.efi" }, "EFI\\Boot\\bootx64.efi", zero);
+        try testJoinMaybeZUefi(&[_][]const u8{ "EFI\\", "\\Boot\\", "\\bootx64.efi" }, "EFI\\Boot\\bootx64.efi", zero);
+
         try testJoinMaybeZWindows(&[_][]const u8{ "c:\\", "a", "b/", "c" }, "c:\\a\\b/c", zero);
         try testJoinMaybeZWindows(&[_][]const u8{ "c:\\a/", "b\\", "/c" }, "c:\\a/b\\c", zero);
 
lib/std/os/uefi/protocols/device_path_protocol.zig
@@ -1,4 +1,7 @@
-const uefi = @import("std").os.uefi;
+const std = @import("std");
+const mem = std.mem;
+const uefi = std.os.uefi;
+const Allocator = mem.Allocator;
 const Guid = uefi.Guid;
 
 pub const DevicePathProtocol = packed struct {
@@ -15,6 +18,59 @@ pub const DevicePathProtocol = packed struct {
         .node = [_]u8{ 0x00, 0xa0, 0xc9, 0x69, 0x72, 0x3b },
     };
 
+    /// Returns the next DevicePathProtocol node in the sequence, if any.
+    pub fn next(self: *DevicePathProtocol) ?*DevicePathProtocol {
+        if (self.type == .End and @intToEnum(EndDevicePath.Subtype, self.subtype) == .EndEntire)
+            return null;
+
+        return @ptrCast(*DevicePathProtocol, @ptrCast([*]u8, self) + self.length);
+    }
+
+    /// Calculates the total length of the device path structure in bytes, including the end of device path node.
+    pub fn size(self: *DevicePathProtocol) usize {
+        var node = self;
+
+        while (node.next()) |next_node| {
+            node = next_node;
+        }
+
+        return (@ptrToInt(node) + node.length) - @ptrToInt(self);
+    }
+
+    /// Creates a file device path from the existing device path and a file path.
+    pub fn create_file_device_path(self: *DevicePathProtocol, allocator: Allocator, path: [:0]const u16) !*DevicePathProtocol {
+        var path_size = self.size();
+
+        // 2 * (path.len + 1) for the path and its null terminator, which are u16s
+        // DevicePathProtocol for the extra node before the end
+        var buf = try allocator.alloc(u8, path_size + 2 * (path.len + 1) + @sizeOf(DevicePathProtocol));
+
+        mem.copy(u8, buf, @ptrCast([*]const u8, self)[0..path_size]);
+
+        // Pointer to the copy of the end node of the current chain, which is - 4 from the buffer
+        // as the end node itself is 4 bytes (type: u8 + subtype: u8 + length: u16).
+        var new = @ptrCast(*MediaDevicePath.FilePathDevicePath, buf.ptr + path_size - 4);
+
+        new.type = .Media;
+        new.subtype = .FilePath;
+        new.length = @sizeOf(MediaDevicePath.FilePathDevicePath) + 2 * (@intCast(u16, path.len) + 1);
+
+        // The same as new.getPath(), but not const as we're filling it in.
+        var ptr = @ptrCast([*:0]u16, @alignCast(2, @ptrCast([*]u8, new)) + @sizeOf(MediaDevicePath.FilePathDevicePath));
+
+        for (path) |s, i|
+            ptr[i] = s;
+
+        ptr[path.len] = 0;
+
+        var end = @ptrCast(*EndDevicePath.EndEntireDevicePath, @ptrCast(*DevicePathProtocol, new).next().?);
+        end.type = .End;
+        end.subtype = .EndEntire;
+        end.length = @sizeOf(EndDevicePath.EndEntireDevicePath);
+
+        return @ptrCast(*DevicePathProtocol, buf.ptr);
+    }
+
     pub fn getDevicePath(self: *const DevicePathProtocol) ?DevicePath {
         return switch (self.type) {
             .Hardware => blk: {
lib/std/os/uefi/protocols/file_protocol.zig
@@ -127,15 +127,6 @@ pub const FileProtocol = extern struct {
         return self._flush(self);
     }
 
-    pub const guid align(8) = Guid{
-        .time_low = 0x09576e92,
-        .time_mid = 0x6d3f,
-        .time_high_and_version = 0x11d2,
-        .clock_seq_high_and_reserved = 0x8e,
-        .clock_seq_low = 0x39,
-        .node = [_]u8{ 0x00, 0xa0, 0xc9, 0x69, 0x72, 0x3b },
-    };
-
     pub const efi_file_mode_read: u64 = 0x0000000000000001;
     pub const efi_file_mode_write: u64 = 0x0000000000000002;
     pub const efi_file_mode_create: u64 = 0x8000000000000000;
@@ -171,4 +162,35 @@ pub const FileInfo = extern struct {
     pub const efi_file_directory: u64 = 0x0000000000000010;
     pub const efi_file_archive: u64 = 0x0000000000000020;
     pub const efi_file_valid_attr: u64 = 0x0000000000000037;
+
+    pub const guid align(8) = Guid{
+        .time_low = 0x09576e92,
+        .time_mid = 0x6d3f,
+        .time_high_and_version = 0x11d2,
+        .clock_seq_high_and_reserved = 0x8e,
+        .clock_seq_low = 0x39,
+        .node = [_]u8{ 0x00, 0xa0, 0xc9, 0x69, 0x72, 0x3b },
+    };
+};
+
+pub const FileSystemInfo = extern struct {
+    size: u64,
+    read_only: bool,
+    volume_size: u64,
+    free_space: u64,
+    block_size: u32,
+    _volume_label: u16,
+
+    pub fn getVolumeLabel(self: *const FileSystemInfo) [*:0]const u16 {
+        return @ptrCast([*:0]const u16, &self._volume_label);
+    }
+
+    pub const guid align(8) = Guid{
+        .time_low = 0x09576e93,
+        .time_mid = 0x6d3f,
+        .time_high_and_version = 0x11d2,
+        .clock_seq_high_and_reserved = 0x8e,
+        .clock_seq_low = 0x39,
+        .node = [_]u8{ 0x00, 0xa0, 0xc9, 0x69, 0x72, 0x3b },
+    };
 };
lib/std/os/uefi/pool_allocator.zig
@@ -0,0 +1,153 @@
+const std = @import("std");
+
+const mem = std.mem;
+const uefi = std.os.uefi;
+
+const assert = std.debug.assert;
+
+const Allocator = mem.Allocator;
+
+const UefiPoolAllocator = struct {
+    fn getHeader(ptr: [*]u8) *[*]align(8) u8 {
+        return @intToPtr(*[*]align(8) u8, @ptrToInt(ptr) - @sizeOf(usize));
+    }
+
+    fn alignedAlloc(len: usize, alignment: usize) ?[*]u8 {
+        var unaligned_ptr: [*]align(8) u8 = undefined;
+
+        if (uefi.system_table.boot_services.?.allocatePool(uefi.efi_pool_memory_type, len, &unaligned_ptr) != .Success)
+            return null;
+
+        const unaligned_addr = @ptrToInt(unaligned_ptr);
+        const aligned_addr = mem.alignForward(unaligned_addr + @sizeOf(usize), alignment);
+
+        var aligned_ptr = unaligned_ptr + (aligned_addr - unaligned_addr);
+        getHeader(aligned_ptr).* = unaligned_ptr;
+
+        return aligned_ptr;
+    }
+
+    fn alignedFree(ptr: [*]u8) void {
+        _ = uefi.system_table.boot_services.?.freePool(getHeader(ptr).*);
+    }
+
+    fn alloc(
+        _: *anyopaque,
+        len: usize,
+        ptr_align: u29,
+        len_align: u29,
+        ret_addr: usize,
+    ) Allocator.Error![]u8 {
+        _ = ret_addr;
+
+        assert(len > 0);
+        assert(std.math.isPowerOfTwo(ptr_align));
+
+        var ptr = alignedAlloc(len, ptr_align) orelse return error.OutOfMemory;
+
+        if (len_align == 0)
+            return ptr[0..len];
+
+        return ptr[0..mem.alignBackwardAnyAlign(len, len_align)];
+    }
+
+    fn resize(
+        _: *anyopaque,
+        buf: []u8,
+        buf_align: u29,
+        new_len: usize,
+        len_align: u29,
+        ret_addr: usize,
+    ) ?usize {
+        _ = buf_align;
+        _ = ret_addr;
+
+        return if (new_len <= buf.len) mem.alignAllocLen(buf.len, new_len, len_align) else null;
+    }
+
+    fn free(
+        _: *anyopaque,
+        buf: []u8,
+        buf_align: u29,
+        ret_addr: usize,
+    ) void {
+        _ = buf_align;
+        _ = ret_addr;
+        alignedFree(buf.ptr);
+    }
+};
+
+/// Supports the full Allocator interface, including alignment.
+/// For a direct call of `allocatePool`, see `raw_pool_allocator`.
+pub const pool_allocator = Allocator{
+    .ptr = undefined,
+    .vtable = &pool_allocator_vtable,
+};
+
+const pool_allocator_vtable = Allocator.VTable{
+    .alloc = UefiPoolAllocator.alloc,
+    .resize = UefiPoolAllocator.resize,
+    .free = UefiPoolAllocator.free,
+};
+
+/// Asserts allocations are 8 byte aligned and calls `boot_services.allocatePool`. 
+pub const raw_pool_allocator = Allocator{
+    .ptr = undefined,
+    .vtable = &raw_pool_allocator_table,
+};
+
+const raw_pool_allocator_table = Allocator.VTable{
+    .alloc = uefi_alloc,
+    .resize = uefi_resize,
+    .free = uefi_free,
+};
+
+fn uefi_alloc(
+    _: *anyopaque,
+    len: usize,
+    ptr_align: u29,
+    len_align: u29,
+    ret_addr: usize,
+) Allocator.Error![]u8 {
+    _ = len_align;
+    _ = ret_addr;
+
+    std.debug.assert(ptr_align <= 8);
+
+    var ptr: [*]align(8) u8 = undefined;
+
+    if (uefi.system_table.boot_services.?.allocatePool(uefi.efi_pool_memory_type, len, &ptr) != .Success) {
+        return error.OutOfMemory;
+    }
+
+    return ptr[0..len];
+}
+
+fn uefi_resize(
+    _: *anyopaque,
+    buf: []u8,
+    old_align: u29,
+    new_len: usize,
+    len_align: u29,
+    ret_addr: usize,
+) ?usize {
+    _ = old_align;
+    _ = ret_addr;
+
+    if (new_len <= buf.len) {
+        return mem.alignAllocLen(buf.len, new_len, len_align);
+    }
+
+    return null;
+}
+
+fn uefi_free(
+    _: *anyopaque,
+    buf: []u8,
+    buf_align: u29,
+    ret_addr: usize,
+) void {
+    _ = buf_align;
+    _ = ret_addr;
+    _ = uefi.system_table.boot_services.?.freePool(@alignCast(8, buf.ptr));
+}
lib/std/os/uefi/protocols.zig
@@ -14,6 +14,7 @@ pub const MessagingDevicePath = @import("protocols/device_path_protocol.zig").Me
 pub const SimpleFileSystemProtocol = @import("protocols/simple_file_system_protocol.zig").SimpleFileSystemProtocol;
 pub const FileProtocol = @import("protocols/file_protocol.zig").FileProtocol;
 pub const FileInfo = @import("protocols/file_protocol.zig").FileInfo;
+pub const FileSystemInfo = @import("protocols/file_protocol.zig").FileSystemInfo;
 
 pub const InputKey = @import("protocols/simple_text_input_ex_protocol.zig").InputKey;
 pub const KeyData = @import("protocols/simple_text_input_ex_protocol.zig").KeyData;
lib/std/os/uefi.zig
@@ -7,6 +7,13 @@ pub const protocols = @import("uefi/protocols.zig");
 pub const Status = @import("uefi/status.zig").Status;
 pub const tables = @import("uefi/tables.zig");
 
+/// The memory type to allocate when using the pool
+/// Defaults to .LoaderData, the default data allocation type
+/// used by UEFI applications to allocate pool memory.
+pub var efi_pool_memory_type: tables.MemoryType = .LoaderData;
+pub const pool_allocator = @import("uefi/pool_allocator.zig").pool_allocator;
+pub const raw_pool_allocator = @import("uefi/pool_allocator.zig").raw_pool_allocator;
+
 /// The EFI image's handle that is passed to its entry point.
 pub var handle: Handle = undefined;