hoelee/cloudflared
1
0
Fork 0
mirror of https://github.com/cloudflare/cloudflared.git synced 2025-07-27 18:29:56 +00:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

Revert "TUN-5184: Make sure outstanding websocket write is finished, and no more writes after shutdown"

Browse Source 
This reverts commit f8fbbcd806.
This commit is contained in:
Nuno Diegues
2021-10-25 19:51:52 +01:00
parent f6f10305a6
commit 573d410606
6 changed files with 52 additions and 195 deletions
Download Patch File Download Diff File Expand all files Collapse all files

37
websocket/connection.go
View File

@@ -3,10 +3,8 @@ package websocket
import (
"bytes"
"context"
"errors"
"fmt"
"io"
"sync"
"time"
gobwas "github.com/gobwas/ws"
@@ -100,17 +98,15 @@ func (c *GorillaConn) pinger(ctx context.Context) {
type Conn struct {
rw io.ReadWriter
log *zerolog.Logger
// writeLock makes sure
// 1. Only one write at a time. The pinger and Stream function can both call write.
// 2. Close only returns after in progress Write is finished, and no more Write will succeed after calling Close.
writeLock sync.Mutex
done bool
// closed is a channel to indicate if Conn has been fully terminated
shutdownC chan struct{}
}
func NewConn(ctx context.Context, rw io.ReadWriter, log *zerolog.Logger) *Conn {
c := &Conn{
rw: rw,
log: log,
rw: rw,
log: log,
shutdownC: make(chan struct{}),
}
go c.pinger(ctx)
return c
@@ -125,22 +121,16 @@ func (c *Conn) Read(reader []byte) (int, error) {
return copy(reader, data), nil
}
// Write will write messages to the websocket connection.
// It will not write to the connection after Close is called to fix TUN-5184
// Write will write messages to the websocket connection
func (c *Conn) Write(p []byte) (int, error) {
c.writeLock.Lock()
defer c.writeLock.Unlock()
if c.done {
return 0, errors.New("Write to closed websocket connection")
}
if err := wsutil.WriteServerBinary(c.rw, p); err != nil {
return 0, err
}
return len(p), nil
}
func (c *Conn) pinger(ctx context.Context) {
defer close(c.shutdownC)
pongMessge := wsutil.Message{
OpCode: gobwas.OpPong,
Payload: []byte{},
@@ -150,12 +140,11 @@ func (c *Conn) pinger(ctx context.Context) {
defer ticker.Stop()
for {
select {
// Ping/Pong messages will not be written after the connection is done
case <-ticker.C:
if err := wsutil.WriteServerMessage(c, gobwas.OpPing, []byte{}); err != nil {
if err := wsutil.WriteServerMessage(c.rw, gobwas.OpPing, []byte{}); err != nil {
c.log.Debug().Err(err).Msgf("failed to write ping message")
}
if err := wsutil.HandleClientControlMessage(c, pongMessge); err != nil {
if err := wsutil.HandleClientControlMessage(c.rw, pongMessge); err != nil {
c.log.Debug().Err(err).Msgf("failed to write pong message")
}
case <-ctx.Done():
@@ -173,9 +162,7 @@ func (c *Conn) pingPeriod(ctx context.Context) time.Duration {
return defaultPingPeriod
}
// Close waits for the current write to finish. Further writes will return error
func (c *Conn) Close() {
c.writeLock.Lock()
defer c.writeLock.Unlock()
c.done = true
// Close waits for pinger to terminate
func (c *Conn) WaitForShutdown() {
<-c.shutdownC
}