From 3ba179673e6e5310e2e8992efec330ea56f12f62 Mon Sep 17 00:00:00 2001
From: Michael Quigley <michael@quigley.com>
Date: Fri, 23 Aug 2024 12:35:19 -0400
Subject: [PATCH] shutdown hook to clean up agent at exit (#463)

---
 agent/agent.go         | 18 +++++++++++++++---
 cmd/zrok/agentStart.go | 15 +++++++++++++++
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/agent/agent.go b/agent/agent.go
index 246e5a3b..2d01211d 100644
--- a/agent/agent.go
+++ b/agent/agent.go
@@ -7,12 +7,14 @@ import (
 	"github.com/sirupsen/logrus"
 	"google.golang.org/grpc"
 	"net"
+	"os"
 )
 
 type Agent struct {
-	root     env_core.Root
-	shares   map[string]*share
-	accesses map[string]*access
+	root        env_core.Root
+	agentSocket string
+	shares      map[string]*share
+	accesses    map[string]*access
 }
 
 func NewAgent(root env_core.Root) (*Agent, error) {
@@ -28,6 +30,7 @@ func NewAgent(root env_core.Root) (*Agent, error) {
 
 func (a *Agent) Run() error {
 	logrus.Infof("started")
+
 	agentSocket, err := a.root.AgentSocket()
 	if err != nil {
 		return err
@@ -36,10 +39,19 @@ func (a *Agent) Run() error {
 	if err != nil {
 		return err
 	}
+	a.agentSocket = agentSocket
+
 	srv := grpc.NewServer()
 	agentGrpc.RegisterAgentServer(srv, &agentGrpcImpl{})
 	if err := srv.Serve(l); err != nil {
 		return err
 	}
+
 	return nil
 }
+
+func (a *Agent) Shutdown() {
+	if err := os.Remove(a.agentSocket); err != nil {
+		logrus.Warnf("unable to remove agent socket: %v", err)
+	}
+}
diff --git a/cmd/zrok/agentStart.go b/cmd/zrok/agentStart.go
index 698863b2..3c3c651d 100644
--- a/cmd/zrok/agentStart.go
+++ b/cmd/zrok/agentStart.go
@@ -5,6 +5,9 @@ import (
 	"github.com/openziti/zrok/environment"
 	"github.com/openziti/zrok/tui"
 	"github.com/spf13/cobra"
+	"os"
+	"os/signal"
+	"syscall"
 )
 
 func init() {
@@ -41,7 +44,19 @@ func (cmd *agentStartCommand) run(_ *cobra.Command, _ []string) {
 		tui.Error("error creating agent", err)
 	}
 
+	c := make(chan os.Signal)
+	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
+	go func() {
+		<-c
+		cmd.shutdown(a)
+		os.Exit(0)
+	}()
+
 	if err := a.Run(); err != nil {
 		tui.Error("agent aborted", err)
 	}
 }
+
+func (cmd *agentStartCommand) shutdown(a *agent.Agent) {
+	a.Shutdown()
+}