Commit b587a42233

Jonathan Marler <johnnymarler@gmail.com>
2020-11-28 08:42:17
add std.os.shutdown function for sockets
1 parent baa075a
lib/std/os/bits/darwin.zig
@@ -1523,3 +1523,7 @@ pub const rlimit = extern struct {
     /// Hard limit
     max: rlim_t,
 };
+
+pub const SHUT_RD = 0;
+pub const SHUT_WR = 1;
+pub const SHUT_RDWR = 2;
lib/std/os/bits/freebsd.zig
@@ -1388,3 +1388,7 @@ pub const rlimit = extern struct {
     /// Hard limit
     max: rlim_t,
 };
+
+pub const SHUT_RD = 0;
+pub const SHUT_WR = 1;
+pub const SHUT_RDWR = 2;
lib/std/os/bits/netbsd.zig
@@ -1196,3 +1196,7 @@ pub const rlimit = extern struct {
     /// Hard limit
     max: rlim_t,
 };
+
+pub const SHUT_RD = 0;
+pub const SHUT_WR = 1;
+pub const SHUT_RDWR = 2;
lib/std/os/bits/openbsd.zig
@@ -1134,3 +1134,7 @@ pub const rlimit = extern struct {
     /// Hard limit
     max: rlim_t,
 };
+
+pub const SHUT_RD = 0;
+pub const SHUT_WR = 1;
+pub const SHUT_RDWR = 2;
lib/std/os/windows/bits.zig
@@ -1602,3 +1602,7 @@ pub const MOUNTMGR_MOUNT_POINTS = extern struct {
     MountPoints: [1]MOUNTMGR_MOUNT_POINT,
 };
 pub const IOCTL_MOUNTMGR_QUERY_POINTS: ULONG = 0x6d0008;
+
+pub const SD_RECEIVE = 0;
+pub const SD_SEND = 1;
+pub const SD_BOTH = 2;
lib/std/os/windows/ws2_32.zig
@@ -877,3 +877,7 @@ pub extern "ws2_32" fn setsockopt(
     optval: ?*const c_void,
     optlen: socklen_t,
 ) callconv(WINAPI) c_int;
+pub extern "ws2_32" fn shutdown(
+    s: SOCKET,
+    how: c_int,
+) callconv(WINAPI) c_int;
lib/std/os/test.zig
@@ -627,3 +627,22 @@ test "getrlimit and setrlimit" {
         try os.setrlimit(resource, limit);
     }
 }
+
+test "shutdown socket" {
+    if (builtin.os.tag == .wasi)
+        return error.SkipZigTest;
+    if (builtin.os.tag == .windows) {
+        _ = try std.os.windows.WSAStartup(2, 2);
+    }
+    defer {
+        if (builtin.os.tag == .windows) {
+            std.os.windows.WSACleanup() catch unreachable;
+        }
+    }
+    const sock = try os.socket(os.AF_INET, os.SOCK_STREAM, 0);
+    os.shutdown(sock, .both) catch |err| switch (err) {
+        error.SocketNotConnected => {},
+        else => |e| return e,
+    };
+    os.closeSocket(sock);
+}
lib/std/c.zig
@@ -145,6 +145,7 @@ pub extern "c" fn ioctl(fd: fd_t, request: c_int, ...) c_int;
 pub extern "c" fn uname(buf: *utsname) c_int;
 
 pub extern "c" fn gethostname(name: [*]u8, len: usize) c_int;
+pub extern "c" fn shutdown(socket: fd_t, how: c_int) c_int;
 pub extern "c" fn bind(socket: fd_t, address: ?*const sockaddr, address_len: socklen_t) c_int;
 pub extern "c" fn socketpair(domain: c_uint, sock_type: c_uint, protocol: c_uint, sv: *[2]fd_t) c_int;
 pub extern "c" fn listen(sockfd: fd_t, backlog: c_uint) c_int;
lib/std/os.zig
@@ -2706,6 +2706,62 @@ pub fn socket(domain: u32, socket_type: u32, protocol: u32) SocketError!socket_t
     }
 }
 
+pub const ShutdownError = error{
+    ConnectionAborted,
+
+    /// Connection was reset by peer, application should close socket as it is no longer usable.
+    ConnectionResetByPeer,
+
+    BlockingOperationInProgress,
+
+    /// The network subsystem has failed.
+    NetworkSubsystemFailed,
+
+    /// The socket is not connected (connection-oriented sockets only).
+    SocketNotConnected,
+
+    SystemResources
+} || UnexpectedError;
+
+pub const ShutdownHow = enum { recv, send, both };
+
+/// Shutdown socket send/receive operations
+pub fn shutdown(sock: socket_t, how: ShutdownHow) ShutdownError!void {
+    if (builtin.os.tag == .windows) {
+        const result = windows.ws2_32.shutdown(sock, switch (how) {
+            .recv => windows.SD_RECEIVE,
+            .send => windows.SD_SEND,
+            .both => windows.SD_BOTH,
+        });
+        if (0 != result) switch (windows.ws2_32.WSAGetLastError()) {
+            .WSAECONNABORTED => return error.ConnectionAborted,
+            .WSAECONNRESET => return error.ConnectionResetByPeer,
+            .WSAEINPROGRESS => return error.BlockingOperationInProgress,
+            .WSAEINVAL => unreachable,
+            .WSAENETDOWN => return error.NetworkSubsystemFailed,
+            .WSAENOTCONN => return error.SocketNotConnected,
+            .WSAENOTSOCK => unreachable,
+            .WSANOTINITIALISED => unreachable,
+            else => |err| return windows.unexpectedWSAError(err),
+        };
+    } else {
+        const rc = system.shutdown(sock, switch (how) {
+            .recv => SHUT_RD,
+            .send => SHUT_WR,
+            .both => SHUT_RDWR,
+        });
+        switch (errno(rc)) {
+            0 => return,
+            EBADF => unreachable,
+            EINVAL => unreachable,
+            ENOTCONN => return error.SocketNotConnected,
+            ENOTSOCK => unreachable,
+            ENOBUFS => return error.SystemResources,
+            else => |err| return unexpectedErrno(err),
+        }
+    }
+}
+
 pub fn closeSocket(sock: socket_t) void {
     if (builtin.os.tag == .windows) {
         windows.closesocket(sock) catch unreachable;