fix: merge upstream fixes

Author: naterfute

src/router/router_server.go

@@ -323,7 +323,8 @@ func destroyServerBackupRepositories(ctx context.Context, s *server.Server, clie
 // Adds any of the JTIs passed through in the body to the deny list for the websocket
 // preventing any JWT generated before the current time from being used to connect to
 // the socket or send along commands.
-// @deprecated superceded by /api/revoke
+//
+// deprecated: prefer /api/deauthorize-user
 func postServerDenyWSTokens(c *gin.Context) {
 	var data struct {
 		JTIs []string `json:"jtis"`

src/router/router_server_ws.go

@@ -2,14 +2,20 @@ package router
 
 import (
 	"context"
+	"emperror.dev/errors"
 	"encoding/json"
+
+	"net/http"
 	"time"
 
 	"github.com/gin-gonic/gin"
 	ws "github.com/gorilla/websocket"
 
 	"github.com/pyrohost/elytra/src/router/middleware"
 	"github.com/pyrohost/elytra/src/router/websocket"
+
+	"github.com/pyrohost/elytra/src/server"
+	"golang.org/x/time/rate"
 )
 
 var expectedCloseCodes = []int{
@@ -25,6 +31,27 @@ func getServerWebsocket(c *gin.Context) {
 	manager := middleware.ExtractManager(c)
 	s, _ := manager.Get(c.Param("server"))
 
+	// Limit the total number of websockets that can be opened at any one time for
+	// a server instance. This applies across all users connected to the server, and
+	// is not applied on a per-user basis.
+	//
+	// todo: it would be great to make this per-user instead, but we need to modify
+	//  how we even request this endpoint in order for that to be possible. Some type
+	//  of signed identifier in the URL that is verified on this end and set by the
+	//  panel using a shared secret is likely the easiest option. The benefit of that
+	//  is that we can both scope things to the user before authentication, and also
+	//  verify that the JWT provided by the panel is assigned to the same user.
+	if s.Websockets().Len() >= 30 {
+		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{
+			"error": "Too many open websocket connections.",
+		})
+
+		return
+	}
+
+	c.Header("Content-Security-Policy", "default-src 'self'")
+	c.Header("X-Frame-Options", "DENY")
+
 	// Create a context that can be canceled when the user disconnects from this
 	// socket that will also cancel listeners running in separate threads. If the
 	// connection itself is terminated listeners using this context will also be
@@ -37,53 +64,101 @@ func getServerWebsocket(c *gin.Context) {
 		middleware.CaptureAndAbort(c, err)
 		return
 	}
-	defer handler.Connection.Close()
 
 	// Track this open connection on the server so that we can close them all programmatically
 	// if the server is deleted.
 	s.Websockets().Push(handler.Uuid(), &cancel)
 	handler.Logger().Debug("opening connection to server websocket")
+	defer s.Websockets().Remove(handler.Uuid())
 
-	defer func() {
-		s.Websockets().Remove(handler.Uuid())
-		handler.Logger().Debug("closing connection to server websocket")
+	go func() {
+		select {
+		// When the main context is canceled (through disconnect, server deletion, or server
+		// suspension) close the connection itself.
+		case <-ctx.Done():
+			handler.Logger().Debug("closing connection to server websocket")
+			if err := handler.Connection.Close(); err != nil {
+				handler.Logger().WithError(err).Error("failed to close websocket connection")
+			}
+			break
+		}
 	}()
 
-	// 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
+			return
+		// 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.
 		case <-s.Context().Done():
-			_ = handler.Connection.WriteControl(ws.CloseMessage, ws.FormatCloseMessage(ws.CloseGoingAway, "server deleted"), time.Now().Add(time.Second*5))
+			cancel()
 			break
 		}
 	}()
 
-	for {
-		j := websocket.Message{}
+	// Due to how websockets are handled we need to connect to the socket
+	// and _then_ abort it if the server is suspended. You cannot capture
+	// the HTTP response in the websocket client, thus we connect and then
+	// immediately close with failure.
+	if s.IsSuspended() {
+		_ = handler.Connection.WriteMessage(ws.CloseMessage, ws.FormatCloseMessage(4409, "server is suspended"))
+
+		return
+	}
 
-		_, p, err := handler.Connection.ReadMessage()
+	// There is a separate rate limiter that applies to individual message types
+	// within the actual websocket logic handler. _This_ rate limiter just exists
+	// to avoid enormous floods of data through the socket since we need to parse
+	// JSON each time. This rate limit realistically should never be hit since this
+	// would require sending 50+ messages a second over the websocket (no more than
+	// 10 per 200ms).
+	var throttled bool
+	rl := rate.NewLimiter(rate.Every(time.Millisecond*200), 10)
+
+	for {
+		t, p, err := handler.Connection.ReadMessage()
 		if err != nil {
 			if ws.IsUnexpectedCloseError(err, expectedCloseCodes...) {
 				handler.Logger().WithField("error", err).Warn("error handling websocket message for server")
 			}
 			break
 		}
 
+		if !rl.Allow() {
+			if !throttled {
+				throttled = true
+				_ = handler.Connection.WriteJSON(websocket.Message{Event: websocket.ThrottledEvent, Args: []string{"global"}})
+			}
+			continue
+		}
+
+		throttled = false
+
+		// If the message isn't a format we expect, or the length of the message is far larger
+		// than we'd ever expect, drop it. The websocket upgrader logic does enforce a maximum
+		// _compressed_ message size of 4Kb but that could decompress to a much larger amount
+		// of data.
+		if t != ws.TextMessage || len(p) > 32_768 {
+			continue
+		}
+
 		// Discard and JSON parse errors into the void and don't continue processing this
 		// specific socket request. If we did a break here the client would get disconnected
 		// from the socket, which is NOT what we want to do.
+		var j websocket.Message
 		if err := json.Unmarshal(p, &j); err != nil {
 			continue
 		}
 
 		go func(msg websocket.Message) {
 			if err := handler.HandleInbound(ctx, msg); err != nil {
-				_ = handler.SendErrorJson(msg, err)
+				if errors.Is(err, server.ErrSuspended) {
+					cancel()
+				} else {
+					_ = handler.SendErrorJson(msg, err)
+				}
 			}
 		}(j)
 	}

src/router/router_system.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/pyrohost/elytra/src/config"
 	"github.com/pyrohost/elytra/src/router/middleware"
+	"github.com/pyrohost/elytra/src/router/tokens"
 	"github.com/pyrohost/elytra/src/server"
 	"github.com/pyrohost/elytra/src/server/installer"
 	"github.com/pyrohost/elytra/src/system"
@@ -164,20 +165,24 @@ func postDeauthorizeUser(c *gin.Context) {
 		User    string   `json:"user"`
 		Servers []string `json:"servers"`
 	}
+
 	if err := c.BindJSON(&data); err != nil {
 		return
 	}
+
 	// todo: disconnect websockets more gracefully
 	m := middleware.ExtractManager(c)
 	if len(data.Servers) > 0 {
 		for _, uuid := range data.Servers {
 			if s, ok := m.Get(uuid); ok {
+				tokens.DenyForServer(s.ID(), data.User)
 				s.Websockets().CancelAll()
 				s.Sftp().Cancel(data.User)
 			}
 		}
 	} else {
 		for _, s := range m.All() {
+			tokens.DenyForServer(s.ID(), data.User)
 			s.Websockets().CancelAll()
 			s.Sftp().Cancel(data.User)
 		}

src/router/tokens/websocket.go

@@ -24,8 +24,20 @@ var elytraBootTime = time.Now()
 // This is used to allow the Panel to revoke tokens en-masse for a given user & server
 // combination since the JTI for tokens is just MD5(user.id + server.uuid). When a server
 // is booted this listing is fetched from the panel and the Websocket is dynamically updated.
+//
+// deprecated: prefer use of userDenylist
 var denylist sync.Map
 
+var userDenylist sync.Map
+
+// DenyForServer adds a user UUID to the denylist marking any existing JWTs issued
+// to the user as being invalid. This is associated with the user.
+func DenyForServer(s string, u string) {
+	log.WithField("user_uuid", u).WithField("server_uuid", s).Debugf("denying all JWTs created at or before current time for user \"%s\"", u)
+
+	userDenylist.Store(strings.Join([]string{s, u}, ":"), time.Now())
+}
+
 // Adds a JTI to the denylist by marking any JWTs generated before the current time as
 // being invalid if they use the same JTI.
 func DenyJTI(jti string) {
@@ -62,7 +74,7 @@ func (p *WebsocketPayload) GetServerUuid() string {
 }
 
 // Check if the JWT has been marked as denied by the instance due to either being issued
-// before Elytra was booted, or because we have denied all tokens with the same JTI
+// before Wings was booted, or because we have denied all tokens with the same JTI
 // occurring before a set time.
 func (p *WebsocketPayload) Denylisted() bool {
 	// If there is no IssuedAt present for the token, we cannot validate the token so
@@ -71,20 +83,29 @@ func (p *WebsocketPayload) Denylisted() bool {
 		return true
 	}
 
-	// If the time that the token was issued is before the time at which Elytra was booted
+	// If the time that the token was issued is before the time at which Wings was booted
 	// then the token is invalid for our purposes, even if the token "has permission".
 	if p.IssuedAt.Time.Before(elytraBootTime) {
 		return true
 	}
 
 	// Finally, if the token was issued before a time that is currently denied for this
 	// token instance, ignore the permissions response.
+	//
+	// This list is deprecated, but we maintain the check here so that custom instances
+	// are able to continue working. We'll remove it in a future release.
 	if t, ok := denylist.Load(p.JWTID); ok {
 		if p.IssuedAt.Time.Before(t.(time.Time)) {
 			return true
 		}
 	}
 
+	if t, ok := userDenylist.Load(strings.Join([]string{p.ServerUUID, p.UserUUID}, ":")); ok {
+		if p.IssuedAt.Time.Before(t.(time.Time)) {
+			return true
+		}
+	}
+
 	return false
 }
 

src/router/websocket/limiter.go

@@ -0,0 +1,91 @@
+package websocket
+
+import (
+	"sync"
+	"time"
+
+	"golang.org/x/time/rate"
+)
+
+type LimiterBucket struct {
+	mu        sync.RWMutex
+	limits    map[Event]*rate.Limiter
+	throttles map[Event]bool
+}
+
+func (h *Handler) IsThrottled(e Event) bool {
+	l := h.limiter.For(e)
+
+	h.limiter.mu.Lock()
+	defer h.limiter.mu.Unlock()
+
+	if l.Allow() {
+		h.limiter.throttles[e] = false
+
+		return false
+	}
+
+	// If not allowed, track the throttling and send an event over the wire
+	// if one wasn't already sent in the same throttling period.
+	if v, ok := h.limiter.throttles[e]; !v || !ok {
+		h.limiter.throttles[e] = true
+		h.Logger().WithField("event", e).Debug("throttling websocket due to event volume")
+
+		_ = h.unsafeSendJson(&Message{Event: ThrottledEvent, Args: []string{string(e)}})
+	}
+
+	return true
+}
+
+func NewLimiter() *LimiterBucket {
+	return &LimiterBucket{
+		limits:    make(map[Event]*rate.Limiter, 4),
+		throttles: make(map[Event]bool, 4),
+	}
+}
+
+// For returns the internal rate limiter for the given event type. In most
+// cases this is a shared rate limiter for events, but certain "heavy" or low-frequency
+// events implement their own limiters.
+func (l *LimiterBucket) For(e Event) *rate.Limiter {
+	name := limiterName(e)
+
+	l.mu.RLock()
+	if v, ok := l.limits[name]; ok {
+		l.mu.RUnlock()
+		return v
+	}
+
+	l.mu.RUnlock()
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	limit, burst := limitValuesFor(e)
+	l.limits[name] = rate.NewLimiter(limit, burst)
+
+	return l.limits[name]
+}
+
+// limitValuesFor returns the underlying limit and burst value for the given event.
+func limitValuesFor(e Event) (rate.Limit, int) {
+	// Twice every five seconds.
+	if e == AuthenticationEvent || e == SendServerLogsEvent {
+		return rate.Every(time.Second * 5), 2
+	}
+
+	// 10 per second.
+	if e == SendCommandEvent {
+		return rate.Every(time.Second), 10
+	}
+
+	// 4 per second.
+	return rate.Every(time.Second), 4
+}
+
+func limiterName(e Event) Event {
+	if e == AuthenticationEvent || e == SendServerLogsEvent || e == SendCommandEvent {
+		return e
+	}
+
+	return "_default"
+}

src/router/websocket/listeners.go

@@ -132,7 +132,7 @@ func (h *Handler) listenForServerEvents(ctx context.Context) error {
 				continue
 			}
 			var sendErr error
-			message := Message{Event: e.Topic}
+			message := Message{Event: Event(e.Topic)}
 			if str, ok := e.Data.(string); ok {
 				message.Args = []string{str}
 			} else if b, ok := e.Data.([]byte); ok {
@@ -150,7 +150,7 @@ func (h *Handler) listenForServerEvents(ctx context.Context) error {
 					continue
 				}
 			}
-			onError(message.Event, sendErr)
+			onError(string(message.Event), sendErr)
 		}
 		break
 	}