Close connected sockets when a server is deleted; closes pterodactyl/panel#2428

This commit is contained in:
Dane Everitt 2020-10-03 20:46:29 -07:00
parent e02c197585
commit 37e59e6928
No known key found for this signature in database
GPG Key ID: EEA66103B3D71F53
5 changed files with 101 additions and 0 deletions

View File

@ -207,6 +207,7 @@ func deleteServer(c *gin.Context) {
// Unsubscribe all of the event listeners. // Unsubscribe all of the event listeners.
s.Events().Destroy() s.Events().Destroy()
s.Throttler().StopTimer() s.Throttler().StopTimer()
s.Websockets().CancelAll()
// Destroy the environment; in Docker this will handle a running container and // Destroy the environment; in Docker this will handle a running container and
// forcibly terminate it before removing the container, so we do not need to handle // forcibly terminate it before removing the container, so we do not need to handle

View File

@ -6,6 +6,7 @@ import (
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
ws "github.com/gorilla/websocket" ws "github.com/gorilla/websocket"
"github.com/pterodactyl/wings/router/websocket" "github.com/pterodactyl/wings/router/websocket"
"time"
) )
// Upgrades a connection to a websocket and passes events along between. // Upgrades a connection to a websocket and passes events along between.
@ -23,6 +24,28 @@ func getServerWebsocket(c *gin.Context) {
ctx, cancel := context.WithCancel(context.Background()) ctx, cancel := context.WithCancel(context.Background())
defer cancel() defer cancel()
// Track this open connection on the server so that we can close them all programtically
// if the server is deleted.
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 programatically (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
}
}
}(ctx, handler.Connection)
go handler.ListenForServerEvents(ctx) go handler.ListenForServerEvents(ctx)
go handler.ListenForExpiration(ctx) go handler.ListenForExpiration(ctx)

View File

@ -34,9 +34,11 @@ const (
type Handler struct { type Handler struct {
sync.RWMutex sync.RWMutex
Connection *websocket.Conn Connection *websocket.Conn
jwt *tokens.WebsocketPayload `json:"-"` jwt *tokens.WebsocketPayload `json:"-"`
server *server.Server server *server.Server
uuid uuid.UUID
} }
var ( var (
@ -99,13 +101,23 @@ func GetHandler(s *server.Server, w http.ResponseWriter, r *http.Request) (*Hand
return nil, err return nil, err
} }
u, err := uuid.NewRandom()
if err != nil {
return nil, errors.WithStack(err)
}
return &Handler{ return &Handler{
Connection: conn, Connection: conn,
jwt: nil, jwt: nil,
server: s, server: s,
uuid: u,
}, nil }, nil
} }
func (h *Handler) Uuid() uuid.UUID {
return h.uuid
}
func (h *Handler) SendJson(v *Message) error { func (h *Handler) SendJson(v *Message) error {
// Do not send JSON down the line if the JWT on the connection is not valid! // Do not send JSON down the line if the JWT on the connection is not valid!
if err := h.TokenValid(); err != nil { if err := h.TokenValid(); err != nil {

View File

@ -54,6 +54,10 @@ type Server struct {
// The console throttler instance used to control outputs. // The console throttler instance used to control outputs.
throttler *ConsoleThrottler throttler *ConsoleThrottler
// Tracks open websocket connections for the server.
wsBag *WebsocketBag
wsBagLocker sync.Mutex
} }
type InstallerDetails struct { type InstallerDetails struct {

61
server/websockets.go Normal file
View File

@ -0,0 +1,61 @@
package server
import (
"context"
"github.com/google/uuid"
"sync"
)
type WebsocketBag struct {
mu sync.Mutex
conns map[uuid.UUID]*context.CancelFunc
}
// Returns the websocket bag which contains all of the currently open websocket connections
// for the server instance.
func (s *Server) Websockets() *WebsocketBag {
s.wsBagLocker.Lock()
defer s.wsBagLocker.Unlock()
if s.wsBag == nil {
s.wsBag = &WebsocketBag{}
}
return s.wsBag
}
// Adds a new websocket connection to the stack.
func (w *WebsocketBag) Push(u uuid.UUID, cancel *context.CancelFunc) {
w.mu.Lock()
defer w.mu.Unlock()
if w.conns == nil {
w.conns = make(map[uuid.UUID]*context.CancelFunc)
}
w.conns[u] = cancel
}
// Removes a connection from the stack.
func (w *WebsocketBag) Remove(u uuid.UUID) {
w.mu.Lock()
delete(w.conns, u)
w.mu.Unlock()
}
// Cancels all of the stored cancel functions which has the effect of disconnecting
// every listening websocket for the server.
func (w *WebsocketBag) CancelAll() {
w.mu.Lock()
w.mu.Unlock()
if w.conns != nil {
for _, cancel := range w.conns {
c := *cancel
c()
}
}
// Reset the connections.
w.conns = make(map[uuid.UUID]*context.CancelFunc)
}