From b0062ec39aad553b88d02aa6a9d928c1e2cf23d2 Mon Sep 17 00:00:00 2001 From: Vishal Rana Date: Fri, 30 Oct 2026 09:28:21 -0700 Subject: [PATCH] fix(chat): register permission prompts before slow broker work Permission hooks used a 200 ms dial context for an authenticated health round trip before registering their prompt. A slow response silently returned the hook to the terminal without Driga ever seeing the request. Prepare hook clients without that preflight and make registration the first authenticated request. Keep socket validation, token checks, bounded request lifetimes and fail-open behavior. General clients still check broker health before use. Add a deterministic cross-platform socket test that blocks health and delays registration beyond the old preflight budget, then verifies prompt visibility and an allow decision. The test fails with the original dial. Claude-Session: https://claude.ai/code/session_01T2vtcW6ihn1ywy9vxUq9a4 --- internal/chat/transport.go | 28 ++++++++--- internal/chat/transport_test.go | 83 +++++++++++++++++++++++++++++++++ internal/cli/hook_prompt.go | 2 +- 3 files changed, 106 insertions(+), 7 deletions(-) diff --git a/internal/chat/transport.go b/internal/chat/transport.go index 67fc0404..e8111c5d 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 43d6de81..6ca3cba4 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 9b92bebf..97240012 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()