Add std.Io.net.Stream.shutdown

5 files changed, 118 insertions, 0 deletions

diff --git a/lib/std/Io.zig b/lib/std/Io.zig
index 7c3aa98e16..77eedccb74 100644
--- a/lib/std/Io.zig
+++ b/lib/std/Io.zig
@@ -734,6 +734,7 @@ pub const VTable = struct {
     netWrite: *const fn (?*anyopaque, dest: net.Socket.Handle, header: []const u8, data: []const []const u8, splat: usize) net.Stream.Writer.Error!usize,
     netWriteFile: *const fn (?*anyopaque, net.Socket.Handle, header: []const u8, *Io.File.Reader, Io.Limit) net.Stream.Writer.WriteFileError!usize,
     netClose: *const fn (?*anyopaque, handle: []const net.Socket.Handle) void,
+    netShutdown: *const fn (?*anyopaque, handle: net.Socket.Handle, how: net.ShutdownHow) net.ShutdownError!void,
     netInterfaceNameResolve: *const fn (?*anyopaque, *const net.Interface.Name) net.Interface.Name.ResolveError!net.Interface,
     netInterfaceName: *const fn (?*anyopaque, net.Interface) net.Interface.NameError!net.Interface.Name,
     netLookup: *const fn (?*anyopaque, net.HostName, *Queue(net.HostName.LookupResult), net.HostName.LookupOptions) net.HostName.LookupError!void,
diff --git a/lib/std/Io/Kqueue.zig b/lib/std/Io/Kqueue.zig
index 26b8298cab..df9fa1dee6 100644
--- a/lib/std/Io/Kqueue.zig
+++ b/lib/std/Io/Kqueue.zig
@@ -900,6 +900,7 @@ pub fn io(k: *Kqueue) Io {
             .netConnectIp = netConnectIp,
             .netConnectUnix = netConnectUnix,
             .netClose = netClose,
+            .netShutdown = netShutdown,
             .netRead = netRead,
             .netWrite = netWrite,
             .netSend = netSend,
@@ -1549,12 +1550,22 @@ fn netWrite(userdata: ?*anyopaque, dest: net.Socket.Handle, header: []const u8,
     _ = splat;
     @panic("TODO");
 }
+
 fn netClose(userdata: ?*anyopaque, handle: net.Socket.Handle) void {
     const k: *Kqueue = @ptrCast(@alignCast(userdata));
     _ = k;
     _ = handle;
     @panic("TODO");
 }
+
+fn netShutdown(userdata: ?*anyopaque, handle: net.Socket.Handle, how: net.ShutdownHow) net.ShutdownError!void {
+    const k: *Kqueue = @ptrCast(@alignCast(userdata));
+    _ = k;
+    _ = handle;
+    _ = how;
+    @panic("TODO");
+}
+
 fn netInterfaceNameResolve(
     userdata: ?*anyopaque,
     name: *const net.Interface.Name,
@@ -1564,12 +1575,14 @@ fn netInterfaceNameResolve(
     _ = name;
     @panic("TODO");
 }
+
 fn netInterfaceName(userdata: ?*anyopaque, interface: net.Interface) net.Interface.NameError!net.Interface.Name {
     const k: *Kqueue = @ptrCast(@alignCast(userdata));
     _ = k;
     _ = interface;
     @panic("TODO");
 }
+
 fn netLookup(
     userdata: ?*anyopaque,
     host_name: net.HostName,
diff --git a/lib/std/Io/Threaded.zig b/lib/std/Io/Threaded.zig
index 9e758be9fd..3a1b514561 100644
--- a/lib/std/Io/Threaded.zig
+++ b/lib/std/Io/Threaded.zig
@@ -891,6 +891,10 @@ pub fn io(t: *Threaded) Io {
                 else => netConnectUnixPosix,
             },
             .netClose = netClose,
+            .netShutdown = switch (native_os) {
+                .windows => netShutdownWindows,
+                else => netShutdownPosix,
+            },
             .netRead = switch (native_os) {
                 .windows => netReadWindows,
                 else => netReadPosix,
@@ -1007,6 +1011,7 @@ pub fn ioBasic(t: *Threaded) Io {
             .netConnectIp = netConnectIpUnavailable,
             .netConnectUnix = netConnectUnixUnavailable,
             .netClose = netCloseUnavailable,
+            .netShutdown = netShutdownUnavailable,
             .netRead = netReadUnavailable,
             .netWrite = netWriteUnavailable,
             .netWriteFile = netWriteFileUnavailable,
@@ -10390,6 +10395,89 @@ fn netCloseUnavailable(userdata: ?*anyopaque, handles: []const net.Socket.Handle
     unreachable; // How you gonna close something that was impossible to open?
 }
 
+fn netShutdownPosix(userdata: ?*anyopaque, handle: net.Socket.Handle, how: net.ShutdownHow) net.ShutdownError!void {
+    if (!have_networking) return error.NetworkDown;
+    const t: *Threaded = @ptrCast(@alignCast(userdata));
+    const current_thread = Thread.getCurrent(t);
+
+    const posix_how: i32 = switch (how) {
+        .recv => posix.SHUT.RD,
+        .send => posix.SHUT.WR,
+        .both => posix.SHUT.RDWR,
+    };
+
+    try current_thread.beginSyscall();
+    while (true) {
+        switch (posix.errno(posix.system.shutdown(handle, posix_how))) {
+            .SUCCESS => {
+                current_thread.endSyscall();
+                return;
+            },
+            .INTR => {
+                try current_thread.checkCancel();
+                continue;
+            },
+            else => |e| {
+                current_thread.endSyscall();
+                switch (e) {
+                    .BADF, .NOTSOCK, .INVAL => |err| return errnoBug(err),
+                    .NOTCONN => return error.SocketUnconnected,
+                    .NOBUFS => return error.SystemResources,
+                    else => |err| return posix.unexpectedErrno(err),
+                }
+            },
+        }
+    }
+}
+
+fn netShutdownWindows(userdata: ?*anyopaque, handle: net.Socket.Handle, how: net.ShutdownHow) net.ShutdownError!void {
+    if (!have_networking) return error.NetworkDown;
+    const t: *Threaded = @ptrCast(@alignCast(userdata));
+    const current_thread = Thread.getCurrent(t);
+
+    const wsa_how: i32 = switch (how) {
+        .recv => ws2_32.SD_RECEIVE,
+        .send => ws2_32.SD_SEND,
+        .both => ws2_32.SD_BOTH,
+    };
+
+    try current_thread.beginSyscall();
+    while (true) {
+        const rc = ws2_32.shutdown(handle, wsa_how);
+        if (rc != ws2_32.SOCKET_ERROR) {
+            current_thread.endSyscall();
+            return;
+        }
+        switch (ws2_32.WSAGetLastError()) {
+            .EINTR => {
+                try current_thread.checkCancel();
+                continue;
+            },
+            .NOTINITIALISED => {
+                try initializeWsa(t);
+                try current_thread.checkCancel();
+                continue;
+            },
+            else => |e| {
+                current_thread.endSyscall();
+                switch (e) {
+                    .ECANCELLED, .E_CANCELLED, .OPERATION_ABORTED => return error.Canceled,
+                    .ECONNABORTED => return error.ConnectionAborted,
+                    .ECONNRESET => return error.ConnectionResetByPeer,
+                    .ENETDOWN => return error.NetworkDown,
+                    .ENOTCONN => return error.SocketUnconnected,
+                    .EINVAL, .ENOTSOCK => |err| return wsaErrorBug(err),
+                    else => |err| return windows.unexpectedWSAError(err),
+                }
+            },
+        }
+    }
+}
+
+fn netShutdownUnavailable(_: ?*anyopaque, _: net.Socket.Handle, _: net.ShutdownHow) net.ShutdownError!void {
+    unreachable; // How you gonna shutdown something that was impossible to open?
+}
+
 fn netInterfaceNameResolve(
     userdata: ?*anyopaque,
     name: *const net.Interface.Name,
diff --git a/lib/std/Io/net.zig b/lib/std/Io/net.zig
index 8b1523fbd3..76a581180d 100644
--- a/lib/std/Io/net.zig
+++ b/lib/std/Io/net.zig
@@ -954,6 +954,16 @@ pub const SendFlags = packed struct(u8) {
     _: u3 = 0,
 };
 
+pub const ShutdownHow = enum { recv, send, both };
+
+pub const ShutdownError = error{
+    ConnectionAborted,
+    ConnectionResetByPeer,
+    NetworkDown,
+    SocketUnconnected,
+    SystemResources,
+} || Io.UnexpectedError || Io.Cancelable;
+
 pub const Interface = struct {
     /// Value 0 indicates `none`.
     index: u32,
@@ -1191,6 +1201,10 @@ pub const Stream = struct {
         io.vtable.netClose(io.userdata, (&s.socket.handle)[0..1]);
     }
 
+    pub fn shutdown(s: *const Stream, io: Io, how: ShutdownHow) ShutdownError!void {
+        return io.vtable.netShutdown(io.userdata, s.socket.handle, how);
+    }
+
     pub const Reader = struct {
         io: Io,
         interface: Io.Reader,
diff --git a/lib/std/Io/net/test.zig b/lib/std/Io/net/test.zig
index 6ef8c15f4f..45c26f9540 100644
--- a/lib/std/Io/net/test.zig
+++ b/lib/std/Io/net/test.zig
@@ -346,6 +346,8 @@ test "non-blocking tcp server" {
     const len = try socket_file.read(&buf);
     const msg = buf[0..len];
     try testing.expect(mem.eql(u8, msg, "hello from server\n"));
+
+    try stream.shutdown(io, .both);
 }
 
 test "decompress compressed DNS name" {