fixed persistence error and error with clients sometimes not showing up as connected.
This commit is contained in:
@@ -4,6 +4,7 @@ import (
|
|||||||
"context"
|
"context"
|
||||||
"errors"
|
"errors"
|
||||||
"testing"
|
"testing"
|
||||||
|
"time"
|
||||||
|
|
||||||
"overlord-client/cmd/agent/config"
|
"overlord-client/cmd/agent/config"
|
||||||
rt "overlord-client/cmd/agent/runtime"
|
rt "overlord-client/cmd/agent/runtime"
|
||||||
@@ -18,6 +19,7 @@ func TestHandlePing(t *testing.T) {
|
|||||||
Conn: writer,
|
Conn: writer,
|
||||||
Cfg: config.Config{},
|
Cfg: config.Config{},
|
||||||
}
|
}
|
||||||
|
env.SetLastPong(0)
|
||||||
|
|
||||||
ctx := context.Background()
|
ctx := context.Background()
|
||||||
envelope := map[string]interface{}{
|
envelope := map[string]interface{}{
|
||||||
@@ -44,6 +46,13 @@ func TestHandlePing(t *testing.T) {
|
|||||||
if pong.TS != expectedTS {
|
if pong.TS != expectedTS {
|
||||||
t.Errorf("Expected timestamp %d, got %d", expectedTS, pong.TS)
|
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) {
|
func TestHandleCommand_Ping(t *testing.T) {
|
||||||
|
|||||||
@@ -9,6 +9,8 @@ import (
|
|||||||
)
|
)
|
||||||
|
|
||||||
func HandlePing(ctx context.Context, env *runtime.Env, envelope map[string]interface{}) error {
|
func HandlePing(ctx context.Context, env *runtime.Env, envelope map[string]interface{}) error {
|
||||||
|
env.SetLastPong(time.Now().UnixMilli())
|
||||||
|
|
||||||
ts := extractTimestamp(envelope["ts"])
|
ts := extractTimestamp(envelope["ts"])
|
||||||
if ts == 0 {
|
if ts == 0 {
|
||||||
ts = time.Now().UnixMilli()
|
ts = time.Now().UnixMilli()
|
||||||
|
|||||||
@@ -15,6 +15,14 @@ import (
|
|||||||
const registryKey = `SOFTWARE\Microsoft\Windows\CurrentVersion\Run`
|
const registryKey = `SOFTWARE\Microsoft\Windows\CurrentVersion\Run`
|
||||||
const registryValueName = "OverlordAgent"
|
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) {
|
func getTargetPath() (string, error) {
|
||||||
appDataDir := os.Getenv("APPDATA")
|
appDataDir := os.Getenv("APPDATA")
|
||||||
if appDataDir == "" {
|
if appDataDir == "" {
|
||||||
@@ -46,7 +54,7 @@ func install(exePath string) error {
|
|||||||
}
|
}
|
||||||
defer k.Close()
|
defer k.Close()
|
||||||
|
|
||||||
err = k.SetStringValue(registryValueName, targetPath)
|
err = k.SetStringValue(registryValueName, formatRunRegistryCommand(targetPath))
|
||||||
if err != nil {
|
if err != nil {
|
||||||
return fmt.Errorf("failed to set registry value: %w", err)
|
return fmt.Errorf("failed to set registry value: %w", err)
|
||||||
}
|
}
|
||||||
@@ -101,7 +109,7 @@ func configure(exePath string) error {
|
|||||||
}
|
}
|
||||||
defer k.Close()
|
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)
|
return fmt.Errorf("failed to set registry value: %w", err)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -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)
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user