From 488f68069ba0a81f46a5afe721b908cfd3c61f76 Mon Sep 17 00:00:00 2001
From: Andrew Kelley <andrew@ziglang.org>
Date: Mon, 28 Jun 2021 12:06:14 -0700
Subject: [PATCH] implement std.fs.File.setLock for Windows

---
 lib/std/fs/file.zig          | 31 ++++++++++++++++++--
 lib/std/os/windows.zig       | 56 ++++++++++++++++++++++++++++++++++++
 lib/std/os/windows/ntdll.zig | 21 ++++++++++++++
 3 files changed, 106 insertions(+), 2 deletions(-)

diff --git a/lib/std/fs/file.zig b/lib/std/fs/file.zig
index 67da68901c..aecfe9e6be 100644
--- a/lib/std/fs/file.zig
+++ b/lib/std/fs/file.zig
@@ -840,10 +840,37 @@ pub const File = struct {
     /// TODO: integrate with async I/O
     pub fn setLock(file: File, lock: Lock, non_blocking: bool) SetLockError!void {
         if (is_windows) {
-            @compileError("TODO implement fs.File.setLock for Windows");
+            const range_off: windows.LARGE_INTEGER = 0;
+            const range_len: windows.LARGE_INTEGER = 1;
+            const exclusive = switch (lock) {
+                .None => return windows.UnlockFile(
+                    file.handle,
+                    null,
+                    &range_off,
+                    &range_len,
+                    null,
+                ) catch |err| switch (err) {
+                    error.RangeNotLocked => return,
+                    else => |e| return e,
+                },
+                .Shared => false,
+                .Exclusive => true,
+            };
+            return windows.LockFile(
+                file.handle,
+                null,
+                null,
+                null,
+                null,
+                &range_off,
+                &range_len,
+                null,
+                @boolToInt(non_blocking),
+                @boolToInt(exclusive),
+            );
         }
         const non_blocking_flag = if (non_blocking) os.LOCK_NB else @as(i32, 0);
-        try os.flock(file.handle, switch (lock) {
+        return os.flock(file.handle, switch (lock) {
             .None => os.LOCK_UN,
             .Shared => os.LOCK_SH | non_blocking_flag,
             .Exclusive => os.LOCK_EX | non_blocking_flag,
diff --git a/lib/std/os/windows.zig b/lib/std/os/windows.zig
index 1fade2a462..a53957aad3 100644
--- a/lib/std/os/windows.zig
+++ b/lib/std/os/windows.zig
@@ -1679,6 +1679,62 @@ pub fn SetFileTime(
     }
 }
 
+pub const LockFileError = error{
+    SystemResources,
+    WouldBlock,
+} || std.os.UnexpectedError;
+
+pub fn LockFile(
+    FileHandle: HANDLE,
+    Event: ?HANDLE,
+    ApcRoutine: ?*IO_APC_ROUTINE,
+    ApcContext: ?*c_void,
+    IoStatusBlock: ?*IO_STATUS_BLOCK,
+    ByteOffset: *const LARGE_INTEGER,
+    Length: *const LARGE_INTEGER,
+    Key: ?*ULONG,
+    FailImmediately: BOOLEAN,
+    ExclusiveLock: BOOLEAN,
+) !void {
+    const rc = ntdll.NtLockFile(
+        FileHandle,
+        Event,
+        ApcRoutine,
+        ApcContext,
+        IoStatusBlock,
+        ByteOffset,
+        Length,
+        Key,
+        FailImmediately,
+        ExclusiveLock,
+    );
+    switch (rc) {
+        .SUCCESS => return,
+        .INSUFFICIENT_RESOURCES => return error.SystemResources,
+        .LOCK_NOT_GRANTED => return error.WouldBlock,
+        else => return unexpectedStatus(rc),
+    }
+}
+
+pub const UnlockFileError = error{
+    RangeNotLocked,
+} || std.os.UnexpectedError;
+
+pub fn UnlockFile(
+    FileHandle: HANDLE,
+    IoStatusBlock: ?*IO_STATUS_BLOCK,
+    ByteOffset: *const LARGE_INTEGER,
+    Length: *const LARGE_INTEGER,
+    Key: ?*ULONG,
+) !void {
+    const rc = ntdll.NtUnlockFile(FileHandle, IoStatusBlock, ByteOffset, Length, Key);
+    switch (rc) {
+        .SUCCESS => return,
+        .RANGE_NOT_LOCKED => return error.RangeNotLocked,
+        else => return unexpectedStatus(rc),
+    }
+}
+
 pub fn teb() *TEB {
     return switch (builtin.target.cpu.arch) {
         .i386 => asm volatile (
diff --git a/lib/std/os/windows/ntdll.zig b/lib/std/os/windows/ntdll.zig
index aa7e382ee0..1a49634f6a 100644
--- a/lib/std/os/windows/ntdll.zig
+++ b/lib/std/os/windows/ntdll.zig
@@ -139,3 +139,24 @@ pub extern "NtDll" fn RtlWaitOnAddress(
     AddressSize: SIZE_T,
     Timeout: ?*const LARGE_INTEGER,
 ) callconv(WINAPI) NTSTATUS;
+
+pub extern "NtDll" fn NtLockFile(
+    FileHandle: HANDLE,
+    Event: ?HANDLE,
+    ApcRoutine: ?*IO_APC_ROUTINE,
+    ApcContext: ?*c_void,
+    IoStatusBlock: ?*IO_STATUS_BLOCK,
+    ByteOffset: *const LARGE_INTEGER,
+    Length: *const LARGE_INTEGER,
+    Key: ?*ULONG,
+    FailImmediately: BOOLEAN,
+    ExclusiveLock: BOOLEAN,
+) callconv(WINAPI) NTSTATUS;
+
+pub extern "NtDll" fn NtUnlockFile(
+    FileHandle: HANDLE,
+    IoStatusBlock: ?*IO_STATUS_BLOCK,
+    ByteOffset: *const LARGE_INTEGER,
+    Length: *const LARGE_INTEGER,
+    Key: ?*ULONG,
+) callconv(WINAPI) NTSTATUS;