From 146681e1c798616518bc80082566e13b0e316853 Mon Sep 17 00:00:00 2001
From: andres-portainer <91705312+andres-portainer@users.noreply.github.com>
Date: Fri, 14 Jul 2023 12:34:50 -0300
Subject: [PATCH] fix(snapshots): increase the chance of taking a snapshot for
 edge environments EE-4795 (#9211)

---
 api/chisel/tunnel.go                          | 14 +++++--
 .../endpointedge_status_inspect.go            |  5 +++
 api/http/server.go                            |  3 ++
 api/internal/snapshot/snapshot.go             | 37 +++++++++++++++++++
 4 files changed, 55 insertions(+), 4 deletions(-)

diff --git a/api/chisel/tunnel.go b/api/chisel/tunnel.go
index 64e0afeb9..45e023da2 100644
--- a/api/chisel/tunnel.go
+++ b/api/chisel/tunnel.go
@@ -126,7 +126,10 @@ func (service *Service) SetTunnelStatusToIdle(endpointID portainer.EndpointID) {
 	credentials := tunnel.Credentials
 	if credentials != "" {
 		tunnel.Credentials = ""
-		service.chiselServer.DeleteUser(strings.Split(credentials, ":")[0])
+
+		if service.chiselServer != nil {
+			service.chiselServer.DeleteUser(strings.Split(credentials, ":")[0])
+		}
 	}
 
 	service.ProxyManager.DeleteEndpointProxy(endpointID)
@@ -161,9 +164,12 @@ func (service *Service) SetTunnelStatusToRequired(endpointID portainer.EndpointI
 
 		username, password := generateRandomCredentials()
 		authorizedRemote := fmt.Sprintf("^R:0.0.0.0:%d$", tunnel.Port)
-		err = service.chiselServer.AddUser(username, password, authorizedRemote)
-		if err != nil {
-			return err
+
+		if service.chiselServer != nil {
+			err = service.chiselServer.AddUser(username, password, authorizedRemote)
+			if err != nil {
+				return err
+			}
 		}
 
 		credentials, err := encryptCredentials(username, password, endpoint.EdgeID)
diff --git a/api/http/handler/endpointedge/endpointedge_status_inspect.go b/api/http/handler/endpointedge/endpointedge_status_inspect.go
index 04fc88030..54c126d58 100644
--- a/api/http/handler/endpointedge/endpointedge_status_inspect.go
+++ b/api/http/handler/endpointedge/endpointedge_status_inspect.go
@@ -138,6 +138,11 @@ func (handler *Handler) inspectStatus(tx dataservices.DataStoreTx, r *http.Reque
 		endpoint.EdgeID = edgeIdentifier
 	}
 
+	// Take an initial snapshot
+	if endpoint.LastCheckInDate == 0 {
+		handler.ReverseTunnelService.SetTunnelStatusToRequired(endpoint.ID)
+	}
+
 	agentPlatform, agentPlatformErr := parseAgentPlatform(r)
 	if agentPlatformErr != nil {
 		return nil, httperror.BadRequest("agent platform header is not valid", err)
diff --git a/api/http/server.go b/api/http/server.go
index 0bdd62ad9..e69dfef39 100644
--- a/api/http/server.go
+++ b/api/http/server.go
@@ -59,6 +59,7 @@ import (
 	"github.com/portainer/portainer/api/http/security"
 	"github.com/portainer/portainer/api/internal/authorization"
 	edgestackservice "github.com/portainer/portainer/api/internal/edge/edgestacks"
+	"github.com/portainer/portainer/api/internal/snapshot"
 	"github.com/portainer/portainer/api/internal/ssl"
 	"github.com/portainer/portainer/api/internal/upgrade"
 	k8s "github.com/portainer/portainer/api/kubernetes"
@@ -367,6 +368,8 @@ func (server *Server) Start() error {
 
 	go shutdown(server.ShutdownCtx, httpsServer)
 
+	go snapshot.NewBackgroundSnapshotter(server.DataStore, server.ReverseTunnelService)
+
 	return httpsServer.ListenAndServeTLS("", "")
 }
 
diff --git a/api/internal/snapshot/snapshot.go b/api/internal/snapshot/snapshot.go
index fdc57714b..8d5418889 100644
--- a/api/internal/snapshot/snapshot.go
+++ b/api/internal/snapshot/snapshot.go
@@ -10,6 +10,7 @@ import (
 	"github.com/portainer/portainer/api/agent"
 	"github.com/portainer/portainer/api/crypto"
 	"github.com/portainer/portainer/api/dataservices"
+	"github.com/portainer/portainer/api/internal/endpointutils"
 	"github.com/portainer/portainer/pkg/featureflags"
 
 	"github.com/rs/zerolog/log"
@@ -44,6 +45,42 @@ func NewService(snapshotIntervalFromFlag string, dataStore dataservices.DataStor
 	}, nil
 }
 
+// NewBackgroundSnapshotter queues snapshots of existing edge environments that
+// do not have one already
+func NewBackgroundSnapshotter(dataStore dataservices.DataStore, tunnelService portainer.ReverseTunnelService) {
+	var endpointIDs []portainer.EndpointID
+
+	err := dataStore.ViewTx(func(tx dataservices.DataStoreTx) error {
+		endpoints, err := tx.Endpoint().Endpoints()
+		if err != nil {
+			return err
+		}
+
+		for _, e := range endpoints {
+			if !endpointutils.IsEdgeEndpoint(&e) {
+				continue
+			}
+
+			s, err := tx.Snapshot().Read(e.ID)
+			if dataservices.IsErrObjectNotFound(err) ||
+				(err == nil && s.Docker == nil && s.Kubernetes == nil) {
+				endpointIDs = append(endpointIDs, e.ID)
+			}
+		}
+
+		return nil
+	})
+	if err != nil {
+		log.Error().Err(err).Msg("background snapshotter failure")
+		return
+	}
+
+	for _, endpointID := range endpointIDs {
+		tunnelService.SetTunnelStatusToActive(endpointID)
+		time.Sleep(10 * time.Second)
+	}
+}
+
 func parseSnapshotFrequency(snapshotInterval string, dataStore dataservices.DataStore) (float64, error) {
 	if snapshotInterval == "" {
 		settings, err := dataStore.Settings().Settings()