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) {
|
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,15 +140,11 @@ func main() {
|
||||||
|
|
||||||
http.HandleFunc("/ws", s.handleWS)
|
http.HandleFunc("/ws", s.handleWS)
|
||||||
|
|
||||||
go func() {
|
log.Printf("listening on %s", *addr)
|
||||||
log.Printf("listening on %s (token=%q)", *addr, *token)
|
|
||||||
log.Printf("connect with: ws://localhost%s/ws", *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):")
|
fmt.Println("Type messages to send to connected clients (Ctrl+C to quit):")
|
||||||
|
|
||||||
|
go func() {
|
||||||
scanner := bufio.NewScanner(os.Stdin)
|
scanner := bufio.NewScanner(os.Stdin)
|
||||||
for scanner.Scan() {
|
for scanner.Scan() {
|
||||||
line := strings.TrimSpace(scanner.Text())
|
line := strings.TrimSpace(scanner.Text())
|
||||||
|
|
@ -159,4 +154,7 @@ func main() {
|
||||||
s.broadcast(line)
|
s.broadcast(line)
|
||||||
log.Printf("[server] sent: %s", line)
|
log.Printf("[server] sent: %s", line)
|
||||||
}
|
}
|
||||||
|
}()
|
||||||
|
|
||||||
|
log.Fatal(http.ListenAndServe(*addr, nil))
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -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()
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue