From e97d82b7439268a70e44d1006d6a86b8e1e17e34 Mon Sep 17 00:00:00 2001
From: Paul Lorenz <paul.lorenz@netfoundry.io>
Date: Mon, 26 Aug 2024 19:59:01 -0400
Subject: [PATCH] Make ziti agent snapshot-db match ziti fabric db snapshot.
 Fixes #2333

---
 controller/agent.go                    | 13 +++--
 ziti/cmd/agentcli/agent.go             |  2 +-
 ziti/cmd/agentcli/agent_snapshot_db.go | 71 ++++++++++++++++++++++++++
 3 files changed, 80 insertions(+), 6 deletions(-)
 create mode 100644 ziti/cmd/agentcli/agent_snapshot_db.go

diff --git a/controller/agent.go b/controller/agent.go
index 2b26dc504..f2b841fc5 100644
--- a/controller/agent.go
+++ b/controller/agent.go
@@ -19,9 +19,10 @@ import (
 const (
 	AgentAppId byte = 1
 
-	AgentIdHeader      = 10
-	AgentAddrHeader    = 11
-	AgentIsVoterHeader = 12
+	AgentIdHeader         = 10
+	AgentAddrHeader       = 11
+	AgentIsVoterHeader    = 12
+	AgentSnapshotFileName = 13
 )
 
 func (self *Controller) RegisterAgentBindHandler(bindHandler channel.BindHandler) {
@@ -67,12 +68,14 @@ func (self *Controller) HandleCustomAgentAsyncOp(conn net.Conn) error {
 }
 
 func (self *Controller) agentOpSnapshotDb(m *channel.Message, ch channel.Channel) {
+	fileName, _ := m.GetStringHeader(AgentSnapshotFileName)
+
 	log := pfxlog.Logger()
-	if err := self.network.SnapshotDatabase(); err != nil {
+	if path, err := self.network.SnapshotDatabaseToFile(fileName); err != nil {
 		log.WithError(err).Error("failed to snapshot db")
 		handler_common.SendOpResult(m, ch, "db.snapshot", err.Error(), false)
 	} else {
-		handler_common.SendOpResult(m, ch, "db.snapshot", "", true)
+		handler_common.SendOpResult(m, ch, "db.snapshot", path, true)
 	}
 }
 
diff --git a/ziti/cmd/agentcli/agent.go b/ziti/cmd/agentcli/agent.go
index 2a99f3f04..ca4be3fc4 100644
--- a/ziti/cmd/agentcli/agent.go
+++ b/ziti/cmd/agentcli/agent.go
@@ -62,7 +62,7 @@ func NewAgentCmd(p common.OptionsProvider) *cobra.Command {
 	}
 
 	agentCmd.AddCommand(ctrlCmd)
-	ctrlCmd.AddCommand(NewSimpleChAgentCustomCmd("snapshot-db", AgentAppController, int32(mgmt_pb.ContentType_SnapshotDbRequestType), p))
+	ctrlCmd.AddCommand(NewAgentSnapshotDb(p))
 	ctrlCmd.AddCommand(NewAgentCtrlInit(p))
 	ctrlCmd.AddCommand(NewAgentCtrlInitFromDb(p))
 
diff --git a/ziti/cmd/agentcli/agent_snapshot_db.go b/ziti/cmd/agentcli/agent_snapshot_db.go
new file mode 100644
index 000000000..0341a1de2
--- /dev/null
+++ b/ziti/cmd/agentcli/agent_snapshot_db.go
@@ -0,0 +1,71 @@
+/*
+	Copyright NetFoundry Inc.
+
+	Licensed under the Apache License, Version 2.0 (the "License");
+	you may not use this file except in compliance with the License.
+	You may obtain a copy of the License at
+
+	https://www.apache.org/licenses/LICENSE-2.0
+
+	Unless required by applicable law or agreed to in writing, software
+	distributed under the License is distributed on an "AS IS" BASIS,
+	WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+	See the License for the specific language governing permissions and
+	limitations under the License.
+*/
+
+package agentcli
+
+import (
+	"fmt"
+	"github.com/openziti/channel/v2"
+	"github.com/openziti/ziti/common/pb/mgmt_pb"
+	"github.com/openziti/ziti/controller"
+	"github.com/openziti/ziti/ziti/cmd/common"
+	"github.com/spf13/cobra"
+)
+
+type AgentSnapshoptDbAction struct {
+	AgentOptions
+}
+
+func NewAgentSnapshotDb(p common.OptionsProvider) *cobra.Command {
+	action := &AgentSnapshoptDbAction{
+		AgentOptions: AgentOptions{
+			CommonOptions: p(),
+		},
+	}
+
+	cmd := &cobra.Command{
+		Args: cobra.ExactArgs(1),
+		Use:  "snapshot-db <snapshot path>",
+		RunE: func(cmd *cobra.Command, args []string) error {
+			action.Cmd = cmd
+			action.Args = args
+			return action.MakeChannelRequest(byte(AgentAppController), action.makeRequest)
+		},
+	}
+	cmd.Args = cobra.MaximumNArgs(1)
+	action.AddAgentOptions(cmd)
+
+	return cmd
+}
+
+func (self *AgentSnapshoptDbAction) makeRequest(ch channel.Channel) error {
+	msg := channel.NewMessage(int32(mgmt_pb.ContentType_SnapshotDbRequestType), nil)
+	if len(self.Args) > 0 {
+		msg.PutStringHeader(controller.AgentSnapshotFileName, self.Args[0])
+	}
+
+	reply, err := msg.WithTimeout(self.timeout).SendForReply(ch)
+	if err != nil {
+		return err
+	}
+	result := channel.UnmarshalResult(reply)
+	if result.Success {
+		fmt.Println(result.Message)
+	} else {
+		fmt.Printf("error: %v\n", result.Message)
+	}
+	return nil
+}