Refactor power handling logic to be more robust and able to handle spam clicking and duplicate power actions

router/router_server.go

@@ -2,6 +2,7 @@ package router
 
 import (
 	"bytes"
+	"context"
 	"github.com/apex/log"
 	"github.com/gin-gonic/gin"
 	"github.com/pkg/errors"
@@ -47,13 +48,15 @@ func getServerLogs(c *gin.Context) {
 func postServerPower(c *gin.Context) {
 	s := GetServer(c.Param("server"))
 
-	var data server.PowerAction
-	// BindJSON sends 400 if the request fails, all we need to do is return
+	var data struct{
+		Action server.PowerAction `json:"action"`
+	}
+
 	if err := c.BindJSON(&data); err != nil {
 		return
 	}
 
-	if !data.IsValid() {
+	if !data.Action.IsValid() {
 		c.AbortWithStatusJSON(http.StatusUnprocessableEntity, gin.H{
 			"error": "The power action provided was not valid, should be one of \"stop\", \"start\", \"restart\", \"kill\"",
 		})
@@ -66,7 +69,7 @@ func postServerPower(c *gin.Context) {
 	//
 	// We don't really care about any of the other actions at this point, they'll all result
 	// in the process being stopped, which should have happened anyways if the server is suspended.
-	if (data.Action == "start" || data.Action == "restart") && s.IsSuspended() {
+	if (data.Action == server.PowerActionStart || data.Action == server.PowerActionRestart) && s.IsSuspended() {
 		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{
 			"error": "Cannot start or restart a server that is suspended.",
 		})
@@ -76,10 +79,15 @@ func postServerPower(c *gin.Context) {
 	// Pass the actual heavy processing off to a seperate thread to handle so that
 	// we can immediately return a response from the server. Some of these actions
 	// can take quite some time, especially stopping or restarting.
-	go func(server *server.Server) {
-		if err := server.HandlePowerAction(data); err != nil {
-			server.Log().WithFields(log.Fields{"action": data, "error": err}).
-				Error("encountered error processing a server power action in the background")
+	go func(s *server.Server) {
+		if err := s.HandlePowerAction(data.Action, 30); err != nil {
+			if errors.Is(err, context.DeadlineExceeded) {
+				s.Log().WithField("action", data.Action).
+					Warn("could not acquire a lock while attempting to perform a power action")
+			} else {
+				s.Log().WithFields(log.Fields{"action": data, "error": err}).
+					Error("encountered error processing a server power action in the background")
+			}
 		}
 	}(s)
 

router/websocket/websocket.go

@@ -1,6 +1,7 @@
 package websocket
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"github.com/apex/log"
@@ -12,7 +13,6 @@ import (
 	"github.com/pterodactyl/wings/router/tokens"
 	"github.com/pterodactyl/wings/server"
 	"net/http"
-	"os"
 	"strings"
 	"sync"
 	"time"
@@ -150,7 +150,7 @@ func (h *Handler) TokenValid() error {
 // Sends an error back to the connected websocket instance by checking the permissions
 // of the token. If the user has the "receive-errors" grant we will send back the actual
 // error message, otherwise we just send back a standard error message.
-func (h *Handler) SendErrorJson(msg Message, err error) error {
+func (h *Handler) SendErrorJson(msg Message, err error, shouldLog ...bool) error {
 	j := h.GetJwt()
 
 	message := "an unexpected error was encountered while handling this request"
@@ -163,9 +163,11 @@ func (h *Handler) SendErrorJson(msg Message, err error) error {
 	wsm := Message{Event: ErrorEvent}
 	wsm.Args = []string{m}
 
-	if !server.IsSuspendedError(err) {
-		h.server.Log().WithFields(log.Fields{"event": msg.Event, "error_identifier": u.String(), "error": err}).
-			Error("failed to handle websocket process; an error was encountered processing an event")
+	if len(shouldLog) == 0 || (len(shouldLog) == 1 && shouldLog[0] == true) {
+		if !server.IsSuspendedError(err) {
+			h.server.Log().WithFields(log.Fields{"event": msg.Event, "error_identifier": u.String(), "error": err}).
+				Error("failed to handle websocket process; an error was encountered processing an event")
+		}
 	}
 
 	return h.unsafeSendJson(wsm)
@@ -271,37 +273,34 @@ func (h *Handler) HandleInbound(m Message) error {
 		}
 	case SetStateEvent:
 		{
-			switch strings.Join(m.Args, "") {
-			case "start":
-				if h.GetJwt().HasPermission(PermissionSendPowerStart) {
-					return h.server.Environment.Start()
-				}
-				break
-			case "stop":
-				if h.GetJwt().HasPermission(PermissionSendPowerStop) {
-					return h.server.Environment.Stop()
-				}
-				break
-			case "restart":
-				if h.GetJwt().HasPermission(PermissionSendPowerRestart) {
-					// If the server is alreay restarting don't do anything. Perhaps we send back an event
-					// in the future for this? For now no reason to knowingly trigger an error by trying to
-					// restart a process already restarting.
-					if h.server.Environment.IsRestarting() {
-						return nil
-					}
+			action := server.PowerAction(strings.Join(m.Args, ""))
 
-					return h.server.Environment.Restart()
+			actions := make(map[server.PowerAction]string)
+			actions[server.PowerActionStart] = PermissionSendPowerStart
+			actions[server.PowerActionStop] = PermissionSendPowerStop
+			actions[server.PowerActionRestart] = PermissionSendPowerRestart
+			actions[server.PowerActionTerminate] = PermissionSendPowerStop
+
+			// Check that they have permission to perform this action if it is needed.
+			if permission, exists := actions[action]; exists {
+				if !h.GetJwt().HasPermission(permission) {
+					return nil
 				}
-				break
-			case "kill":
-				if h.GetJwt().HasPermission(PermissionSendPowerStop) {
-					return h.server.Environment.Terminate(os.Kill)
-				}
-				break
 			}
 
-			return nil
+			err := h.server.HandlePowerAction(action)
+			if errors.Is(err, context.DeadlineExceeded) {
+				m, _ := h.GetErrorMessage("another power action is currently being processed for this server, please try again later")
+
+				h.SendJson(&Message{
+					Event: ErrorEvent,
+					Args:  []string{m},
+				})
+
+				return nil
+			}
+
+			return err
 		}
 	case SendServerLogsEvent:
 		{