diff --git a/internal/chat/transport.go b/internal/chat/transport.go index 67fc040..e8111c5 100644 --- a/internal/chat/transport.go +++ b/internal/chat/transport.go @@ -273,23 +273,39 @@ type Remote struct { } func Dial(ctx context.Context, dataDir string) (*Remote, error) { - path := SocketPath(dataDir) - c, err := dialPath(ctx, path) + c, err := DialHook(ctx, dataDir) + if err != nil { + return nil, err + } + if err := c.call(ctx, "GET", "/health", nil, nil); err != nil { + c.Close() + return nil, err + } + return c, nil +} + +// DialHook prepares an authenticated hook client without a health preflight. +// Registration must be the first request: a slow health response must not +// consume the hook's short connection budget and silently lose the prompt. +// The request itself verifies the token and bounds the actual socket dial. +func DialHook(ctx context.Context, dataDir string) (*Remote, error) { + if err := ctx.Err(); err != nil { + return nil, err + } + c, err := dialPath(ctx, SocketPath(dataDir)) if err != nil { return nil, err } c.keyPath = filepath.Join(dataDir, "chat.key") key, err := os.ReadFile(c.keyPath) - if err == nil { - c.token = string(key) - err = c.call(ctx, "GET", "/health", nil, nil) - } if err != nil { c.Close() return nil, err } + c.token = string(key) return c, nil } + func dialPath(ctx context.Context, path string) (*Remote, error) { if err := fleet.CheckSocketDir(path); err != nil { return nil, err diff --git a/internal/chat/transport_test.go b/internal/chat/transport_test.go index 43d6de8..6ca3cba 100644 --- a/internal/chat/transport_test.go +++ b/internal/chat/transport_test.go @@ -2,13 +2,17 @@ package chat import ( "context" + "fmt" "github.com/labstack/driga/internal/fleet" "github.com/labstack/driga/internal/secfile" + "net" + "net/http" "os" "path/filepath" "runtime" "strings" "testing" + "time" ) func TestSocketAndFallback(t *testing.T) { @@ -156,3 +160,82 @@ func TestChatRefusesUnsafeFallbackDirectory(t *testing.T) { }) } } + +// A loaded collector can take longer than the hook's old 200 ms health +// budget even though it is listening and able to register the permission. +func TestHookRegistersWithSlowBroker(t *testing.T) { + b, _, _ := fixture(t, true) + dir := t.TempDir() + if err := secfile.MkdirPrivate(dir); err != nil { + t.Fatal(err) + } + if err := secfile.Protect(dir); err != nil { + t.Fatal(err) + } + path := SocketPath(dir) + if err := fleet.EnsureSocketDir(path); err != nil { + t.Fatal(err) + } + l, err := net.Listen("unix", path) + if err != nil { + t.Fatal(err) + } + defer os.Remove(path) + if runtime.GOOS != "windows" { + if err := os.Chmod(path, 0600); err != nil { + t.Fatal(err) + } + } + b.clientToken = "test-hook-token" + if err := os.WriteFile(filepath.Join(dir, "chat.key"), []byte(b.clientToken), 0600); err != nil { + t.Fatal(err) + } + release := make(chan struct{}) + defer close(release) + h := handler(b) + server := &http.Server{Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path == "/health" { + // Never finish a preflight before its deadline. No timing race or + // platform-specific scheduling is needed to reproduce the lost prompt. + select { + case <-release: + case <-r.Context().Done(): + } + return + } + if r.URL.Path == "/hook-prompt" { + select { + case <-time.After(300 * time.Millisecond): + case <-r.Context().Done(): + return + } + } + h.ServeHTTP(w, r) + })} + defer server.Close() + go server.Serve(l) + dialCtx, cancelDial := context.WithTimeout(context.Background(), 200*time.Millisecond) + c, err := DialHook(dialCtx, dir) + cancelDial() + if err != nil { + t.Fatalf("hook did not register prompt: %v", err) + } + defer c.Close() + ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second) + defer cancel() + done := make(chan error, 1) + go func() { + out, err := c.WaitPrompt(ctx, testPrompt("slow")) + if err == nil && out.Behavior != "allow" { + err = fmt.Errorf("decision: %+v", out) + } + done <- err + }() + waitPrompt(t, b, "slow") + if err := b.AnswerPrompt(ctx, "slow", true); err != nil { + t.Fatal(err) + } + if err := <-done; err != nil { + t.Fatal(err) + } +} diff --git a/internal/cli/hook_prompt.go b/internal/cli/hook_prompt.go index 9b92beb..9724001 100644 --- a/internal/cli/hook_prompt.go +++ b/internal/cli/hook_prompt.go @@ -20,7 +20,7 @@ type permissionPayload struct { func promptBroker() (*chat.Remote, error) { ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) defer cancel() - return chat.Dial(ctx, config.DataDir()) + return chat.DialHook(ctx, config.DataDir()) } func endPromptSession(session, cwd string) { c, err := promptBroker()