From: Brad Fitzpatrick <bradfitz@golang.org>
Date: Wed, 28 Sep 2011 15:12:38 +0000 (-0700)
Subject: net: add shutdown: TCPConn.CloseWrite and CloseRead
X-Git-Tag: weekly.2011-10-06~77
X-Git-Url: http://www.git.cypherpunks.su/?a=commitdiff_plain;h=394842e2a56e49cb4e81f1d3cc875195cfe507a1;p=gostls13.git

net: add shutdown: TCPConn.CloseWrite and CloseRead

R=golang-dev, rsc, iant
CC=golang-dev
https://golang.org/cl/5136052
---

diff --git a/src/pkg/net/fd.go b/src/pkg/net/fd.go
index 9084e88755..a0c56f78ee 100644
--- a/src/pkg/net/fd.go
+++ b/src/pkg/net/fd.go
@@ -358,6 +358,22 @@ func (fd *netFD) Close() os.Error {
 	return nil
 }
 
+func (fd *netFD) CloseRead() os.Error {
+	if fd == nil || fd.sysfile == nil {
+		return os.EINVAL
+	}
+	syscall.Shutdown(fd.sysfd, syscall.SHUT_RD)
+	return nil
+}
+
+func (fd *netFD) CloseWrite() os.Error {
+	if fd == nil || fd.sysfile == nil {
+		return os.EINVAL
+	}
+	syscall.Shutdown(fd.sysfd, syscall.SHUT_WR)
+	return nil
+}
+
 func (fd *netFD) Read(p []byte) (n int, err os.Error) {
 	if fd == nil {
 		return 0, os.EINVAL
diff --git a/src/pkg/net/fd_windows.go b/src/pkg/net/fd_windows.go
index b025bddea0..8155d04aae 100644
--- a/src/pkg/net/fd_windows.go
+++ b/src/pkg/net/fd_windows.go
@@ -312,6 +312,22 @@ func (fd *netFD) Close() os.Error {
 	return nil
 }
 
+func (fd *netFD) CloseRead() os.Error {
+	if fd == nil || fd.sysfd == syscall.InvalidHandle {
+		return os.EINVAL
+	}
+	syscall.Shutdown(fd.sysfd, syscall.SHUT_RD)
+	return nil
+}
+
+func (fd *netFD) CloseWrite() os.Error {
+	if fd == nil || fd.sysfd == syscall.InvalidHandle {
+		return os.EINVAL
+	}
+	syscall.Shutdown(fd.sysfd, syscall.SHUT_WR)
+	return nil
+}
+
 // Read from network.
 
 type readOp struct {
diff --git a/src/pkg/net/net_test.go b/src/pkg/net/net_test.go
index 698a845277..e4d7a253e2 100644
--- a/src/pkg/net/net_test.go
+++ b/src/pkg/net/net_test.go
@@ -6,6 +6,7 @@ package net
 
 import (
 	"flag"
+	"os"
 	"regexp"
 	"testing"
 )
@@ -119,3 +120,46 @@ func TestReverseAddress(t *testing.T) {
 		}
 	}
 }
+
+func TestShutdown(t *testing.T) {
+	l, err := Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		if l, err = Listen("tcp6", "[::1]:0"); err != nil {
+			t.Fatalf("ListenTCP on :0: %v", err)
+		}
+	}
+
+	go func() {
+		c, err := l.Accept()
+		if err != nil {
+			t.Fatalf("Accept: %v", err)
+		}
+		var buf [10]byte
+		n, err := c.Read(buf[:])
+		if n != 0 || err != os.EOF {
+			t.Fatalf("server Read = %d, %v; want 0, os.EOF", n, err)
+		}
+		c.Write([]byte("response"))
+		c.Close()
+	}()
+
+	c, err := Dial("tcp", l.Addr().String())
+	if err != nil {
+		t.Fatalf("Dial: %v", err)
+	}
+	defer c.Close()
+
+	err = c.(*TCPConn).CloseWrite()
+	if err != nil {
+		t.Fatalf("CloseWrite: %v", err)
+	}
+	var buf [10]byte
+	n, err := c.Read(buf[:])
+	if err != nil {
+		t.Fatalf("client Read: %d, %v", n, err)
+	}
+	got := string(buf[:n])
+	if got != "response" {
+		t.Errorf("read = %q, want \"response\"", got)
+	}
+}
diff --git a/src/pkg/net/tcpsock_posix.go b/src/pkg/net/tcpsock_posix.go
index 35d536c319..740a63d303 100644
--- a/src/pkg/net/tcpsock_posix.go
+++ b/src/pkg/net/tcpsock_posix.go
@@ -100,6 +100,24 @@ func (c *TCPConn) Close() os.Error {
 	return err
 }
 
+// CloseRead shuts down the reading side of the TCP connection.
+// Most callers should just use Close.
+func (c *TCPConn) CloseRead() os.Error {
+	if !c.ok() {
+		return os.EINVAL
+	}
+	return c.fd.CloseRead()
+}
+
+// CloseWrite shuts down the writing side of the TCP connection.
+// Most callers should just use Close.
+func (c *TCPConn) CloseWrite() os.Error {
+	if !c.ok() {
+		return os.EINVAL
+	}
+	return c.fd.CloseWrite()
+}
+
 // LocalAddr returns the local network address, a *TCPAddr.
 func (c *TCPConn) LocalAddr() Addr {
 	if !c.ok() {