TUN-5184: Handle errors in bidrectional streaming (websocket#Stream) gracefully when 1 side has ended

This commit is contained in:
Nuno Diegues 2022-02-10 12:59:15 +00:00
parent db01127191
commit ff4cfeda0c
2 changed files with 54 additions and 30 deletions

View File

@ -78,16 +78,3 @@ func (sp *socksProxyOverWSConnection) Stream(ctx context.Context, tunnelConn io.
func (sp *socksProxyOverWSConnection) Close() { func (sp *socksProxyOverWSConnection) Close() {
} }
// wsProxyConnection represents a bidirectional stream for a websocket connection to the origin
type wsProxyConnection struct {
rwc io.ReadWriteCloser
}
func (conn *wsProxyConnection) Stream(ctx context.Context, tunnelConn io.ReadWriter, log *zerolog.Logger) {
websocket.Stream(tunnelConn, conn.rwc, log)
}
func (conn *wsProxyConnection) Close() {
conn.rwc.Close()
}

View File

@ -8,6 +8,7 @@ import (
"fmt" "fmt"
"io" "io"
"net/http" "net/http"
"sync/atomic"
"time" "time"
"github.com/gorilla/websocket" "github.com/gorilla/websocket"
@ -28,28 +29,64 @@ func NewResponseHeader(req *http.Request) http.Header {
return header return header
} }
type bidirectionalStreamStatus struct {
doneChan chan struct{}
anyDone uint32
}
func newBiStreamStatus() *bidirectionalStreamStatus {
return &bidirectionalStreamStatus{
doneChan: make(chan struct{}, 2),
anyDone: 0,
}
}
func (s *bidirectionalStreamStatus) markUniStreamDone() {
atomic.StoreUint32(&s.anyDone, 1)
s.doneChan <- struct{}{}
}
func (s *bidirectionalStreamStatus) waitAnyDone() {
<-s.doneChan
}
func (s *bidirectionalStreamStatus) isAnyDone() bool {
return atomic.LoadUint32(&s.anyDone) > 0
}
// Stream copies copy data to & from provided io.ReadWriters. // Stream copies copy data to & from provided io.ReadWriters.
func Stream(tunnelConn, originConn io.ReadWriter, log *zerolog.Logger) { func Stream(tunnelConn, originConn io.ReadWriter, log *zerolog.Logger) {
proxyDone := make(chan struct{}, 2) status := newBiStreamStatus()
go func() { go unidirectionalStream(tunnelConn, originConn, "origin->tunnel", status, log)
_, err := copyData(tunnelConn, originConn, "origin->tunnel") go unidirectionalStream(originConn, tunnelConn, "tunnel->origin", status, log)
if err != nil {
log.Debug().Msgf("origin to tunnel copy: %v", err)
}
proxyDone <- struct{}{}
}()
go func() {
_, err := copyData(originConn, tunnelConn, "tunnel->origin")
if err != nil {
log.Debug().Msgf("tunnel to origin copy: %v", err)
}
proxyDone <- struct{}{}
}()
// If one side is done, we are done. // If one side is done, we are done.
<-proxyDone status.waitAnyDone()
}
func unidirectionalStream(dst io.Writer, src io.Reader, dir string, status *bidirectionalStreamStatus, log *zerolog.Logger) {
defer func() {
// The bidirectional streaming spawns 2 goroutines to stream each direction.
// If any ends, the callstack returns, meaning the Tunnel request/stream (depending on http2 vs quic) will
// close. In such case, if the other direction did not stop (due to application level stopping, e.g., if a
// server/origin listens forever until closure), it may read/write from the underlying ReadWriter (backed by
// the Edge<->cloudflared transport) in an unexpected state.
if status.isAnyDone() {
// Because of this, we set this recover() logic, which kicks-in *only* if any stream is known to have
// exited. In such case, we stop a possible panic from propagating upstream.
if r := recover(); r != nil {
// We handle such unexpected errors only when we detect that one side of the streaming is done.
log.Debug().Msgf("Handled gracefully error %v in Streaming for %s", r, dir)
}
}
}()
_, err := copyData(dst, src, dir)
if err != nil {
log.Debug().Msgf("%s copy: %v", dir, err)
}
status.markUniStreamDone()
} }
// when set to true, enables logging of content copied to/from origin and tunnel // when set to true, enables logging of content copied to/from origin and tunnel