feat(auth): fail-fast for unauthenticated requests before MCP call

- 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
This commit is contained in:
fantiu committed 2026-04-01 15:48:23 +08:00
1 parent 0257d1f084
commit 14f558facf
6 files changed
+424 -13

No files matched your search

+1 -1
View File
@@ -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 {
+11 -2
View File
@@ -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 {
+25 -4
View File
@@ -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)
}
}()
+37 -6
View File
@@ -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 {
+177
View File
@@ -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) != ""
}
+173
View File
@@ -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")
}
}