From 36fc3ea253c56d4b557ea355b21e7a95b5bf6be7 Mon Sep 17 00:00:00 2001
From: Owen <owen@txv.io>
Date: Mon, 3 Nov 2025 14:15:16 -0800
Subject: [PATCH] Add exit call

Former-commit-id: 4a89915826b9e0ed36d58562b0277504741ed708
---
 api/api.go | 34 ++++++++++++++++++++++++++++++++++
 olm/olm.go | 12 ++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/api/api.go b/api/api.go
index c7dfcf3..050902c 100644
--- a/api/api.go
+++ b/api/api.go
@@ -44,6 +44,7 @@ type API struct {
 	listener       net.Listener
 	server         *http.Server
 	connectionChan chan ConnectionRequest
+	shutdownChan   chan struct{}
 	statusMu       sync.RWMutex
 	peerStatuses   map[int]*PeerStatus
 	connectedAt    time.Time
@@ -57,6 +58,7 @@ func NewAPI(addr string) *API {
 	s := &API{
 		addr:           addr,
 		connectionChan: make(chan ConnectionRequest, 1),
+		shutdownChan:   make(chan struct{}, 1),
 		peerStatuses:   make(map[int]*PeerStatus),
 	}
 
@@ -68,6 +70,7 @@ func NewAPISocket(socketPath string) *API {
 	s := &API{
 		socketPath:     socketPath,
 		connectionChan: make(chan ConnectionRequest, 1),
+		shutdownChan:   make(chan struct{}, 1),
 		peerStatuses:   make(map[int]*PeerStatus),
 	}
 
@@ -79,6 +82,7 @@ func (s *API) Start() error {
 	mux := http.NewServeMux()
 	mux.HandleFunc("/connect", s.handleConnect)
 	mux.HandleFunc("/status", s.handleStatus)
+	mux.HandleFunc("/exit", s.handleExit)
 
 	s.server = &http.Server{
 		Handler: mux,
@@ -132,6 +136,11 @@ func (s *API) GetConnectionChannel() <-chan ConnectionRequest {
 	return s.connectionChan
 }
 
+// GetShutdownChannel returns the channel for receiving shutdown requests
+func (s *API) GetShutdownChannel() <-chan struct{} {
+	return s.shutdownChan
+}
+
 // UpdatePeerStatus updates the status of a peer including endpoint and relay info
 func (s *API) UpdatePeerStatus(siteID int, connected bool, rtt time.Duration, endpoint string, isRelay bool) {
 	s.statusMu.Lock()
@@ -255,3 +264,28 @@ func (s *API) handleStatus(w http.ResponseWriter, r *http.Request) {
 	w.Header().Set("Content-Type", "application/json")
 	json.NewEncoder(w).Encode(resp)
 }
+
+// handleExit handles the /exit endpoint
+func (s *API) handleExit(w http.ResponseWriter, r *http.Request) {
+	if r.Method != http.MethodPost {
+		http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
+		return
+	}
+
+	logger.Info("Received exit request via API")
+
+	// Send shutdown signal
+	select {
+	case s.shutdownChan <- struct{}{}:
+		// Signal sent successfully
+	default:
+		// Channel already has a signal, don't block
+	}
+
+	// Return a success response
+	w.Header().Set("Content-Type", "application/json")
+	w.WriteHeader(http.StatusOK)
+	json.NewEncoder(w).Encode(map[string]string{
+		"status": "shutdown initiated",
+	})
+}
diff --git a/olm/olm.go b/olm/olm.go
index 7c77f69..3942199 100644
--- a/olm/olm.go
+++ b/olm/olm.go
@@ -58,6 +58,10 @@ type Config struct {
 }
 
 func Run(ctx context.Context, config Config) {
+	// Create a cancellable context for internal shutdown control
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
 	// Extract commonly used values from config for convenience
 	var (
 		endpoint      = config.Endpoint
@@ -108,6 +112,14 @@ func Run(ctx context.Context, config Config) {
 		if err := apiServer.Start(); err != nil {
 			logger.Fatal("Failed to start HTTP server: %v", err)
 		}
+
+		// Listen for shutdown requests from the API
+		go func() {
+			<-apiServer.GetShutdownChannel()
+			logger.Info("Shutdown requested via API")
+			// Cancel the context to trigger graceful shutdown
+			cancel()
+		}()
 	}
 
 	// 	// Use a goroutine to handle connection requests