Close internal Pipes in websocketconn.Conn Close.bug33367-websocketconn

Unless something externally called Write after Close, the
writeLoop(ws, pr2) goroutine would run forever, because nothing would
ever close pw2/pr2.
https://bugs.torproject.org/33367#comment:4

2 files changed, 30 insertions, 0 deletions

diff --git a/common/websocketconn/websocketconn.go b/common/websocketconn/websocketconn.go
index 46bb977..c745522 100644
--- a/common/websocketconn/websocketconn.go
+++ b/common/websocketconn/websocketconn.go
@@ -24,6 +24,8 @@ func (conn *Conn) Write(b []byte) (n int, err error) {
 }
 
 func (conn *Conn) Close() error {
+	conn.Reader.(*io.PipeReader).Close()
+	conn.Writer.(*io.PipeWriter).Close()
 	// Ignore any error in trying to write a Close frame.
 	_ = conn.Conn.WriteControl(websocket.CloseMessage, []byte{}, time.Now().Add(time.Second))
 	return conn.Conn.Close()
diff --git a/common/websocketconn/websocketconn_test.go b/common/websocketconn/websocketconn_test.go
index ad6f100..92774d4 100644
--- a/common/websocketconn/websocketconn_test.go
+++ b/common/websocketconn/websocketconn_test.go
@@ -233,3 +233,31 @@ func TestConcurrentWrite(t *testing.T) {
 		t.Fatalf("Write: %v", err)
 	}
 }
+
+// Test that Read and Write methods return errors after Close.
+func TestClose(t *testing.T) {
+	s, c, err := connPair()
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer c.Close()
+
+	err = s.Close()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var buf [10]byte
+	n, err := s.Read(buf[:])
+	if n != 0 || err == nil {
+		t.Fatalf("Read after Close returned (%v, %v), expected (%v, non-nil)", n, err, 0)
+	}
+
+	_, err = s.Write([]byte{1, 2, 3})
+	// Here we break the abstraction a little and look for a specific error,
+	// io.ErrClosedPipe. This is because we know the Conn uses an io.Pipe
+	// internally.
+	if err != io.ErrClosedPipe {
+		t.Fatalf("Write after Close returned %v, expected %v", err, io.ErrClosedPipe)
+	}
+}