MCorange/wings
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Use the request context for cancelation, not a background context

Browse Source 
This also fixes an improperly written server deletion listener to look at the correct context cancelation.

Theoretically this should help address the issues in pterodactyl/panel#3596 but I'm not really sure how that happens, and theres no steps for reproduction.
This commit is contained in:
Dane Everitt
2021-09-12 11:14:00 -07:00
parent ee91224eb6
commit 09e1ba6f34
3 changed files with 41 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
router/router_server_ws.go
View File

@@ -12,6 +12,14 @@ import (
"github.com/pterodactyl/wings/router/websocket"
)
var expectedCloseCodes = []int{
ws.CloseGoingAway,
ws.CloseAbnormalClosure,
ws.CloseNormalClosure,
ws.CloseNoStatusReceived,
ws.CloseServiceRestart,
}
// Upgrades a connection to a websocket and passes events along between.
func getServerWebsocket(c *gin.Context) {
manager := middleware.ExtractManager(c)
@@ -24,8 +32,10 @@ func getServerWebsocket(c *gin.Context) {
defer handler.Connection.Close()
// Create a context that can be canceled when the user disconnects from this
// socket that will also cancel listeners running in separate threads.
ctx, cancel := context.WithCancel(context.Background())
// socket that will also cancel listeners running in separate threads. If the
// connection itself is terminated listeners using this context will also be
// closed.
ctx, cancel := context.WithCancel(c.Request.Context())
defer cancel()
// Track this open connection on the server so that we can close them all programmatically
@@ -33,22 +43,19 @@ func getServerWebsocket(c *gin.Context) {
s.Websockets().Push(handler.Uuid(), &cancel)
defer s.Websockets().Remove(handler.Uuid())
// Listen for the context being canceled and then close the websocket connection. This normally
// just happens because you're disconnecting from the socket in the browser, however in some
// cases we close the connections programmatically (e.g. deleting the server) and need to send
// a close message to the websocket so it disconnects.
go func(ctx context.Context, c *ws.Conn) {
ListenerLoop:
for {
select {
case <-ctx.Done():
handler.Connection.WriteControl(ws.CloseMessage, ws.FormatCloseMessage(ws.CloseGoingAway, "server deleted"), time.Now().Add(time.Second*5))
// A break right here without defining the specific loop would only break the select
// and not actually break the for loop, thus causing this routine to stick around forever.
break ListenerLoop
}
// If the server is deleted we need to send a close message to the connected client
// so that they disconnect since there will be no more events sent along. Listen for
// the request context being closed to break this loop, otherwise this routine will
// be left hanging in the background.
go func() {
select {
case <-ctx.Done():
break
case <-s.Context().Done():
handler.Connection.WriteControl(ws.CloseMessage, ws.FormatCloseMessage(ws.CloseGoingAway, "server deleted"), time.Now().Add(time.Second*5))
break
}
}(ctx, handler.Connection)
}()
go handler.ListenForServerEvents(ctx)
go handler.ListenForExpiration(ctx)
@@ -58,14 +65,7 @@ func getServerWebsocket(c *gin.Context) {
_, p, err := handler.Connection.ReadMessage()
if err != nil {
if !ws.IsCloseError(
err,
ws.CloseNormalClosure,
ws.CloseGoingAway,
ws.CloseNoStatusReceived,
ws.CloseServiceRestart,
ws.CloseAbnormalClosure,
) {
if ws.IsUnexpectedCloseError(err, expectedCloseCodes...) {
s.Log().WithField("error", err).Warn("error handling websocket message for server")
}
break