Fix client exit failure with an explicit Close

Author: belak

client.go

@@ -50,6 +50,7 @@ type cap struct {
 // much simpler.
 type Client struct {
 	*Conn
+	rwc    io.ReadWriteCloser
 	config ClientConfig
 
 	// Internal state
@@ -63,9 +64,10 @@ type Client struct {
 }
 
 // NewClient creates a client given an io stream and a client config.
-func NewClient(rw io.ReadWriter, config ClientConfig) *Client {
+func NewClient(rwc io.ReadWriteCloser, config ClientConfig) *Client {
 	c := &Client{
-		Conn:    NewConn(rw),
+		Conn:    NewConn(rwc),
+		rwc:     rwc,
 		config:  config,
 		errChan: make(chan error, 1),
 		caps:    make(map[string]cap),
@@ -312,6 +314,7 @@ func (c *Client) RunContext(ctx context.Context) error {
 	}
 
 	close(exiting)
+	c.rwc.Close()
 	wg.Wait()
 
 	return err

client_test.go

@@ -408,4 +408,16 @@ func TestPingLoop(t *testing.T) {
 		SendLine("001 :hello_world\r\n"),
 		Delay(25 * time.Millisecond),
 	})
+
+	// See if we can get the client to hang
+	runClientTest(t, config, errors.New("test error"), nil, []TestAction{
+		ExpectLine("PASS :test_pass\r\n"),
+		ExpectLine("NICK :test_nick\r\n"),
+		ExpectLine("USER test_user 0 * :test_name\r\n"),
+		// We queue this up a line early because the next write will happen after the delay.
+		QueueWriteError(errors.New("test error")),
+		SendLine("001 :hello_world\r\n"),
+		Delay(2 * time.Second),
+		AssertClosed(),
+	})
 }

stream_test.go

@@ -18,6 +18,14 @@ func SendLine(output string) TestAction {
 	return SendLineWithTimeout(output, 1*time.Second)
 }
 
+func AssertClosed() TestAction {
+	return func(t *testing.T, rw *testReadWriter) {
+		if !rw.closed {
+			assert.Fail(t, "Expected conn to be closed")
+		}
+	}
+}
+
 func SendLineWithTimeout(output string, timeout time.Duration) TestAction {
 	return func(t *testing.T, rw *testReadWriter) {
 		waitChan := time.After(timeout)
@@ -116,6 +124,7 @@ type testReadWriter struct {
 	readEmptyChan  chan struct{}
 	exiting        chan struct{}
 	clientDone     chan struct{}
+	closed         bool
 	serverBuffer   bytes.Buffer
 }
 
@@ -182,6 +191,20 @@ func (rw *testReadWriter) Write(buf []byte) (int, error) {
 	}
 }
 
+func (rw *testReadWriter) Close() error {
+	select {
+	case <-rw.exiting:
+		return errors.New("Connection closed")
+	default:
+		// Ensure no double close
+		if !rw.closed {
+			rw.closed = true
+			close(rw.exiting)
+		}
+		return nil
+	}
+}
+
 func newTestReadWriter(actions []TestAction) *testReadWriter {
 	return &testReadWriter{
 		actions:        actions,
@@ -223,7 +246,7 @@ func runTest(t *testing.T, rw *testReadWriter, actions []TestAction) {
 	// TODO: Make sure there are no more incoming messages
 
 	// Ask everything to shut down
-	close(rw.exiting)
+	rw.Close()
 
 	// Wait for the client to stop
 	select {