potential leak when write and tcp closed concurrently

[okhowang]

Summary

if client close tcp connection directly when server write data with background context.
the goroutine do Write may block forever, and succeed Write may block too.

Solution

DO NOT call Write method with context.Background,
AND DO Read in goroutine once connection established.

analysis

in Conn.writeFrame if real write to tcp failed.
in defer func, it wait closed or context.Done.
if there is no read goroutine, no closed event

	defer func() {
		if err != nil {
			select {
			case <-c.closed:
				err = c.closeErr
			case <-ctx.Done():
				err = ctx.Err()
			}
			c.close(err)
			err = fmt.Errorf("failed to write frame: %w", err)
		}
	}()

first stucked write stack may like

#	0x14837f8	nhooyr.io/websocket.(*Conn).writeFrame.func1+0x98								/go/pkg/mod/nhooyr.io/websocket@v1.8.7/write.go:276
#	0x148349a	nhooyr.io/websocket.(*Conn).writeFrame+0x6ba									/go/pkg/mod/nhooyr.io/websocket@v1.8.7/write.go:318
#	0x1481e5b	nhooyr.io/websocket.(*Conn).write+0x19b										/go/pkg/mod/nhooyr.io/websocket@v1.8.7/write.go:129
#	0x1481944	nhooyr.io/websocket.(*Conn).Write+0x24										/go/pkg/mod/nhooyr.io/websocket@v1.8.7/write.go:42

succeed stack may like

#	0x147d46c	nhooyr.io/websocket.(*mu).lock+0xac										/go/pkg/mod/nhooyr.io/websocket@v1.8.7/conn_notjs.go:238
#	0x1481fd6	nhooyr.io/websocket.(*msgWriterState).reset+0x36								/go/pkg/mod/nhooyr.io/websocket@v1.8.7/write.go:142
#	0x1481c2b	nhooyr.io/websocket.(*Conn).writer+0x2b										/go/pkg/mod/nhooyr.io/websocket@v1.8.7/write.go:111
#	0x1481d2c	nhooyr.io/websocket.(*Conn).write+0x6c										/go/pkg/mod/nhooyr.io/websocket@v1.8.7/write.go:122
#	0x1481944	nhooyr.io/websocket.(*Conn).Write+0x24										/go/pkg/mod/nhooyr.io/websocket@v1.8.7/write.go:42

reproduce code

package main

import (
	"context"
	"flag"
	ws "github.com/gorilla/websocket"
	"log"
	"net/http"
	_ "net/http/pprof"
	"nhooyr.io/websocket"
	"sync"
	"time"
)

var (
	concurrency = flag.Int("c", 1, "concurrency")
	wait        = flag.Int("w", 1, "wait milliseconds before close")
)

func startupServer(connWg *sync.WaitGroup) {
	http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
		connWg.Add(1)
		defer connWg.Done()
		c, err := websocket.Accept(w, r, &websocket.AcceptOptions{
			InsecureSkipVerify: true,
			CompressionMode:    websocket.CompressionDisabled,
		})
		if err != nil {
			log.Printf("websocket accept error: %v", err)
			return
		}
		var wg sync.WaitGroup
		// uncomment below for resolving the issue
		//wg.Add(1)
		//go func() { // read goroutine
		//	defer wg.Done()
		//	time.Sleep(time.Millisecond * time.Duration(*wait) * 2)
		//	for {
		//		_, _, err := c.Read(context.Background())
		//		if err != nil {
		//			return
		//		}
		//	}
		//}()
		for i := 0; i < *concurrency; i++ { // multi write goroutine
			wg.Add(1)
			go func() {
				defer wg.Done()
				for j := 0; j < 100; j++ {
					// NOTICE background may block forever
					if err := c.Write(context.Background(), websocket.MessageText, []byte("hello")); err != nil {
						return
					}
				}
			}()
		}
		wg.Wait() // wait all write goroutine end
		// do some read after write
		for {
			_, _, err := c.Read(context.Background())
			if err != nil {
				return
			}
		}
	})
	go http.ListenAndServe("127.0.0.1:40000", nil)
}

func main() {
	log.SetFlags(log.Lmicroseconds)
	flag.Parse()

	var connWg sync.WaitGroup
	startupServer(&connWg)

	c, _, err := ws.DefaultDialer.Dial("ws://127.0.0.1:40000/ws", nil)
	if err != nil {
		panic(err)
	}
	go func() {
		for {
			_, _, err := c.ReadMessage()
			if err != nil {
				log.Println("client, read error", err)
				return
			}
		}
	}()
	time.Sleep(time.Millisecond * time.Duration(*wait))
	log.Println("client close", c.Close())
	connWg.Wait() // wait for server end
}

[nhooyr]

Good catch!

[nhooyr]

This shouldn't have affected anyone using the library correctly as you'd want to use c.CloseRead if you're not reading from the connection to have the ctx for write cancel when read fails. But will still fix as yes it's not intended.

[nhooyr]

Fixed in dev.