From 14f558facf89e8323d709abe5539a43a2f8f3da9 Mon Sep 17 00:00:00 2001 From: fantiu Date: Wed, 1 Apr 2026 15:48:23 +0800 Subject: [PATCH] feat(auth): fail-fast for unauthenticated requests before MCP call MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add token validation in executeInvocation to reject requests early - Return clear Chinese error message: '未登录，请先执行 dws auth login' - Include actionable hint and suggested command in error response - Avoid wasting network resources on HTTP 400 from MCP gateway Also includes: - Add DWS_PERF_TIMING env for CLI execution timing analysis - Add TimingCollector to track cmd_init, auth_token, mcp_call durations Test changes: - Add --token flag to tests that require authentication - Add TestRuntimeRunnerRejectsUnauthenticatedRequest test case --- internal/app/recovery_command_test.go | 2 +- internal/app/root.go | 13 +- internal/app/runner.go | 29 ++++- internal/app/runner_test.go | 43 ++++++- internal/app/timing.go | 177 ++++++++++++++++++++++++++ internal/app/timing_test.go | 173 +++++++++++++++++++++++++ 6 files changed, 424 insertions(+), 13 deletions(-) create mode 100644 internal/app/timing.go create mode 100644 internal/app/timing_test.go diff --git a/internal/app/recovery_command_test.go b/internal/app/recovery_command_test.go index 32fa6cf7..dd317b78 100644 --- a/internal/app/recovery_command_test.go +++ b/internal/app/recovery_command_test.go @@ -261,7 +261,7 @@ func TestExecuteWritesRecoveryEventIDToStderrOnCapturedFailure(t *testing.T) { oldArgs := os.Args defer func() { os.Args = oldArgs }() - os.Args = []string{"dws", "mcp", "doc", "search_documents", "--json", `{"keyword":"design"}`} + os.Args = []string{"dws", "mcp", "doc", "search_documents", "--json", `{"keyword":"design"}`, "--token", "test-token"} stdoutR, stdoutW, err := os.Pipe() if err != nil { diff --git a/internal/app/root.go b/internal/app/root.go index 46488b64..331a0c0b 100644 --- a/internal/app/root.go +++ b/internal/app/root.go @@ -52,7 +52,9 @@ const recoveryEventStderrPrefix = "RECOVERY_EVENT_ID=" // Execute runs the root command and returns the process exit code. func Execute() int { totalStart := time.Now() + timing := NewTimingCollector() defer func() { + timing.PrintIfEnabled() if os.Getenv("DWS_PERF_DEBUG") != "" { _, _ = fmt.Fprintf(os.Stderr, "[PERF] Execute total: %v\n", time.Since(totalStart)) } @@ -61,12 +63,17 @@ func Execute() int { ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) defer cancel() + // Attach timing collector to context for use by child components + ctx = WithTimingCollector(ctx, timing) + initStart := time.Now() recovery.ResetRuntimeState() engine := newPipelineEngine() root := NewRootCommandWithEngine(ctx, engine) + initDuration := time.Since(initStart) + timing.Record("cmd_init", initDuration) if os.Getenv("DWS_PERF_DEBUG") != "" { - _, _ = fmt.Fprintf(os.Stderr, "[PERF] command init: %v\n", time.Since(initStart)) + _, _ = fmt.Fprintf(os.Stderr, "[PERF] command init: %v\n", initDuration) } // Run PreParse handlers on raw argv before Cobra parses flags. @@ -76,8 +83,10 @@ func Execute() int { execStart := time.Now() executed, err := root.ExecuteC() + execDuration := time.Since(execStart) + timing.Record("cobra_exec", execDuration) if os.Getenv("DWS_PERF_DEBUG") != "" { - _, _ = fmt.Fprintf(os.Stderr, "[PERF] cobra ExecuteC: %v\n", time.Since(execStart)) + _, _ = fmt.Fprintf(os.Stderr, "[PERF] cobra ExecuteC: %v\n", execDuration) } if err != nil { if executed == nil { diff --git a/internal/app/runner.go b/internal/app/runner.go index 12544542..f72cd6cb 100644 --- a/internal/app/runner.go +++ b/internal/app/runner.go @@ -101,7 +101,9 @@ func (r *runtimeRunner) Run(ctx context.Context, invocation executor.Invocation) } } + catalogStart := time.Now() catalog, err := r.loader.Load(ctx) + RecordTiming(ctx, "catalog_load", time.Since(catalogStart)) if err != nil { return executor.Result{}, err } @@ -132,9 +134,11 @@ func (r *runtimeRunner) executeInvocation(ctx context.Context, endpoint string, } authStart := time.Now() - tc := r.transport.WithAuth(r.resolveAuthToken(ctx), resolveIdentityHeaders()) + authToken := r.resolveAuthToken(ctx) + authDuration := time.Since(authStart) + RecordTiming(ctx, "auth_token", authDuration) if os.Getenv("DWS_PERF_DEBUG") != "" { - _, _ = fmt.Fprintf(os.Stderr, "[PERF] resolveAuthToken: %v\n", time.Since(authStart)) + _, _ = fmt.Fprintf(os.Stderr, "[PERF] resolveAuthToken: %v\n", authDuration) } if invocation.DryRun { @@ -166,10 +170,25 @@ func (r *runtimeRunner) executeInvocation(ctx context.Context, endpoint string, }, nil } + // Fail-fast: reject unauthenticated requests before making network calls. + // This provides a clear error message instead of cryptic HTTP 400 from MCP. + if strings.TrimSpace(authToken) == "" { + return executor.Result{}, apperrors.NewAuth( + "未登录，请先执行 dws auth login", + apperrors.WithReason("not_authenticated"), + apperrors.WithHint("运行 'dws auth login' 完成登录后重试"), + apperrors.WithActions("dws auth login"), + ) + } + + tc := r.transport.WithAuth(authToken, resolveIdentityHeaders()) + callStart := time.Now() callResult, err := tc.CallTool(ctx, endpoint, invocation.Tool, invocation.Params) + callDuration := time.Since(callStart) + RecordTiming(ctx, "mcp_call", callDuration) if os.Getenv("DWS_PERF_DEBUG") != "" { - _, _ = fmt.Fprintf(os.Stderr, "[PERF] MCP CallTool: %v\n", time.Since(callStart)) + _, _ = fmt.Fprintf(os.Stderr, "[PERF] MCP CallTool: %v\n", callDuration) } if err != nil { captureRuntimeFailure(invocation, err, err) @@ -247,8 +266,10 @@ func getCachedRuntimeToken(ctx context.Context) string { cachedRuntimeTokenOnce.Do(func() { loadStart := time.Now() defer func() { + loadDuration := time.Since(loadStart) + RecordTiming(ctx, "keychain_load", loadDuration) if os.Getenv("DWS_PERF_DEBUG") != "" { - _, _ = fmt.Fprintf(os.Stderr, "[PERF] getCachedRuntimeToken (first load): %v\n", time.Since(loadStart)) + _, _ = fmt.Fprintf(os.Stderr, "[PERF] getCachedRuntimeToken (first load): %v\n", loadDuration) } }() diff --git a/internal/app/runner_test.go b/internal/app/runner_test.go index 8434084b..d197a102 100644 --- a/internal/app/runner_test.go +++ b/internal/app/runner_test.go @@ -45,7 +45,7 @@ func TestRuntimeRunnerIncludesContentScanReportWhenEnabled(t *testing.T) { var out bytes.Buffer cmd.SetOut(&out) cmd.SetErr(&out) - cmd.SetArgs([]string{"mcp", "doc", "search_documents", "--json", `{"keyword":"design"}`}) + cmd.SetArgs([]string{"mcp", "doc", "search_documents", "--json", `{"keyword":"design"}`, "--token", "test-token"}) if err := cmd.Execute(); err != nil { t.Fatalf("Execute() error = %v", err) @@ -90,7 +90,7 @@ func TestRuntimeRunnerBlocksUnsafeContentWhenEnforced(t *testing.T) { cmd := NewRootCommand() cmd.SetOut(&bytes.Buffer{}) cmd.SetErr(&bytes.Buffer{}) - cmd.SetArgs([]string{"mcp", "doc", "search_documents", "--json", `{"keyword":"design"}`}) + cmd.SetArgs([]string{"mcp", "doc", "search_documents", "--json", `{"keyword":"design"}`, "--token", "test-token"}) err := cmd.Execute() if err == nil { @@ -121,7 +121,7 @@ func TestCanonicalCommandUsesRuntimeRunnerWhenEnabled(t *testing.T) { var out bytes.Buffer cmd.SetOut(&out) cmd.SetErr(&out) - cmd.SetArgs([]string{"mcp", "doc", "create_document", "--json", `{"title":"Quarterly"}`, "--yes"}) + cmd.SetArgs([]string{"mcp", "doc", "create_document", "--json", `{"title":"Quarterly"}`, "--yes", "--token", "test-token"}) if err := cmd.Execute(); err != nil { t.Fatalf("Execute() error = %v", err) @@ -255,6 +255,37 @@ func TestRuntimeRunnerInjectsAuthTokenFromFlag(t *testing.T) { } } +// TestRuntimeRunnerRejectsUnauthenticatedRequest verifies that requests without +// a valid token are rejected with a clear error before making any network call. +func TestRuntimeRunnerRejectsUnauthenticatedRequest(t *testing.T) { + setupRuntimeCommandTest(t) + server := mockmcp.DefaultServer() + defer server.Close() + + t.Setenv(cli.CatalogFixtureEnv, writeDocCatalogFixture(t, server.RemoteURL("/server/doc"), false)) + + cmd := NewRootCommand() + var stdout, stderr bytes.Buffer + cmd.SetOut(&stdout) + cmd.SetErr(&stderr) + // No --token flag, should be rejected + cmd.SetArgs([]string{"mcp", "doc", "search_documents", "--json", `{"keyword":"design"}`}) + + err := cmd.Execute() + if err == nil { + t.Fatal("Execute() error = nil, want authentication error") + } + + // Verify we get a clear auth error, not a cryptic HTTP 400 + errMsg := err.Error() + if !strings.Contains(errMsg, "未登录") { + t.Fatalf("Execute() error = %v, want error containing '未登录'", err) + } + if !strings.Contains(errMsg, "auth login") { + t.Fatalf("Execute() error = %v, want error containing 'auth login'", err) + } +} + func TestRuntimeRunnerFallsBackForUnavailableProduct(t *testing.T) { setupRuntimeCommandTest(t) server := mockmcp.DefaultServer() @@ -415,7 +446,7 @@ func TestCanonicalSensitiveToolAcceptsInteractiveConfirmation(t *testing.T) { cmd.SetOut(&out) cmd.SetErr(&errOut) cmd.SetIn(strings.NewReader("yes\n")) - cmd.SetArgs([]string{"mcp", "doc", "create_document", "--json", `{"title":"Quarterly"}`}) + cmd.SetArgs([]string{"mcp", "doc", "create_document", "--json", `{"title":"Quarterly"}`, "--token", "test-token"}) if err := cmd.Execute(); err != nil { t.Fatalf("Execute() error = %v", err) @@ -465,7 +496,7 @@ func TestRuntimeRunnerUsesProductEndpointOverride(t *testing.T) { var out bytes.Buffer cmd.SetOut(&out) cmd.SetErr(&out) - cmd.SetArgs([]string{"mcp", "doc", "search_documents", "--json", `{"keyword":"design"}`}) + cmd.SetArgs([]string{"mcp", "doc", "search_documents", "--json", `{"keyword":"design"}`, "--token", "test-token"}) if err := cmd.Execute(); err != nil { t.Fatalf("Execute() error = %v", err) @@ -628,7 +659,7 @@ func TestRuntimeRunnerReturnsErrorWhenMCPIsErrorTrue(t *testing.T) { cmd := NewRootCommand() cmd.SetOut(&bytes.Buffer{}) cmd.SetErr(&bytes.Buffer{}) - cmd.SetArgs([]string{"mcp", "doc", "search_documents", "--json", `{"keyword":"design"}`}) + cmd.SetArgs([]string{"mcp", "doc", "search_documents", "--json", `{"keyword":"design"}`, "--token", "test-token"}) err := cmd.Execute() if err == nil { diff --git a/internal/app/timing.go b/internal/app/timing.go new file mode 100644 index 00000000..9f3ff5e3 --- /dev/null +++ b/internal/app/timing.go @@ -0,0 +1,177 @@ +// Copyright 2026 Alibaba Group +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package app + +import ( + "context" + "fmt" + "io" + "os" + "sort" + "sync" + "time" +) + +// Environment variable to enable performance timing output. +const PerfTimingEnv = "DWS_PERF_TIMING" + +// timingContextKey is the context key for TimingCollector. +type timingContextKey struct{} + +// TimingEntry represents a single timing measurement. +type TimingEntry struct { + Name string + Duration time.Duration + Timestamp time.Time + Seq int // insertion order +} + +// TimingCollector collects timing measurements for a single command execution. +// It is safe for concurrent use. +type TimingCollector struct { + mu sync.Mutex + start time.Time + entries []TimingEntry + seq int +} + +// NewTimingCollector creates a new collector with the start time set to now. +func NewTimingCollector() *TimingCollector { + return &TimingCollector{ + start: time.Now(), + entries: make([]TimingEntry, 0, 16), + } +} + +// Record adds a timing entry with the given name and duration. +func (tc *TimingCollector) Record(name string, d time.Duration) { + if tc == nil { + return + } + tc.mu.Lock() + defer tc.mu.Unlock() + tc.entries = append(tc.entries, TimingEntry{ + Name: name, + Duration: d, + Timestamp: time.Now(), + Seq: tc.seq, + }) + tc.seq++ +} + +// StartTimer returns a function that, when called, records the elapsed time +// since StartTimer was called. This is convenient for defer usage: +// +// defer tc.StartTimer("operation")() +func (tc *TimingCollector) StartTimer(name string) func() { + if tc == nil { + return func() {} + } + start := time.Now() + return func() { + tc.Record(name, time.Since(start)) + } +} + +// Total returns the total elapsed time since the collector was created. +func (tc *TimingCollector) Total() time.Duration { + if tc == nil { + return 0 + } + return time.Since(tc.start) +} + +// Entries returns a copy of all recorded entries in insertion order. +func (tc *TimingCollector) Entries() []TimingEntry { + if tc == nil { + return nil + } + tc.mu.Lock() + defer tc.mu.Unlock() + result := make([]TimingEntry, len(tc.entries)) + copy(result, tc.entries) + sort.Slice(result, func(i, j int) bool { + return result[i].Seq < result[j].Seq + }) + return result +} + +// Print writes a summary of all timing entries to the given writer. +func (tc *TimingCollector) Print(w io.Writer) { + if tc == nil || w == nil { + return + } + entries := tc.Entries() + if len(entries) == 0 { + fmt.Fprintf(w, "\n[Timing] Total: %v (no detailed entries)\n", tc.Total().Truncate(time.Millisecond)) + return + } + + fmt.Fprintln(w) + fmt.Fprintln(w, "[Timing] Execution breakdown:") + for _, e := range entries { + fmt.Fprintf(w, " %-30s %v\n", e.Name, e.Duration.Truncate(time.Millisecond)) + } + fmt.Fprintf(w, " %-30s %v\n", "──────────────────────────────", "──────────") + fmt.Fprintf(w, " %-30s %v\n", "Total", tc.Total().Truncate(time.Millisecond)) +} + +// PrintIfEnabled prints timing info to stderr if DWS_PERF_TIMING is set. +func (tc *TimingCollector) PrintIfEnabled() { + if tc == nil { + return + } + if os.Getenv(PerfTimingEnv) == "" { + return + } + tc.Print(os.Stderr) +} + +// WithTimingCollector returns a new context with the TimingCollector attached. +func WithTimingCollector(ctx context.Context, tc *TimingCollector) context.Context { + return context.WithValue(ctx, timingContextKey{}, tc) +} + +// TimingCollectorFromContext extracts the TimingCollector from context, or nil. +func TimingCollectorFromContext(ctx context.Context) *TimingCollector { + if ctx == nil { + return nil + } + tc, _ := ctx.Value(timingContextKey{}).(*TimingCollector) + return tc +} + +// RecordTiming is a convenience function to record timing to the collector in context. +func RecordTiming(ctx context.Context, name string, d time.Duration) { + if tc := TimingCollectorFromContext(ctx); tc != nil { + tc.Record(name, d) + } +} + +// StartTiming is a convenience function that returns a stop function for defer usage. +// Example: +// +// defer StartTiming(ctx, "operation")() +func StartTiming(ctx context.Context, name string) func() { + tc := TimingCollectorFromContext(ctx) + if tc == nil { + return func() {} + } + return tc.StartTimer(name) +} + +// IsPerfTimingEnabled returns true if performance timing output is enabled. +func IsPerfTimingEnabled() bool { + return os.Getenv(PerfTimingEnv) != "" +} diff --git a/internal/app/timing_test.go b/internal/app/timing_test.go new file mode 100644 index 00000000..f57fb9e3 --- /dev/null +++ b/internal/app/timing_test.go @@ -0,0 +1,173 @@ +// Copyright 2026 Alibaba Group +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package app + +import ( + "bytes" + "context" + "os" + "strings" + "testing" + "time" +) + +func TestTimingCollector_Basic(t *testing.T) { + tc := NewTimingCollector() + if tc == nil { + t.Fatal("NewTimingCollector returned nil") + } + + // Record some timings + tc.Record("op1", 10*time.Millisecond) + tc.Record("op2", 20*time.Millisecond) + + entries := tc.Entries() + if len(entries) != 2 { + t.Errorf("expected 2 entries, got %d", len(entries)) + } + + // Check ordering + if entries[0].Name != "op1" { + t.Errorf("expected first entry to be 'op1', got %q", entries[0].Name) + } + if entries[1].Name != "op2" { + t.Errorf("expected second entry to be 'op2', got %q", entries[1].Name) + } +} + +func TestTimingCollector_StartTimer(t *testing.T) { + tc := NewTimingCollector() + + stop := tc.StartTimer("timed_op") + time.Sleep(5 * time.Millisecond) + stop() + + entries := tc.Entries() + if len(entries) != 1 { + t.Fatalf("expected 1 entry, got %d", len(entries)) + } + if entries[0].Name != "timed_op" { + t.Errorf("expected entry name 'timed_op', got %q", entries[0].Name) + } + if entries[0].Duration < 5*time.Millisecond { + t.Errorf("expected duration >= 5ms, got %v", entries[0].Duration) + } +} + +func TestTimingCollector_NilSafe(t *testing.T) { + var tc *TimingCollector + + // Should not panic on nil collector + tc.Record("op", 10*time.Millisecond) + stop := tc.StartTimer("op") + stop() + _ = tc.Total() + _ = tc.Entries() + tc.Print(nil) + tc.PrintIfEnabled() +} + +func TestTimingCollector_Print(t *testing.T) { + tc := NewTimingCollector() + tc.Record("auth_token", 44*time.Millisecond) + tc.Record("mcp_call", 150*time.Millisecond) + + var buf bytes.Buffer + tc.Print(&buf) + + output := buf.String() + if !strings.Contains(output, "[Timing]") { + t.Error("output should contain [Timing] header") + } + if !strings.Contains(output, "auth_token") { + t.Error("output should contain 'auth_token'") + } + if !strings.Contains(output, "mcp_call") { + t.Error("output should contain 'mcp_call'") + } + if !strings.Contains(output, "Total") { + t.Error("output should contain 'Total'") + } +} + +func TestTimingCollector_PrintIfEnabled(t *testing.T) { + // Set environment variable + os.Setenv(PerfTimingEnv, "1") + defer os.Unsetenv(PerfTimingEnv) + + tc := NewTimingCollector() + tc.Record("test_op", 10*time.Millisecond) + + // This should not panic and should print to stderr + tc.PrintIfEnabled() +} + +func TestTimingCollector_ContextIntegration(t *testing.T) { + tc := NewTimingCollector() + ctx := WithTimingCollector(context.Background(), tc) + + // Retrieve from context + retrieved := TimingCollectorFromContext(ctx) + if retrieved != tc { + t.Error("TimingCollectorFromContext should return the same collector") + } + + // Use convenience functions + RecordTiming(ctx, "ctx_op", 30*time.Millisecond) + stop := StartTiming(ctx, "ctx_timed") + time.Sleep(2 * time.Millisecond) + stop() + + entries := tc.Entries() + if len(entries) != 2 { + t.Errorf("expected 2 entries, got %d", len(entries)) + } +} + +func TestTimingCollectorFromContext_NilContext(t *testing.T) { + tc := TimingCollectorFromContext(nil) + if tc != nil { + t.Error("TimingCollectorFromContext(nil) should return nil") + } +} + +func TestTimingCollectorFromContext_NoCollector(t *testing.T) { + tc := TimingCollectorFromContext(context.Background()) + if tc != nil { + t.Error("TimingCollectorFromContext with no collector should return nil") + } +} + +func TestStartTiming_NoCollector(t *testing.T) { + ctx := context.Background() + stop := StartTiming(ctx, "no_collector") + // Should not panic + stop() +} + +func TestIsPerfTimingEnabled(t *testing.T) { + // Clear the env var first + os.Unsetenv(PerfTimingEnv) + + if IsPerfTimingEnabled() { + t.Error("IsPerfTimingEnabled should return false when env var is not set") + } + + os.Setenv(PerfTimingEnv, "1") + defer os.Unsetenv(PerfTimingEnv) + + if !IsPerfTimingEnabled() { + t.Error("IsPerfTimingEnabled should return true when env var is set") + } +}