fix(pico): address PR #1198 review — goroutine leak, race, auth

- Add per-connection context cancel to picoConn to prevent pingLoop
  goroutine leak on disconnect
- Re-acquire mutex in StartTyping stop closure to avoid stale conn race
- Remove query-param token auth from echo server (header-only)
- Move ListenAndServe to main goroutine where log.Fatal is safe

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Amir Mamaghani 2026-03-07 18:25:00 +01:00
parent 524fed7b8d
commit bd17ca1c18
3 changed files with 34 additions and 23 deletions

View file

@ -43,8 +43,7 @@ type server struct {
func (s *server) handleWS(w http.ResponseWriter, r *http.Request) { func (s *server) handleWS(w http.ResponseWriter, r *http.Request) {
if s.token != "" { if s.token != "" {
auth := r.Header.Get("Authorization") auth := r.Header.Get("Authorization")
qtoken := r.URL.Query().Get("token") if auth != "Bearer "+s.token {
if auth != "Bearer "+s.token && qtoken != s.token {
http.Error(w, "unauthorized", http.StatusUnauthorized) http.Error(w, "unauthorized", http.StatusUnauthorized)
return return
} }
@ -141,22 +140,21 @@ func main() {
http.HandleFunc("/ws", s.handleWS) http.HandleFunc("/ws", s.handleWS)
log.Printf("listening on %s", *addr)
log.Printf("connect with: ws://localhost%s/ws", *addr)
fmt.Println("Type messages to send to connected clients (Ctrl+C to quit):")
go func() { go func() {
log.Printf("listening on %s (token=%q)", *addr, *token) scanner := bufio.NewScanner(os.Stdin)
log.Printf("connect with: ws://localhost%s/ws", *addr) for scanner.Scan() {
if err := http.ListenAndServe(*addr, nil); err != nil { line := strings.TrimSpace(scanner.Text())
log.Fatal(err) if line == "" {
continue
}
s.broadcast(line)
log.Printf("[server] sent: %s", line)
} }
}() }()
fmt.Println("Type messages to send to connected clients (Ctrl+C to quit):") log.Fatal(http.ListenAndServe(*addr, nil))
scanner := bufio.NewScanner(os.Stdin)
for scanner.Scan() {
line := strings.TrimSpace(scanner.Text())
if line == "" {
continue
}
s.broadcast(line)
log.Printf("[server] sent: %s", line)
}
} }

View file

@ -93,10 +93,13 @@ func (c *PicoClientChannel) dial() error {
return err return err
} }
connCtx, connCancel := context.WithCancel(c.ctx)
pc := &picoConn{ pc := &picoConn{
id: uuid.New().String(), id: uuid.New().String(),
conn: ws, conn: ws,
sessionID: c.config.SessionID, sessionID: c.config.SessionID,
cancel: connCancel,
} }
if pc.sessionID == "" { if pc.sessionID == "" {
pc.sessionID = uuid.New().String() pc.sessionID = uuid.New().String()
@ -106,7 +109,7 @@ func (c *PicoClientChannel) dial() error {
c.conn = pc c.conn = pc
c.mu.Unlock() c.mu.Unlock()
go c.readLoop(pc) go c.readLoop(connCtx, pc)
return nil return nil
} }
@ -148,7 +151,7 @@ func (c *PicoClientChannel) reconnectLoop() {
} }
} }
func (c *PicoClientChannel) readLoop(pc *picoConn) { func (c *PicoClientChannel) readLoop(connCtx context.Context, pc *picoConn) {
defer pc.close() defer pc.close()
readTimeout := time.Duration(c.config.ReadTimeout) * time.Second readTimeout := time.Duration(c.config.ReadTimeout) * time.Second
@ -165,11 +168,11 @@ func (c *PicoClientChannel) readLoop(pc *picoConn) {
if pingInterval <= 0 { if pingInterval <= 0 {
pingInterval = 30 * time.Second pingInterval = 30 * time.Second
} }
go c.pingLoop(pc, pingInterval) go c.pingLoop(connCtx, pc, pingInterval)
for { for {
select { select {
case <-c.ctx.Done(): case <-connCtx.Done():
return return
default: default:
} }
@ -199,12 +202,12 @@ func (c *PicoClientChannel) readLoop(pc *picoConn) {
} }
} }
func (c *PicoClientChannel) pingLoop(pc *picoConn, interval time.Duration) { func (c *PicoClientChannel) pingLoop(connCtx context.Context, pc *picoConn, interval time.Duration) {
ticker := time.NewTicker(interval) ticker := time.NewTicker(interval)
defer ticker.Stop() defer ticker.Stop()
for { for {
select { select {
case <-c.ctx.Done(): case <-connCtx.Done():
return return
case <-ticker.C: case <-ticker.C:
if pc.closed.Load() { if pc.closed.Load() {
@ -303,8 +306,14 @@ func (c *PicoClientChannel) StartTyping(ctx context.Context, chatID string) (fun
return func() {}, err return func() {}, err
} }
return func() { return func() {
c.mu.Lock()
currentPC := c.conn
c.mu.Unlock()
if currentPC == nil {
return
}
stopMsg := newMessage(TypeTypingStop, nil) stopMsg := newMessage(TypeTypingStop, nil)
stopMsg.SessionID = strings.TrimPrefix(chatID, "pico_client:") stopMsg.SessionID = strings.TrimPrefix(chatID, "pico_client:")
pc.writeJSON(stopMsg) currentPC.writeJSON(stopMsg)
}, nil }, nil
} }

View file

@ -27,6 +27,7 @@ type picoConn struct {
sessionID string sessionID string
writeMu sync.Mutex writeMu sync.Mutex
closed atomic.Bool closed atomic.Bool
cancel context.CancelFunc // cancels per-connection goroutines (e.g. pingLoop)
} }
// writeJSON sends a JSON message to the connection with write locking. // writeJSON sends a JSON message to the connection with write locking.
@ -42,6 +43,9 @@ func (pc *picoConn) writeJSON(v any) error {
// close closes the connection. // close closes the connection.
func (pc *picoConn) close() { func (pc *picoConn) close() {
if pc.closed.CompareAndSwap(false, true) { if pc.closed.CompareAndSwap(false, true) {
if pc.cancel != nil {
pc.cancel()
}
pc.conn.Close() pc.conn.Close()
} }
} }