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

server/crash.go

@@ -82,5 +82,5 @@ func (s *Server) handleServerCrash() error {
 
 	s.crasher.SetLastCrash(time.Now())
 
-	return s.Environment.Start()
+	return s.HandlePowerAction(PowerActionStart)
 }

server/environment.go

@@ -36,7 +36,6 @@ type Environment interface {
 	// unnecessary double/triple/quad looping issues if multiple people press restart or spam the
 	// button to restart.
 	Restart() error
-	IsRestarting() bool
 
 	// Waits for a server instance to stop gracefully. If the server is still detected
 	// as running after seconds, an error will be returned, or the server will be terminated

server/environment_docker.go

@@ -342,13 +342,13 @@ func (d *DockerEnvironment) acquireRestartLock() error {
 // start command. This will return an error if there is already a restart process executing for the
 // server. The lock is released when the process is stopped and a start has begun.
 func (d *DockerEnvironment) Restart() error {
-	d.Server.Log().Debug("attempting to acquire restart lock...")
+	d.Server.Log().Debug("acquiring process restart lock...")
 	if err := d.acquireRestartLock(); err != nil {
 		d.Server.Log().Warn("failed to acquire restart lock; already acquired by a different process")
 		return err
 	}
 
-	d.Server.Log().Debug("acquired restart lock")
+	d.Server.Log().Info("acquired process lock; beginning restart process...")
 
 	err := d.WaitForStop(60, false)
 	if err != nil {
@@ -496,12 +496,14 @@ func (d *DockerEnvironment) Attach() error {
 	}
 
 	var err error
+	d.Lock()
 	d.stream, err = d.Client.ContainerAttach(context.Background(), d.Server.Id(), types.ContainerAttachOptions{
 		Stdin:  true,
 		Stdout: true,
 		Stderr: true,
 		Stream: true,
 	})
+	d.Unlock()
 
 	if err != nil {
 		return errors.WithStack(err)

server/power.go

@@ -1,12 +1,80 @@
 package server
 
-type PowerAction struct {
-	Action string `json:"action"`
+import (
+	"context"
+	"github.com/pkg/errors"
+	"golang.org/x/sync/semaphore"
+	"os"
+	"time"
+)
+
+type PowerAction string
+
+// The power actions that can be performed for a given server. This taps into the given server
+// environment and performs them in a way that prevents a race condition from occurring. For
+// example, sending two "start" actions back to back will not process the second action until
+// the first action has been completed.
+//
+// This utilizes a workerpool with a limit of one worker so that all of the actions execute
+// in a sync manner.
+const (
+	PowerActionStart     = "start"
+	PowerActionStop      = "stop"
+	PowerActionRestart   = "restart"
+	PowerActionTerminate = "kill"
+)
+
+// Checks if the power action being received is valid.
+func (pa PowerAction) IsValid() bool {
+	return pa == PowerActionStart ||
+		pa == PowerActionStop ||
+		pa == PowerActionTerminate ||
+		pa == PowerActionRestart
 }
 
-func (pr *PowerAction) IsValid() bool {
-	return pr.Action == "start" ||
-		pr.Action == "stop" ||
-		pr.Action == "kill" ||
-		pr.Action == "restart"
+// Helper function that can receive a power action and then process the actions that need
+// to occur for it. This guards against someone calling Start() twice at the same time, or
+// trying to restart while another restart process is currently running.
+//
+// However, the code design for the daemon does depend on the user correctly calling this
+// function rather than making direct calls to the start/stop/restart functions on the
+// environment struct.
+func (s *Server) HandlePowerAction(action PowerAction, waitSeconds ...int) error {
+	if s.powerLock == nil {
+		s.powerLock = semaphore.NewWeighted(1)
+	}
+
+	// Determines if we should wait for the lock or not. If a value greater than 0 is passed
+	// into this function we will wait that long for a lock to be acquired.
+	if len(waitSeconds) > 0 && waitSeconds[0] != 0 {
+		ctx, _ := context.WithTimeout(context.Background(), time.Second*time.Duration(waitSeconds[0]))
+		// Attempt to acquire a lock on the power action lock for up to 30 seconds. If more
+		// time than that passes an error will be propagated back up the chain and this
+		// request will be aborted.
+		if err := s.powerLock.Acquire(ctx, 1); err != nil {
+			return errors.WithMessage(err, "could not acquire lock on power state")
+		}
+	} else {
+		// If no wait duration was provided we will attempt to immediately acquire the lock
+		// and bail out with a context deadline error if it is not acquired immediately.
+		if ok := s.powerLock.TryAcquire(1); !ok {
+			return errors.WithMessage(context.DeadlineExceeded, "could not acquire lock on power state")
+		}
+	}
+
+	// Release the lock once the process being requested has finished executing.
+	defer s.powerLock.Release(1)
+
+	switch action {
+	case PowerActionStart:
+		return s.Environment.Start()
+	case PowerActionStop:
+		return s.Environment.Stop()
+	case PowerActionRestart:
+		return s.Environment.Restart()
+	case PowerActionTerminate:
+		return s.Environment.Terminate(os.Kill)
+	}
+
+	return errors.New("attempting to handle unknown power action")
 }

server/server.go

@@ -8,7 +8,6 @@ import (
 	"github.com/pkg/errors"
 	"github.com/pterodactyl/wings/api"
 	"golang.org/x/sync/semaphore"
-	"os"
 	"strings"
 	"sync"
 	"time"
@@ -20,6 +19,7 @@ type Server struct {
 	// writing the configuration to the disk.
 	sync.RWMutex
 	emitterLock sync.Mutex
+	powerLock   *semaphore.Weighted
 
 	// Maintains the configuration for the server. This is the data that gets returned by the Panel
 	// such as build settings and container images.
@@ -158,23 +158,6 @@ func (s *Server) GetProcessConfiguration() (*api.ServerConfigurationResponse, *a
 	return api.NewRequester().GetServerConfiguration(s.Id())
 }
 
-// Helper function that can receive a power action and then process the
-// actions that need to occur for it.
-func (s *Server) HandlePowerAction(action PowerAction) error {
-	switch action.Action {
-	case "start":
-		return s.Environment.Start()
-	case "restart":
-		return s.Environment.Restart()
-	case "stop":
-		return s.Environment.Stop()
-	case "kill":
-		return s.Environment.Terminate(os.Kill)
-	default:
-		return errors.New("an invalid power action was provided")
-	}
-}
-
 // Checks if the server is marked as being suspended or not on the system.
 func (s *Server) IsSuspended() bool {
 	return s.Config().Suspended