From 471886dd34fa6f620f614feed8c001de1813a704 Mon Sep 17 00:00:00 2001
From: Matthew Penner <me@matthewp.io>
Date: Fri, 12 Mar 2021 16:19:35 -0700
Subject: [PATCH] internally mark if a server is restoring to restrict actions

---
 router/router_server_backup.go | 20 +++++++++++++++++---
 server/errors.go               |  1 +
 server/install.go              |  8 ++++++++
 server/power.go                |  4 ++++
 server/server.go               |  2 ++
 5 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/router/router_server_backup.go b/router/router_server_backup.go
index b3e4ad8..1cbf881 100644
--- a/router/router_server_backup.go
+++ b/router/router_server_backup.go
@@ -63,7 +63,7 @@ func postServerBackup(c *gin.Context) {
 // This endpoint will block until the backup is fully restored allowing for a
 // spinner to be displayed in the Panel UI effectively.
 //
-// TODO: stop the server if it is running; internally mark it as suspended
+// TODO: stop the server if it is running
 func postServerRestoreBackup(c *gin.Context) {
 	s := middleware.ExtractServer(c)
 	client := middleware.ExtractApiClient(c)
@@ -84,9 +84,19 @@ func postServerRestoreBackup(c *gin.Context) {
 		return
 	}
 
+	s.SetRestoring(true)
+	hasError := true
+	defer func() {
+		if !hasError {
+			return
+		}
+
+		s.SetRestoring(false)
+	}()
+
 	logger.Info("processing server backup restore request")
 	if data.TruncateDirectory {
-		logger.Info(`recieved "truncate_directory" flag in request: deleting server files`)
+		logger.Info("received \"truncate_directory\" flag in request: deleting server files")
 		if err := s.Filesystem().TruncateRootDirectory(); err != nil {
 			middleware.CaptureAndAbort(c, err)
 			return
@@ -109,7 +119,9 @@ func postServerRestoreBackup(c *gin.Context) {
 			s.Events().Publish(server.DaemonMessageEvent, "Completed server restoration from local backup.")
 			s.Events().Publish(server.BackupRestoreCompletedEvent, "")
 			logger.Info("completed server restoration from local backup")
+			s.SetRestoring(false)
 		}(s, b, logger)
+		hasError = false
 		c.Status(http.StatusAccepted)
 		return
 	}
@@ -136,7 +148,7 @@ func postServerRestoreBackup(c *gin.Context) {
 	}
 	// Don't allow content types that we know are going to give us problems.
 	if res.Header.Get("Content-Type") == "" || !strings.Contains("application/x-gzip application/gzip", res.Header.Get("Content-Type")) {
-		res.Body.Close()
+		_ = res.Body.Close()
 		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{
 			"error": "The provided backup link is not a supported content type. \"" + res.Header.Get("Content-Type") + "\" is not application/x-gzip.",
 		})
@@ -151,8 +163,10 @@ func postServerRestoreBackup(c *gin.Context) {
 		s.Events().Publish(server.DaemonMessageEvent, "Completed server restoration from S3 backup.")
 		s.Events().Publish(server.BackupRestoreCompletedEvent, "")
 		logger.Info("completed server restoration from S3 backup")
+		s.SetRestoring(false)
 	}(s, c.Param("backup"), logger)
 
+	hasError = false
 	c.Status(http.StatusAccepted)
 }
 
diff --git a/server/errors.go b/server/errors.go
index ffdc7fa..cd14314 100644
--- a/server/errors.go
+++ b/server/errors.go
@@ -9,6 +9,7 @@ var (
 	ErrSuspended            = errors.New("server is currently in a suspended state")
 	ErrServerIsInstalling   = errors.New("server is currently installing")
 	ErrServerIsTransferring = errors.New("server is currently being transferred")
+	ErrServerIsRestoring    = errors.New("server is currently being restored")
 )
 
 type crashTooFrequent struct {
diff --git a/server/install.go b/server/install.go
index 1d54e99..77f1d16 100644
--- a/server/install.go
+++ b/server/install.go
@@ -151,6 +151,14 @@ func (s *Server) SetTransferring(state bool) {
 	s.transferring.Store(state)
 }
 
+func (s *Server) IsRestoring() bool {
+	return s.restoring.Load()
+}
+
+func (s *Server) SetRestoring(state bool) {
+	s.restoring.Store(state)
+}
+
 // Removes the installer container for the server.
 func (ip *InstallationProcess) RemoveContainer() error {
 	err := ip.client.ContainerRemove(ip.context, ip.Server.Id()+"_installer", types.ContainerRemoveOptions{
diff --git a/server/power.go b/server/power.go
index e587c8a..69486ff 100644
--- a/server/power.go
+++ b/server/power.go
@@ -70,6 +70,10 @@ func (s *Server) HandlePowerAction(action PowerAction, waitSeconds ...int) error
 		return ErrServerIsTransferring
 	}
 
+	if s.IsRestoring() {
+		return ErrServerIsRestoring
+	}
+
 	if s.powerLock == nil {
 		s.powerLock = semaphore.NewWeighted(1)
 	}
diff --git a/server/server.go b/server/server.go
index 895bb8e..0bbb336 100644
--- a/server/server.go
+++ b/server/server.go
@@ -60,6 +60,7 @@ type Server struct {
 	// installer process is still running.
 	installing   *system.AtomicBool
 	transferring *system.AtomicBool
+	restoring    *system.AtomicBool
 
 	// The console throttler instance used to control outputs.
 	throttler *ConsoleThrottler
@@ -79,6 +80,7 @@ func New(client remote.Client) (*Server, error) {
 		client:       client,
 		installing:   system.NewAtomicBool(false),
 		transferring: system.NewAtomicBool(false),
+		restoring:    system.NewAtomicBool(false),
 	}
 	if err := defaults.Set(&s); err != nil {
 		return nil, errors.Wrap(err, "server: could not set default values for struct")