TUN-5141: Make sure websocket pinger returns before streaming returns

This commit is contained in:
cthuang
2021-09-22 17:33:05 +01:00
parent f985ed567f
commit 6238fd9022
4 changed files with 95 additions and 13 deletions

View File

@@ -48,7 +48,12 @@ type tcpOverWSConnection struct {
}
func (wc *tcpOverWSConnection) Stream(ctx context.Context, tunnelConn io.ReadWriter, log *zerolog.Logger) {
wc.streamHandler(websocket.NewConn(ctx, tunnelConn, log), wc.conn, log)
wsCtx, cancel := context.WithCancel(ctx)
wsConn := websocket.NewConn(wsCtx, tunnelConn, log)
wc.streamHandler(wsConn, wc.conn, log)
cancel()
// Makes sure wsConn stops sending ping before terminating the stream
wsConn.WaitForShutdown()
}
func (wc *tcpOverWSConnection) Close() {
@@ -63,7 +68,12 @@ type socksProxyOverWSConnection struct {
}
func (sp *socksProxyOverWSConnection) Stream(ctx context.Context, tunnelConn io.ReadWriter, log *zerolog.Logger) {
socks.StreamNetHandler(websocket.NewConn(ctx, tunnelConn, log), sp.accessPolicy, log)
wsCtx, cancel := context.WithCancel(ctx)
wsConn := websocket.NewConn(wsCtx, tunnelConn, log)
socks.StreamNetHandler(wsConn, sp.accessPolicy, log)
cancel()
// Makes sure wsConn stops sending ping before terminating the stream
wsConn.WaitForShutdown()
}
func (sp *socksProxyOverWSConnection) Close() {