fix: add IsActive checks to Writer methods to prevent panic after Close (#416)

Author: xiaost

connection_impl.go

@@ -243,6 +243,9 @@ func (c *connection) ReadByte() (b byte, err error) {
 
 // Malloc implements Connection.
 func (c *connection) Malloc(n int) (buf []byte, err error) {
+	if !c.IsActive() {
+		return nil, Exception(ErrConnClosed, "when malloc")
+	}
 	return c.outputBuffer.Malloc(n)
 }
 
@@ -273,31 +276,49 @@ func (c *connection) Flush() error {
 
 // MallocAck implements Connection.
 func (c *connection) MallocAck(n int) (err error) {
+	if !c.IsActive() {
+		return Exception(ErrConnClosed, "when malloc ack")
+	}
 	return c.outputBuffer.MallocAck(n)
 }
 
 // Append implements Connection.
 func (c *connection) Append(w Writer) (err error) {
+	if !c.IsActive() {
+		return Exception(ErrConnClosed, "when append")
+	}
 	return c.outputBuffer.Append(w)
 }
 
 // WriteString implements Connection.
 func (c *connection) WriteString(s string) (n int, err error) {
+	if !c.IsActive() {
+		return 0, Exception(ErrConnClosed, "when write string")
+	}
 	return c.outputBuffer.WriteString(s)
 }
 
 // WriteBinary implements Connection.
 func (c *connection) WriteBinary(b []byte) (n int, err error) {
+	if !c.IsActive() {
+		return 0, Exception(ErrConnClosed, "when write binary")
+	}
 	return c.outputBuffer.WriteBinary(b)
 }
 
 // WriteDirect implements Connection.
 func (c *connection) WriteDirect(p []byte, remainCap int) (err error) {
+	if !c.IsActive() {
+		return Exception(ErrConnClosed, "when write direct")
+	}
 	return c.outputBuffer.WriteDirect(p, remainCap)
 }
 
 // WriteByte implements Connection.
 func (c *connection) WriteByte(b byte) (err error) {
+	if !c.IsActive() {
+		return Exception(ErrConnClosed, "when write byte")
+	}
 	return c.outputBuffer.WriteByte(b)
 }
 

connection_test.go

@@ -858,6 +858,45 @@ func TestConnectionServerClose(t *testing.T) {
 	wg.Wait()
 }
 
+func TestWriterAfterClose(t *testing.T) {
+	r, w := GetSysFdPairs()
+	rconn, wconn := &connection{}, &connection{}
+	rconn.init(&netFD{fd: r}, nil)
+	wconn.init(&netFD{fd: w}, nil)
+
+	err := wconn.Close()
+	MustNil(t, err)
+
+	for wconn.IsActive() {
+		runtime.Gosched()
+	}
+
+	methods := []struct {
+		name string
+		fn   func() error
+	}{
+		{"Malloc", func() error { _, err := wconn.Malloc(1); return err }},
+		{"MallocAck", func() error { return wconn.MallocAck(0) }},
+		{"WriteBinary", func() error { _, err := wconn.WriteBinary([]byte("hi")); return err }},
+		{"WriteString", func() error { _, err := wconn.WriteString("hi"); return err }},
+		{"WriteByte", func() error { return wconn.WriteByte('a') }},
+		{"WriteDirect", func() error { return wconn.WriteDirect([]byte("hi"), 0) }},
+		{"Flush", func() error { return wconn.Flush() }},
+	}
+	for _, tc := range methods {
+		t.Run(tc.name, func(t *testing.T) {
+			defer func() {
+				if r := recover(); r != nil {
+					t.Fatalf("Writer.%s panicked after Close: %v", tc.name, r)
+				}
+			}()
+			err := tc.fn()
+			Assert(t, err != nil, fmt.Sprintf("Writer.%s should return error after Close", tc.name))
+		})
+	}
+	rconn.Close()
+}
+
 func TestConnectionDailTimeoutAndClose(t *testing.T) {
 	ln := createTestTCPListener(t)
 	defer ln.Close()

mux/shard_queue.go

@@ -172,6 +172,9 @@ func (q *ShardQueue) foreach() {
 
 // deal is used to get deal of netpoll.Writer.
 func (q *ShardQueue) deal(gts []WriterGetter) {
+	if !q.conn.IsActive() {
+		return
+	}
 	writer := q.conn.Writer()
 	for _, gt := range gts {
 		buf, isNil := gt()

netpoll_unix_test.go

@@ -436,8 +436,6 @@ func TestServerReadAndClose(t *testing.T) {
 		runtime.Gosched() // wait for poller close connection
 	}
 	_, err = conn.Writer().WriteBinary(sendMsg)
-	MustNil(t, err)
-	err = conn.Writer().Flush()
 	Assert(t, errors.Is(err, ErrConnClosed), err)
 
 	err = loop.Shutdown(context.Background())