Wait for the power lock when stopping or killing a booting server

Fixes pterodactyl/panel#5712

router/router_server.go

@@ -86,7 +86,10 @@ func postServerPower(c *gin.Context) {
 	// we can immediately return a response from the server. Some of these actions
 	// can take quite some time, especially stopping or restarting.
 	go func(s *server.Server) {
-		if data.WaitSeconds < 0 || data.WaitSeconds > 300 {
+		// The response has already been sent, so an action that cannot get the power lock
+		// right away would be dropped without the caller knowing. Wait for the lock by default
+		// so that a stop or kill sent while the server is booting runs once the boot finishes.
+		if data.WaitSeconds <= 0 || data.WaitSeconds > 300 {
 			data.WaitSeconds = 30
 		}
 		if err := s.HandlePowerAction(data.Action, data.WaitSeconds); err != nil {

router/router_server_power_test.go

@@ -0,0 +1,68 @@
+package router
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/apex/log"
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/require"
+
+	"github.com/pterodactyl/wings/environment"
+	wserver "github.com/pterodactyl/wings/server"
+)
+
+// powerTestEnvironment reports each stop on stops and holds it until release is closed.
+type powerTestEnvironment struct {
+	environment.ProcessEnvironment
+	stops   chan struct{}
+	release chan struct{}
+}
+
+func (e *powerTestEnvironment) WaitForStop(context.Context, time.Duration, bool) error {
+	e.stops <- struct{}{}
+	<-e.release
+	return nil
+}
+
+// A stop sent without wait_seconds while another power action holds the lock runs
+// once the lock is released, instead of being dropped after the 202 response.
+func TestStopWithoutWaitSecondsWaitsForPowerLock(t *testing.T) {
+	s, err := wserver.New(nil)
+	require.NoError(t, err)
+	t.Cleanup(s.CtxCancel)
+	env := &powerTestEnvironment{stops: make(chan struct{}, 2), release: make(chan struct{})}
+	s.Environment = env
+
+	go func() { _ = s.HandlePowerAction(wserver.PowerActionStop) }()
+	<-env.stops
+
+	gin.SetMode(gin.TestMode)
+	w := httptest.NewRecorder()
+	c, _ := gin.CreateTestContext(w)
+	c.Request = httptest.NewRequest(http.MethodPost, "/api/servers/test/power", strings.NewReader(`{"action":"stop"}`))
+	c.Request.Header.Set("Content-Type", "application/json")
+	c.Set("server", s)
+	c.Set("logger", log.WithField("test", t.Name()))
+
+	postServerPower(c)
+	require.Equal(t, http.StatusAccepted, c.Writer.Status())
+
+	select {
+	case <-env.stops:
+		t.Fatal("the stop ran while another power action held the lock")
+	case <-time.After(100 * time.Millisecond):
+	}
+
+	close(env.release)
+
+	select {
+	case <-env.stops:
+	case <-time.After(2 * time.Second):
+		t.Fatal("the stop was dropped instead of waiting for the power lock")
+	}
+}

server/power.go

@@ -113,12 +113,27 @@ func (s *Server) HandlePowerAction(action PowerAction, waitSeconds ...int) error
 		//
 		// If we did successfully acquire the lock, make sure we release it once we're done
 		// executiong the power actions.
-		if err := s.powerLock.Acquire(); err == nil {
+		err := s.powerLock.Acquire()
+
+		// A lock held while the server is booting belongs to a start or restart that has not
+		// created the container yet, so there is nothing to kill. Wait for the boot like any
+		// other action instead of marking a server that is still starting as offline.
+		booting := errors.Is(err, system.ErrLockerLocked) && s.isBooting()
+		if booting && wait > 0 {
+			ctx, cancel := context.WithTimeout(s.ctx, time.Second*time.Duration(wait))
+			defer cancel()
+
+			err = s.powerLock.TryAcquire(ctx)
+		}
+
+		if err == nil {
 			log.Info("acquired exclusive lock on power actions, processing event...")
 			defer cleanup()
 		} else if errors.Is(err, system.ErrLockerDestroyed) {
 			// The server is being deleted.
 			return err
+		} else if booting {
+			return errors.Wrap(err, "failed to acquire exclusive lock for power actions while the server is starting")
 		} else {
 			log.Warn("failed to acquire exclusive lock, ignoring failure for termination event")
 		}
@@ -170,6 +185,14 @@ func (s *Server) HandlePowerAction(action PowerAction, waitSeconds ...int) error
 	return errors.New("attempting to handle unknown power action")
 }
 
+// isBooting reports whether the server process is offline or starting. When the power lock
+// is held in one of these states, a start or restart is running its pre-boot steps or
+// creating and starting the container.
+func (s *Server) isBooting() bool {
+	st := s.Environment.State()
+	return st == environment.ProcessOfflineState || st == environment.ProcessStartingState
+}
+
 // Execute a few functions before actually calling the environment start commands. This ensures
 // that everything is ready to go for environment booting, and that the server can even be started.
 func (s *Server) onBeforeStart() error {

server/power_boot_test.go

@@ -0,0 +1,93 @@
+package server
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"emperror.dev/errors"
+	"github.com/stretchr/testify/require"
+
+	"github.com/pterodactyl/wings/environment"
+	"github.com/pterodactyl/wings/system"
+)
+
+// bootTestEnvironment reports a fixed process state and records kill signals.
+type bootTestEnvironment struct {
+	environment.ProcessEnvironment
+	state      string
+	terminated chan struct{}
+}
+
+func (e *bootTestEnvironment) State() string { return e.state }
+
+func (e *bootTestEnvironment) Terminate(context.Context, string) error {
+	close(e.terminated)
+	return nil
+}
+
+func newBootTestServer(t *testing.T, state string) (*Server, *bootTestEnvironment) {
+	s, err := New(nil)
+	require.NoError(t, err)
+	t.Cleanup(s.CtxCancel)
+	env := &bootTestEnvironment{state: state, terminated: make(chan struct{})}
+	s.Environment = env
+	return s, env
+}
+
+// A kill sent while a boot holds the power lock runs once the boot releases it.
+func TestKillWhileBootingWaitsForBoot(t *testing.T) {
+	s, env := newBootTestServer(t, environment.ProcessStartingState)
+	require.NoError(t, s.powerLock.Acquire())
+
+	done := make(chan error, 1)
+	go func() { done <- s.HandlePowerAction(PowerActionTerminate, 5) }()
+
+	select {
+	case <-env.terminated:
+		t.Fatal("the server was killed before the boot released the power lock")
+	case <-time.After(100 * time.Millisecond):
+	}
+
+	s.powerLock.Release()
+
+	select {
+	case <-env.terminated:
+	case <-time.After(2 * time.Second):
+		t.Fatal("the server was not killed after the boot released the power lock")
+	}
+	require.NoError(t, <-done)
+	require.False(t, s.ExecutingPowerAction())
+}
+
+// A kill sent while a boot holds the power lock, without a wait, reports the lock
+// instead of marking the server offline.
+func TestKillWhileBootingWithoutWaitReturnsLockError(t *testing.T) {
+	s, env := newBootTestServer(t, environment.ProcessOfflineState)
+	require.NoError(t, s.powerLock.Acquire())
+
+	err := s.HandlePowerAction(PowerActionTerminate)
+	require.True(t, errors.Is(err, system.ErrLockerLocked), "got %v", err)
+
+	select {
+	case <-env.terminated:
+		t.Fatal("the server was killed while it was booting")
+	default:
+	}
+}
+
+// A kill still passes through the power lock held by a stop, so a server that is
+// stuck stopping can be killed.
+func TestKillWhileStoppingSkipsLock(t *testing.T) {
+	s, env := newBootTestServer(t, environment.ProcessStoppingState)
+	require.NoError(t, s.powerLock.Acquire())
+
+	require.NoError(t, s.HandlePowerAction(PowerActionTerminate, 5))
+
+	select {
+	case <-env.terminated:
+	default:
+		t.Fatal("the server was not killed while stopping")
+	}
+	require.True(t, s.ExecutingPowerAction())
+}