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:
parent
524fed7b8d
commit
bd17ca1c18
3 changed files with 34 additions and 23 deletions
|
|
@ -43,8 +43,7 @@ type server struct {
|
|||
func (s *server) handleWS(w http.ResponseWriter, r *http.Request) {
|
||||
if s.token != "" {
|
||||
auth := r.Header.Get("Authorization")
|
||||
qtoken := r.URL.Query().Get("token")
|
||||
if auth != "Bearer "+s.token && qtoken != s.token {
|
||||
if auth != "Bearer "+s.token {
|
||||
http.Error(w, "unauthorized", http.StatusUnauthorized)
|
||||
return
|
||||
}
|
||||
|
|
@ -141,15 +140,11 @@ func main() {
|
|||
|
||||
http.HandleFunc("/ws", s.handleWS)
|
||||
|
||||
go func() {
|
||||
log.Printf("listening on %s (token=%q)", *addr, *token)
|
||||
log.Printf("listening on %s", *addr)
|
||||
log.Printf("connect with: ws://localhost%s/ws", *addr)
|
||||
if err := http.ListenAndServe(*addr, nil); err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
}()
|
||||
|
||||
fmt.Println("Type messages to send to connected clients (Ctrl+C to quit):")
|
||||
|
||||
go func() {
|
||||
scanner := bufio.NewScanner(os.Stdin)
|
||||
for scanner.Scan() {
|
||||
line := strings.TrimSpace(scanner.Text())
|
||||
|
|
@ -159,4 +154,7 @@ func main() {
|
|||
s.broadcast(line)
|
||||
log.Printf("[server] sent: %s", line)
|
||||
}
|
||||
}()
|
||||
|
||||
log.Fatal(http.ListenAndServe(*addr, nil))
|
||||
}
|
||||
|
|
|
|||
|
|
@ -93,10 +93,13 @@ func (c *PicoClientChannel) dial() error {
|
|||
return err
|
||||
}
|
||||
|
||||
connCtx, connCancel := context.WithCancel(c.ctx)
|
||||
|
||||
pc := &picoConn{
|
||||
id: uuid.New().String(),
|
||||
conn: ws,
|
||||
sessionID: c.config.SessionID,
|
||||
cancel: connCancel,
|
||||
}
|
||||
if pc.sessionID == "" {
|
||||
pc.sessionID = uuid.New().String()
|
||||
|
|
@ -106,7 +109,7 @@ func (c *PicoClientChannel) dial() error {
|
|||
c.conn = pc
|
||||
c.mu.Unlock()
|
||||
|
||||
go c.readLoop(pc)
|
||||
go c.readLoop(connCtx, pc)
|
||||
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()
|
||||
|
||||
readTimeout := time.Duration(c.config.ReadTimeout) * time.Second
|
||||
|
|
@ -165,11 +168,11 @@ func (c *PicoClientChannel) readLoop(pc *picoConn) {
|
|||
if pingInterval <= 0 {
|
||||
pingInterval = 30 * time.Second
|
||||
}
|
||||
go c.pingLoop(pc, pingInterval)
|
||||
go c.pingLoop(connCtx, pc, pingInterval)
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-c.ctx.Done():
|
||||
case <-connCtx.Done():
|
||||
return
|
||||
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)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-c.ctx.Done():
|
||||
case <-connCtx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
if pc.closed.Load() {
|
||||
|
|
@ -303,8 +306,14 @@ func (c *PicoClientChannel) StartTyping(ctx context.Context, chatID string) (fun
|
|||
return func() {}, err
|
||||
}
|
||||
return func() {
|
||||
c.mu.Lock()
|
||||
currentPC := c.conn
|
||||
c.mu.Unlock()
|
||||
if currentPC == nil {
|
||||
return
|
||||
}
|
||||
stopMsg := newMessage(TypeTypingStop, nil)
|
||||
stopMsg.SessionID = strings.TrimPrefix(chatID, "pico_client:")
|
||||
pc.writeJSON(stopMsg)
|
||||
currentPC.writeJSON(stopMsg)
|
||||
}, nil
|
||||
}
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ type picoConn struct {
|
|||
sessionID string
|
||||
writeMu sync.Mutex
|
||||
closed atomic.Bool
|
||||
cancel context.CancelFunc // cancels per-connection goroutines (e.g. pingLoop)
|
||||
}
|
||||
|
||||
// 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.
|
||||
func (pc *picoConn) close() {
|
||||
if pc.closed.CompareAndSwap(false, true) {
|
||||
if pc.cancel != nil {
|
||||
pc.cancel()
|
||||
}
|
||||
pc.conn.Close()
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue