Merge pull request #4214 from bestony/agent/devbox-coding/25c66071-1783957460

feat: add opt-in Server-Timing for Admin UI APIs
This commit is contained in:
Wesley Liddick
2026-07-14 10:14:17 +08:00
committed by GitHub
43 changed files with 2023 additions and 26 deletions
+5
View File
@@ -601,6 +601,7 @@ type ServerConfig struct {
Host string `mapstructure:"host"`
Port int `mapstructure:"port"`
Mode string `mapstructure:"mode"` // debug/release
EnableServerTiming bool `mapstructure:"enable_server_timing"` // Admin UI Server-Timing response header
FrontendURL string `mapstructure:"frontend_url"` // 前端基础 URL,用于生成邮件中的外部链接
ReadHeaderTimeout int `mapstructure:"read_header_timeout"` // 读取请求头超时(秒)
IdleTimeout int `mapstructure:"idle_timeout"` // 空闲连接超时(秒)
@@ -1459,6 +1460,9 @@ func load(allowMissingJWTSecret bool) (*Config, error) {
// 环境变量支持
viper.AutomaticEnv()
viper.SetEnvKeyReplacer(strings.NewReplacer(".", "_"))
if err := viper.BindEnv("server.enable_server_timing", "ENABLE_SERVER_TIMING"); err != nil {
return nil, fmt.Errorf("bind ENABLE_SERVER_TIMING: %w", err)
}
// 默认值
setDefaults()
@@ -1614,6 +1618,7 @@ func setDefaults() {
viper.SetDefault("server.host", "0.0.0.0")
viper.SetDefault("server.port", 8080)
viper.SetDefault("server.mode", "release")
viper.SetDefault("server.enable_server_timing", false)
viper.SetDefault("server.frontend_url", "")
viper.SetDefault("server.read_header_timeout", 30) // 30秒读取请求头
viper.SetDefault("server.idle_timeout", 120) // 120秒空闲超时
+17
View File
@@ -17,6 +17,23 @@ func resetViperWithJWTSecret(t *testing.T) {
t.Setenv("JWT_SECRET", strings.Repeat("x", 32))
}
func TestLoadServerTimingConfig(t *testing.T) {
t.Run("disabled by default", func(t *testing.T) {
resetViperWithJWTSecret(t)
cfg, err := Load()
require.NoError(t, err)
require.False(t, cfg.Server.EnableServerTiming)
})
t.Run("enabled by exact environment variable", func(t *testing.T) {
resetViperWithJWTSecret(t)
t.Setenv("ENABLE_SERVER_TIMING", "true")
cfg, err := Load()
require.NoError(t, err)
require.True(t, cfg.Server.EnableServerTiming)
})
}
func TestLoadForBootstrapAllowsMissingJWTSecret(t *testing.T) {
viper.Reset()
t.Setenv("JWT_SECRET", "")
@@ -16,6 +16,7 @@ import (
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/gorilla/websocket"
@@ -323,7 +324,7 @@ func (h *OpsHandler) QPSWSHandler(c *gin.Context) {
// If realtime monitoring is disabled, prefer a successful WS upgrade followed by a clean close
// with a deterministic close code. This prevents clients from spinning on 404/1006 reconnect loops.
if !h.opsService.IsRealtimeMonitoringEnabled(c.Request.Context()) {
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
conn, err := upgrader.Upgrade(c.Writer, c.Request, servermiddleware.ServerTimingResponseHeader(c))
if err != nil {
c.JSON(http.StatusNotFound, gin.H{"error": "ops realtime monitoring is disabled"})
return
@@ -358,7 +359,7 @@ func (h *OpsHandler) QPSWSHandler(c *gin.Context) {
defer releaseOpsWSIPSlot(clientIP)
}
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
conn, err := upgrader.Upgrade(c.Writer, c.Request, servermiddleware.ServerTimingResponseHeader(c))
if err != nil {
logger.LegacyPrintf("handler.admin.ops_ws", "[OpsWS] upgrade failed: %v", err)
return
+9 -9
View File
@@ -17,6 +17,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
)
// ForbiddenError 表示上游返回 403 Forbidden
@@ -279,7 +280,6 @@ func NewClient(proxyURL string) (*Client, error) {
}
client.Transport = transport
}
return &Client{
httpClient: client,
}, nil
@@ -341,7 +341,7 @@ func (c *Client) ExchangeCode(ctx context.Context, code, codeVerifier string) (*
}
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
resp, err := c.httpClient.Do(req)
resp, err := servertiming.Do(c.httpClient, req)
if err != nil {
return nil, fmt.Errorf("token 交换请求失败: %w", err)
}
@@ -383,7 +383,7 @@ func (c *Client) RefreshToken(ctx context.Context, refreshToken string) (*TokenR
}
req.Header.Set("Content-Type", "application/x-www-form-urlencoded")
resp, err := c.httpClient.Do(req)
resp, err := servertiming.Do(c.httpClient, req)
if err != nil {
return nil, fmt.Errorf("token 刷新请求失败: %w", err)
}
@@ -414,7 +414,7 @@ func (c *Client) GetUserInfo(ctx context.Context, accessToken string) (*UserInfo
}
req.Header.Set("Authorization", "Bearer "+accessToken)
resp, err := c.httpClient.Do(req)
resp, err := servertiming.Do(c.httpClient, req)
if err != nil {
return nil, fmt.Errorf("用户信息请求失败: %w", err)
}
@@ -465,7 +465,7 @@ func (c *Client) LoadCodeAssist(ctx context.Context, accessToken string) (*LoadC
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
resp, err := c.httpClient.Do(req)
resp, err := servertiming.Do(c.httpClient, req)
if err != nil {
lastErr = fmt.Errorf("loadCodeAssist 请求失败: %w", err)
if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
@@ -544,7 +544,7 @@ func (c *Client) OnboardUser(ctx context.Context, accessToken, tierID string) (s
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
resp, err := c.httpClient.Do(req)
resp, err := servertiming.Do(c.httpClient, req)
if err != nil {
lastErr = fmt.Errorf("onboardUser 请求失败: %w", err)
if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
@@ -683,7 +683,7 @@ func (c *Client) FetchAvailableModels(ctx context.Context, accessToken, projectI
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", GetUserAgentForContext(ctx))
resp, err := fetchClient.Do(req)
resp, err := servertiming.Do(fetchClient, req)
if err != nil {
lastErr = fmt.Errorf("fetchAvailableModels 请求失败: %w", err)
if shouldFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
@@ -842,7 +842,7 @@ func (c *Client) SetUserSettings(ctx context.Context, accessToken string) (*SetU
req.Header.Set("X-Goog-Api-Client", "gl-node/22.21.1")
req.Host = "daily-cloudcode-pa.googleapis.com"
resp, err := c.httpClient.Do(req)
resp, err := servertiming.Do(c.httpClient, req)
if err != nil {
return nil, fmt.Errorf("setUserSettings 请求失败: %w", err)
}
@@ -885,7 +885,7 @@ func (c *Client) FetchUserInfo(ctx context.Context, accessToken, projectID strin
req.Header.Set("X-Goog-Api-Client", "gl-node/22.21.1")
req.Host = "daily-cloudcode-pa.googleapis.com"
resp, err := c.httpClient.Do(req)
resp, err := servertiming.Do(c.httpClient, req)
if err != nil {
return nil, fmt.Errorf("fetchUserInfo 请求失败: %w", err)
}
+2
View File
@@ -25,6 +25,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
"github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
)
@@ -92,6 +93,7 @@ func buildClient(opts Options) (*http.Client, error) {
if opts.ValidateResolvedIP && !opts.AllowPrivateHosts {
rt = newValidatedTransport(transport)
}
rt = servertiming.WrapRoundTripper(rt)
return &http.Client{
Transport: rt,
Timeout: opts.Timeout,
@@ -0,0 +1,348 @@
package servertiming
import (
"context"
"fmt"
"sort"
"strconv"
"strings"
"sync"
"time"
)
const (
HeaderName = "Server-Timing"
AdminUIHeader = "X-Admin-UI-Request"
MetricDatabase = "db"
MetricRedis = "redis"
dependencyPrefix = "dep_"
maxMetricNameLength = 48
maxIntervals = 2048
maxHeaderLength = 4096
)
type contextKey struct{}
type interval struct {
start time.Time
end time.Time
}
type metric struct {
count int64
intervals []interval
}
// Collector stores request-scoped timing samples. It is safe for concurrent use.
type Collector struct {
startedAt time.Time
mu sync.Mutex
metrics map[string]*metric
cacheStatus string
}
// New creates a collector whose total duration starts at startedAt.
func New(startedAt time.Time) *Collector {
if startedAt.IsZero() {
startedAt = time.Now()
}
return &Collector{
startedAt: startedAt,
metrics: make(map[string]*metric),
}
}
// WithCollector attaches a collector to a context.
func WithCollector(ctx context.Context, collector *Collector) context.Context {
if ctx == nil {
ctx = context.Background()
}
if collector == nil {
return ctx
}
return context.WithValue(ctx, contextKey{}, collector)
}
// FromContext returns the request timing collector, when one is active.
func FromContext(ctx context.Context) (*Collector, bool) {
if ctx == nil {
return nil, false
}
collector, ok := ctx.Value(contextKey{}).(*Collector)
return collector, ok && collector != nil
}
// Active reports whether timing collection is enabled for this request.
func Active(ctx context.Context) bool {
_, ok := FromContext(ctx)
return ok
}
// Record adds a completed interval and operation count to a metric.
func Record(ctx context.Context, name string, startedAt, endedAt time.Time, count int) {
collector, ok := FromContext(ctx)
if !ok {
return
}
collector.Record(name, startedAt, endedAt, count)
}
// RecordInterval adds timing without incrementing the operation count. It is
// useful when one logical operation has multiple blocking driver calls.
func RecordInterval(ctx context.Context, name string, startedAt, endedAt time.Time) {
collector, ok := FromContext(ctx)
if !ok {
return
}
collector.record(name, startedAt, endedAt, 0)
}
// Record adds a completed interval directly to the collector.
func (c *Collector) Record(name string, startedAt, endedAt time.Time, count int) {
if count <= 0 {
count = 1
}
c.record(name, startedAt, endedAt, count)
}
func (c *Collector) record(name string, startedAt, endedAt time.Time, count int) {
name = normalizeMetricName(name)
if c == nil || name == "" || startedAt.IsZero() || endedAt.Before(startedAt) {
return
}
if count < 0 {
count = 0
}
c.mu.Lock()
m := c.metrics[name]
if m == nil {
m = &metric{}
c.metrics[name] = m
}
m.count += int64(count)
if len(m.intervals) < maxIntervals {
m.intervals = append(m.intervals, interval{start: startedAt, end: endedAt})
}
c.mu.Unlock()
}
// Observe starts a metric span and returns an idempotent completion function.
func Observe(ctx context.Context, name string) func() {
collector, ok := FromContext(ctx)
name = normalizeMetricName(name)
if !ok || name == "" {
return func() {}
}
startedAt := time.Now()
var once sync.Once
return func() {
once.Do(func() {
collector.Record(name, startedAt, time.Now(), 1)
})
}
}
// ObserveDependency starts a named external dependency span.
func ObserveDependency(ctx context.Context, module string) func() {
return Observe(ctx, dependencyMetricName(module))
}
// RecordDependency records a completed external dependency interval.
func RecordDependency(ctx context.Context, module string, startedAt, endedAt time.Time) {
Record(ctx, dependencyMetricName(module), startedAt, endedAt, 1)
}
// SetCacheStatus records the response-cache outcome for the request.
func SetCacheStatus(ctx context.Context, status string) {
collector, ok := FromContext(ctx)
if !ok {
return
}
status = normalizeCacheStatus(status)
if status == "" {
return
}
collector.mu.Lock()
collector.cacheStatus = status
collector.mu.Unlock()
}
// HeaderValue renders a bounded, deterministic Server-Timing header.
func HeaderValue(ctx context.Context, endedAt time.Time, cacheStatus string) string {
collector, ok := FromContext(ctx)
if !ok {
return ""
}
return collector.HeaderValue(endedAt, cacheStatus)
}
// HeaderValue renders a bounded, deterministic Server-Timing header.
func (c *Collector) HeaderValue(endedAt time.Time, cacheStatus string) string {
if c == nil {
return ""
}
if endedAt.IsZero() {
endedAt = time.Now()
}
if endedAt.Before(c.startedAt) {
endedAt = c.startedAt
}
c.mu.Lock()
metrics := make(map[string]metric, len(c.metrics))
allIntervals := make([]interval, 0)
dependencyIntervals := make([]interval, 0)
var dependencyCount int64
for name, source := range c.metrics {
copied := metric{count: source.count, intervals: append([]interval(nil), source.intervals...)}
metrics[name] = copied
allIntervals = append(allIntervals, copied.intervals...)
if strings.HasPrefix(name, dependencyPrefix) {
dependencyIntervals = append(dependencyIntervals, copied.intervals...)
dependencyCount += copied.count
}
}
storedCacheStatus := c.cacheStatus
c.mu.Unlock()
total := endedAt.Sub(c.startedAt)
blocked := unionDuration(allIntervals, c.startedAt, endedAt)
app := total - blocked
if app < 0 {
app = 0
}
cacheStatus = normalizeCacheStatus(cacheStatus)
if cacheStatus == "" {
cacheStatus = normalizeCacheStatus(storedCacheStatus)
}
if cacheStatus == "" {
cacheStatus = "bypass"
}
database := metrics[MetricDatabase]
redisMetric := metrics[MetricRedis]
parts := []string{
"total;dur=" + formatDuration(total),
"app;dur=" + formatDuration(app),
fmt.Sprintf("db;dur=%s;desc=\"queries=%d\"", formatDuration(unionDuration(database.intervals, c.startedAt, endedAt)), database.count),
fmt.Sprintf("redis;dur=%s;desc=\"commands=%d\"", formatDuration(unionDuration(redisMetric.intervals, c.startedAt, endedAt)), redisMetric.count),
"cache;desc=\"" + cacheStatus + "\"",
fmt.Sprintf("deps;dur=%s;desc=\"calls=%d\"", formatDuration(unionDuration(dependencyIntervals, c.startedAt, endedAt)), dependencyCount),
}
dependencyNames := make([]string, 0)
for name := range metrics {
if strings.HasPrefix(name, dependencyPrefix) {
dependencyNames = append(dependencyNames, name)
}
}
sort.Strings(dependencyNames)
for _, name := range dependencyNames {
m := metrics[name]
part := fmt.Sprintf("%s;dur=%s;desc=\"calls=%d\"", name, formatDuration(unionDuration(m.intervals, c.startedAt, endedAt)), m.count)
candidate := strings.Join(append(parts, part), ", ")
if len(candidate) > maxHeaderLength {
break
}
parts = append(parts, part)
}
return strings.Join(parts, ", ")
}
func dependencyMetricName(module string) string {
module = normalizeMetricName(module)
module = strings.TrimPrefix(module, dependencyPrefix)
if module == "" {
module = "http"
}
return dependencyPrefix + module
}
func normalizeMetricName(name string) string {
name = strings.ToLower(strings.TrimSpace(name))
if name == "" {
return ""
}
var b strings.Builder
b.Grow(min(len(name), maxMetricNameLength))
for _, r := range name {
if b.Len() >= maxMetricNameLength {
break
}
switch {
case r >= 'a' && r <= 'z', r >= '0' && r <= '9':
_, _ = b.WriteRune(r)
case r == '_' || r == '-':
_ = b.WriteByte('_')
}
}
return strings.Trim(b.String(), "_")
}
func normalizeCacheStatus(status string) string {
switch strings.ToLower(strings.TrimSpace(status)) {
case "hit":
return "hit"
case "miss":
return "miss"
case "bypass":
return "bypass"
default:
return ""
}
}
func unionDuration(intervals []interval, lowerBound, upperBound time.Time) time.Duration {
if len(intervals) == 0 || !upperBound.After(lowerBound) {
return 0
}
normalized := make([]interval, 0, len(intervals))
for _, item := range intervals {
start := item.start
end := item.end
if start.Before(lowerBound) {
start = lowerBound
}
if end.After(upperBound) {
end = upperBound
}
if end.After(start) {
normalized = append(normalized, interval{start: start, end: end})
}
}
if len(normalized) == 0 {
return 0
}
sort.Slice(normalized, func(i, j int) bool {
return normalized[i].start.Before(normalized[j].start)
})
currentStart := normalized[0].start
currentEnd := normalized[0].end
var total time.Duration
for _, item := range normalized[1:] {
if !item.start.After(currentEnd) {
if item.end.After(currentEnd) {
currentEnd = item.end
}
continue
}
total += currentEnd.Sub(currentStart)
currentStart = item.start
currentEnd = item.end
}
total += currentEnd.Sub(currentStart)
return total
}
func formatDuration(value time.Duration) string {
if value < 0 {
value = 0
}
return strconv.FormatFloat(float64(value)/float64(time.Millisecond), 'f', 1, 64)
}
@@ -0,0 +1,129 @@
package servertiming
import (
"context"
"fmt"
"strings"
"sync"
"testing"
"time"
)
func TestCollectorHeaderValueAggregatesIntervals(t *testing.T) {
startedAt := time.Unix(100, 0)
collector := New(startedAt)
collector.Record(MetricDatabase, startedAt.Add(10*time.Millisecond), startedAt.Add(40*time.Millisecond), 2)
collector.Record(MetricRedis, startedAt.Add(30*time.Millisecond), startedAt.Add(50*time.Millisecond), 3)
collector.Record(dependencyMetricName("openai"), startedAt.Add(70*time.Millisecond), startedAt.Add(100*time.Millisecond), 1)
collector.Record(dependencyMetricName("github"), startedAt.Add(60*time.Millisecond), startedAt.Add(90*time.Millisecond), 1)
got := collector.HeaderValue(startedAt.Add(120*time.Millisecond), "miss")
want := `total;dur=120.0, app;dur=40.0, db;dur=30.0;desc="queries=2", redis;dur=20.0;desc="commands=3", cache;desc="miss", deps;dur=40.0;desc="calls=2", dep_github;dur=30.0;desc="calls=1", dep_openai;dur=30.0;desc="calls=1"`
if got != want {
t.Fatalf("HeaderValue() = %q, want %q", got, want)
}
}
func TestRecordIntervalDoesNotIncrementCount(t *testing.T) {
startedAt := time.Unix(200, 0)
collector := New(startedAt)
ctx := WithCollector(context.Background(), collector)
Record(ctx, MetricDatabase, startedAt.Add(10*time.Millisecond), startedAt.Add(20*time.Millisecond), 1)
RecordInterval(ctx, MetricDatabase, startedAt.Add(30*time.Millisecond), startedAt.Add(40*time.Millisecond))
header := HeaderValue(ctx, startedAt.Add(100*time.Millisecond), "hit")
if !strings.Contains(header, `db;dur=20.0;desc="queries=1"`) {
t.Fatalf("header %q does not contain one query with both blocking intervals", header)
}
if !strings.Contains(header, "app;dur=80.0") {
t.Fatalf("header %q does not subtract the interval union from app time", header)
}
}
func TestCollectorCacheStatusFallback(t *testing.T) {
startedAt := time.Unix(300, 0)
collector := New(startedAt)
ctx := WithCollector(context.Background(), collector)
SetCacheStatus(ctx, " HIT ")
if got := HeaderValue(ctx, startedAt.Add(time.Millisecond), "invalid"); !strings.Contains(got, `cache;desc="hit"`) {
t.Fatalf("HeaderValue() = %q, want stored cache hit", got)
}
other := New(startedAt)
if got := other.HeaderValue(startedAt.Add(time.Millisecond), "invalid"); !strings.Contains(got, `cache;desc="bypass"`) {
t.Fatalf("HeaderValue() = %q, want cache bypass", got)
}
}
func TestCollectorSanitizesDependencyMetric(t *testing.T) {
startedAt := time.Unix(400, 0)
collector := New(startedAt)
ctx := WithCollector(context.Background(), collector)
RecordDependency(ctx, "GitHub API\r\nInjected;dur=999", startedAt, startedAt.Add(time.Millisecond))
header := HeaderValue(ctx, startedAt.Add(2*time.Millisecond), "bypass")
if strings.ContainsAny(header, "\r\n") || strings.Contains(header, ";dur=999") {
t.Fatalf("unsafe metric content reached header: %q", header)
}
if !strings.Contains(header, "dep_githubapiinjecteddur999;dur=1.0") {
t.Fatalf("sanitized dependency metric missing from header: %q", header)
}
}
func TestCollectorBoundsHeaderLength(t *testing.T) {
startedAt := time.Unix(500, 0)
collector := New(startedAt)
for i := 0; i < 300; i++ {
collector.Record(
dependencyMetricName(fmt.Sprintf("module_%03d_with_a_deliberately_long_name", i)),
startedAt,
startedAt.Add(time.Millisecond),
1,
)
}
header := collector.HeaderValue(startedAt.Add(2*time.Millisecond), "bypass")
if len(header) > maxHeaderLength {
t.Fatalf("header length = %d, want <= %d", len(header), maxHeaderLength)
}
if !strings.Contains(header, "total;dur=2.0") || !strings.Contains(header, "deps;dur=1.0") {
t.Fatalf("bounded header lost fixed metrics: %q", header)
}
}
func TestCollectorConcurrentRecording(t *testing.T) {
startedAt := time.Now()
collector := New(startedAt)
ctx := WithCollector(context.Background(), collector)
const workers = 25
const recordsPerWorker = 100
var wg sync.WaitGroup
wg.Add(workers)
for i := 0; i < workers; i++ {
go func() {
defer wg.Done()
for j := 0; j < recordsPerWorker; j++ {
Record(ctx, MetricDatabase, startedAt, startedAt.Add(time.Microsecond), 1)
}
}()
}
wg.Wait()
header := HeaderValue(ctx, startedAt.Add(time.Millisecond), "bypass")
want := fmt.Sprintf(`queries=%d`, workers*recordsPerWorker)
if !strings.Contains(header, want) {
t.Fatalf("header %q does not contain %q", header, want)
}
}
func TestContextHelpersHandleMissingCollector(t *testing.T) {
if Active(context.Background()) {
t.Fatal("context without collector reported active")
}
if got := HeaderValue(context.Background(), time.Now(), "hit"); got != "" {
t.Fatalf("HeaderValue() = %q without collector, want empty", got)
}
}
+104
View File
@@ -0,0 +1,104 @@
package servertiming
import (
"context"
"net/http"
"strings"
"time"
)
type dependencyModuleKey struct{}
type timingRoundTripper struct {
base http.RoundTripper
}
// WithDependencyModule overrides the safe module name used for an outbound call.
func WithDependencyModule(ctx context.Context, module string) context.Context {
if ctx == nil {
ctx = context.Background()
}
module = strings.TrimPrefix(normalizeMetricName(module), dependencyPrefix)
if module == "" {
return ctx
}
return context.WithValue(ctx, dependencyModuleKey{}, module)
}
// WrapRoundTripper records outbound response-header latency for active requests.
func WrapRoundTripper(base http.RoundTripper) http.RoundTripper {
if base == nil {
base = http.DefaultTransport
}
if _, ok := base.(*timingRoundTripper); ok {
return base
}
return &timingRoundTripper{base: base}
}
// InstrumentClient returns a shallow client copy with an instrumented transport.
func InstrumentClient(client *http.Client) *http.Client {
if client == nil {
client = &http.Client{}
}
copyClient := *client
copyClient.Transport = WrapRoundTripper(copyClient.Transport)
return &copyClient
}
// Do records response-header latency without changing the client's transport
// type. Use it for clients whose callers inspect or configure *http.Transport.
func Do(client *http.Client, req *http.Request) (*http.Response, error) {
if client == nil {
client = http.DefaultClient
}
if req == nil || !Active(req.Context()) {
return client.Do(req)
}
startedAt := time.Now()
response, err := client.Do(req)
RecordDependency(req.Context(), dependencyModule(req), startedAt, time.Now())
return response, err
}
func (t *timingRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
if req == nil || !Active(req.Context()) {
return t.base.RoundTrip(req)
}
startedAt := time.Now()
response, err := t.base.RoundTrip(req)
RecordDependency(req.Context(), dependencyModule(req), startedAt, time.Now())
return response, err
}
func dependencyModule(req *http.Request) string {
if req != nil {
if module, ok := req.Context().Value(dependencyModuleKey{}).(string); ok && module != "" {
return module
}
}
if req == nil || req.URL == nil {
return "http"
}
host := strings.ToLower(req.URL.Hostname())
switch {
case strings.Contains(host, "github"):
return "github"
case strings.Contains(host, "openai"):
return "openai"
case strings.Contains(host, "anthropic"):
return "anthropic"
case strings.Contains(host, "generativelanguage") || strings.Contains(host, "gemini"):
return "gemini"
case strings.Contains(host, "cloudcode") || strings.Contains(host, "antigravity"):
return "antigravity"
case strings.Contains(host, "googleapis") || strings.Contains(host, "google"):
return "google"
case strings.Contains(host, "amazonaws") || strings.Contains(host, "cloudflarestorage") || strings.Contains(host, "s3"):
return "s3"
case strings.Contains(host, "stripe") || strings.Contains(host, "airwallex") || strings.Contains(host, "alipay") || strings.Contains(host, "wechatpay") || strings.Contains(host, "paypal"):
return "payment"
default:
return "http"
}
}
@@ -0,0 +1,168 @@
package servertiming
import (
"context"
"io"
"net/http"
"strings"
"testing"
"time"
)
type roundTripFunc func(*http.Request) (*http.Response, error)
func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
return f(req)
}
type trackingBody struct {
read bool
}
func (b *trackingBody) Read(_ []byte) (int, error) {
b.read = true
return 0, io.EOF
}
func (b *trackingBody) Close() error { return nil }
func TestWrapRoundTripperRecordsResponseHeaderLatency(t *testing.T) {
startedAt := time.Now()
collector := New(startedAt)
body := &trackingBody{}
baseCalled := false
base := roundTripFunc(func(req *http.Request) (*http.Response, error) {
baseCalled = true
return &http.Response{
StatusCode: http.StatusOK,
Body: body,
Header: make(http.Header),
Request: req,
}, nil
})
req, err := http.NewRequestWithContext(WithCollector(context.Background(), collector), http.MethodGet, "https://api.github.com/repos/example/project", nil)
if err != nil {
t.Fatal(err)
}
resp, err := WrapRoundTripper(base).RoundTrip(req)
if err != nil {
t.Fatal(err)
}
defer func() { _ = resp.Body.Close() }()
if !baseCalled {
t.Fatal("base RoundTripper was not called")
}
if body.read {
t.Fatal("RoundTripper instrumentation read the response body; timing must stop at response headers")
}
header := collector.HeaderValue(time.Now(), "bypass")
if !strings.Contains(header, `dep_github;dur=`) || !strings.Contains(header, `deps;dur=`) {
t.Fatalf("dependency metrics missing from header: %q", header)
}
}
func TestWrapRoundTripperUsesContextModuleOverride(t *testing.T) {
collector := New(time.Now())
ctx := WithDependencyModule(WithCollector(context.Background(), collector), "data-managementd")
req, err := http.NewRequestWithContext(ctx, http.MethodGet, "https://private.example.test/path", nil)
if err != nil {
t.Fatal(err)
}
base := roundTripFunc(func(req *http.Request) (*http.Response, error) {
return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Header: make(http.Header), Request: req}, nil
})
if _, err := WrapRoundTripper(base).RoundTrip(req); err != nil {
t.Fatal(err)
}
header := collector.HeaderValue(time.Now(), "bypass")
if !strings.Contains(header, "dep_data_managementd") {
t.Fatalf("module override missing from header: %q", header)
}
if strings.Contains(header, "private.example") {
t.Fatalf("raw host leaked into header: %q", header)
}
}
func TestWrapRoundTripperSkipsInactiveContext(t *testing.T) {
called := false
base := roundTripFunc(func(req *http.Request) (*http.Response, error) {
called = true
return &http.Response{StatusCode: http.StatusNoContent, Body: http.NoBody, Header: make(http.Header), Request: req}, nil
})
req, err := http.NewRequest(http.MethodGet, "https://api.openai.com/v1/models", nil)
if err != nil {
t.Fatal(err)
}
if _, err := WrapRoundTripper(base).RoundTrip(req); err != nil {
t.Fatal(err)
}
if !called {
t.Fatal("inactive request did not reach base RoundTripper")
}
}
func TestDoRecordsWithoutChangingTransportType(t *testing.T) {
base := roundTripFunc(func(req *http.Request) (*http.Response, error) {
return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Header: make(http.Header), Request: req}, nil
})
client := &http.Client{Transport: base}
collector := New(time.Now())
req, err := http.NewRequestWithContext(WithCollector(context.Background(), collector), http.MethodGet, "https://api.openai.com/v1/models", nil)
if err != nil {
t.Fatal(err)
}
if _, err := Do(client, req); err != nil {
t.Fatal(err)
}
if _, ok := client.Transport.(roundTripFunc); !ok {
t.Fatalf("Do changed client transport type to %T", client.Transport)
}
if header := collector.HeaderValue(time.Now(), "bypass"); !strings.Contains(header, "dep_openai;dur=") {
t.Fatalf("dependency metric missing from header: %q", header)
}
}
func TestDependencyModuleClassification(t *testing.T) {
tests := map[string]string{
"https://api.github.com/repos/a/b": "github",
"https://api.openai.com/v1/models": "openai",
"https://api.anthropic.com/v1/messages": "anthropic",
"https://generativelanguage.googleapis.com/v1/models": "gemini",
"https://cloudcode-pa.googleapis.com/v1internal": "antigravity",
"https://storage.googleapis.com/bucket/object": "google",
"https://bucket.s3.amazonaws.com/object": "s3",
"https://api.stripe.com/v1/refunds": "payment",
"https://dependency.example.test/path": "http",
}
for rawURL, want := range tests {
req, err := http.NewRequest(http.MethodGet, rawURL, nil)
if err != nil {
t.Fatalf("NewRequest(%q): %v", rawURL, err)
}
if got := dependencyModule(req); got != want {
t.Errorf("dependencyModule(%q) = %q, want %q", rawURL, got, want)
}
}
}
func TestClientInstrumentationDoesNotMutateOriginal(t *testing.T) {
base := roundTripFunc(func(req *http.Request) (*http.Response, error) {
return &http.Response{StatusCode: http.StatusOK, Body: http.NoBody, Header: make(http.Header), Request: req}, nil
})
original := &http.Client{Transport: base, Timeout: time.Second}
instrumented := InstrumentClient(original)
if instrumented == original {
t.Fatal("InstrumentClient returned the original client")
}
if _, ok := original.Transport.(roundTripFunc); !ok {
t.Fatalf("InstrumentClient mutated the original transport to %T", original.Transport)
}
if instrumented.Timeout != original.Timeout {
t.Fatal("InstrumentClient did not preserve client settings")
}
if WrapRoundTripper(instrumented.Transport) != instrumented.Transport {
t.Fatal("WrapRoundTripper wrapped an already instrumented transport twice")
}
}
@@ -13,6 +13,7 @@ import (
"github.com/aws/aws-sdk-go-v2/credentials"
"github.com/aws/aws-sdk-go-v2/service/s3"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
"github.com/Wei-Shaw/sub2api/internal/service"
)
@@ -63,12 +64,14 @@ func (s *S3BackupStore) Upload(ctx context.Context, key string, body io.Reader,
return 0, fmt.Errorf("read body: %w", err)
}
finish := servertiming.ObserveDependency(ctx, "s3")
_, err = s.client.PutObject(ctx, &s3.PutObjectInput{
Bucket: &s.bucket,
Key: &key,
Body: bytes.NewReader(data),
ContentType: &contentType,
})
finish()
if err != nil {
return 0, fmt.Errorf("S3 PutObject: %w", err)
}
@@ -76,10 +79,12 @@ func (s *S3BackupStore) Upload(ctx context.Context, key string, body io.Reader,
}
func (s *S3BackupStore) Download(ctx context.Context, key string) (io.ReadCloser, error) {
finish := servertiming.ObserveDependency(ctx, "s3")
result, err := s.client.GetObject(ctx, &s3.GetObjectInput{
Bucket: &s.bucket,
Key: &key,
})
finish()
if err != nil {
return nil, fmt.Errorf("S3 GetObject: %w", err)
}
@@ -87,10 +92,12 @@ func (s *S3BackupStore) Download(ctx context.Context, key string) (io.ReadCloser
}
func (s *S3BackupStore) Delete(ctx context.Context, key string) error {
finish := servertiming.ObserveDependency(ctx, "s3")
_, err := s.client.DeleteObject(ctx, &s3.DeleteObjectInput{
Bucket: &s.bucket,
Key: &key,
})
finish()
return err
}
@@ -107,9 +114,11 @@ func (s *S3BackupStore) PresignURL(ctx context.Context, key string, expiry time.
}
func (s *S3BackupStore) HeadBucket(ctx context.Context) error {
finish := servertiming.ObserveDependency(ctx, "s3")
_, err := s.client.HeadBucket(ctx, &s3.HeadBucketInput{
Bucket: &s.bucket,
})
finish()
if err != nil {
return fmt.Errorf("S3 HeadBucket failed: %w", err)
}
@@ -276,5 +276,5 @@ func createReqClient(proxyURL string) (*req.Client, error) {
client.SetProxyURL(trimmed)
}
return client, nil
return instrumentReqClient(client), nil
}
+14 -4
View File
@@ -15,7 +15,7 @@ import (
"entgo.io/ent/dialect"
entsql "entgo.io/ent/dialect/sql"
_ "github.com/lib/pq" // PostgreSQL 驱动,通过副作用导入注册驱动
"github.com/lib/pq"
)
// InitEnt 初始化 Ent ORM 客户端并返回客户端实例和底层的 *sql.DB。
@@ -48,9 +48,19 @@ func InitEnt(cfg *config.Config) (*ent.Client, *sql.DB, error) {
// 使用 Ent 的 SQL 驱动打开 PostgreSQL 连接。
// dialect.Postgres 指定使用 PostgreSQL 方言进行 SQL 生成。
drv, err := entsql.Open(dialect.Postgres, dsn)
if err != nil {
return nil, nil, err
var drv *entsql.Driver
if cfg.Server.EnableServerTiming {
connector, err := pq.NewConnector(dsn)
if err != nil {
return nil, nil, err
}
drv = entsql.OpenDB(dialect.Postgres, sql.OpenDB(newServerTimingConnector(connector)))
} else {
var err error
drv, err = entsql.Open(dialect.Postgres, dsn)
if err != nil {
return nil, nil, err
}
}
applyDBPoolSettings(drv.DB(), cfg)
+3 -2
View File
@@ -25,6 +25,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/Wei-Shaw/sub2api/internal/util/urlvalidator"
@@ -186,7 +187,7 @@ func (s *httpUpstreamService) Do(req *http.Request, proxyURL string, accountID i
}
// 执行请求
resp, err := entry.client.Do(req)
resp, err := servertiming.Do(entry.client, req)
if err != nil {
s.recordOpenAIHTTP2Failure(profile, entry.protocolMode, entry.proxyKey, err)
// 请求失败,立即减少计数
@@ -243,7 +244,7 @@ func (s *httpUpstreamService) DoWithTLS(req *http.Request, proxyURL string, acco
return nil, err
}
resp, err := entry.client.Do(req)
resp, err := servertiming.Do(entry.client, req)
if err != nil {
atomic.AddInt64(&entry.inFlight, -1)
atomic.StoreInt64(&entry.lastUsed, time.Now().UnixNano())
+5 -1
View File
@@ -21,7 +21,11 @@ import (
// 2. MinIdleConns: 保持最小空闲连接,减少冷启动延迟(默认 10)
// 3. DialTimeout/ReadTimeout/WriteTimeout: 精确控制各阶段超时
func InitRedis(cfg *config.Config) *redis.Client {
return redis.NewClient(buildRedisOptions(cfg))
client := redis.NewClient(buildRedisOptions(cfg))
if cfg.Server.EnableServerTiming {
client.AddHook(serverTimingRedisHook{})
}
return client
}
// buildRedisOptions 构建 Redis 连接选项
@@ -2,11 +2,13 @@ package repository
import (
"fmt"
"net/http"
"strings"
"sync"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
"github.com/imroc/req/v3"
)
@@ -57,6 +59,7 @@ func getSharedReqClient(opts reqClientOptions) (*req.Client, error) {
if trimmed != "" {
client.SetProxyURL(trimmed)
}
client = instrumentReqClient(client)
actual, _ := sharedReqClients.LoadOrStore(key, client)
if c, ok := actual.(*req.Client); ok {
@@ -65,6 +68,17 @@ func getSharedReqClient(opts reqClientOptions) (*req.Client, error) {
return client, nil
}
func instrumentReqClient(client *req.Client) *req.Client {
if client == nil {
return nil
}
client.GetTransport().WrapRoundTripFunc(func(rt http.RoundTripper) req.HttpRoundTripFunc {
timed := servertiming.WrapRoundTripper(rt)
return timed.RoundTrip
})
return client
}
func buildReqClientKey(opts reqClientOptions) string {
return fmt.Sprintf("%s|%s|%t|%t",
strings.TrimSpace(opts.ProxyURL),
@@ -1,12 +1,17 @@
package repository
import (
"context"
"net/http"
"net/http/httptest"
"reflect"
"strings"
"sync"
"testing"
"time"
"unsafe"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
"github.com/imroc/req/v3"
"github.com/stretchr/testify/require"
)
@@ -118,3 +123,20 @@ func TestCreateGeminiReqClient_ForceHTTP2Disabled(t *testing.T) {
require.NoError(t, err)
require.Equal(t, "", forceHTTPVersion(t, client))
}
func TestInstrumentReqClientRecordsDependency(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusNoContent)
}))
defer server.Close()
collector := servertiming.New(time.Now())
ctx := servertiming.WithCollector(context.Background(), collector)
client := instrumentReqClient(req.C())
response, err := client.R().SetContext(ctx).Get(server.URL)
require.NoError(t, err)
require.Equal(t, http.StatusNoContent, response.StatusCode)
header := collector.HeaderValue(time.Now(), "bypass")
require.True(t, strings.Contains(header, "dep_http;dur="), header)
}
@@ -0,0 +1,39 @@
package repository
import (
"context"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
"github.com/redis/go-redis/v9"
)
type serverTimingRedisHook struct{}
func (serverTimingRedisHook) DialHook(next redis.DialHook) redis.DialHook {
return next
}
func (serverTimingRedisHook) ProcessHook(next redis.ProcessHook) redis.ProcessHook {
return func(ctx context.Context, cmd redis.Cmder) error {
if !servertiming.Active(ctx) {
return next(ctx, cmd)
}
startedAt := time.Now()
err := next(ctx, cmd)
servertiming.Record(ctx, servertiming.MetricRedis, startedAt, time.Now(), 1)
return err
}
}
func (serverTimingRedisHook) ProcessPipelineHook(next redis.ProcessPipelineHook) redis.ProcessPipelineHook {
return func(ctx context.Context, cmds []redis.Cmder) error {
if !servertiming.Active(ctx) {
return next(ctx, cmds)
}
startedAt := time.Now()
err := next(ctx, cmds)
servertiming.Record(ctx, servertiming.MetricRedis, startedAt, time.Now(), len(cmds))
return err
}
}
@@ -0,0 +1,63 @@
package repository
import (
"context"
"errors"
"strings"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
"github.com/redis/go-redis/v9"
)
func TestServerTimingRedisHookRecordsCommands(t *testing.T) {
collector := servertiming.New(time.Now())
ctx := servertiming.WithCollector(context.Background(), collector)
hook := serverTimingRedisHook{}
process := hook.ProcessHook(func(context.Context, redis.Cmder) error {
time.Sleep(time.Millisecond)
return errors.New("redis failure")
})
if err := process(ctx, redis.NewStringCmd(ctx, "get", "sensitive-key")); err == nil {
t.Fatal("ProcessHook did not return the underlying error")
}
pipeline := hook.ProcessPipelineHook(func(context.Context, []redis.Cmder) error {
time.Sleep(time.Millisecond)
return nil
})
commands := []redis.Cmder{
redis.NewStringCmd(ctx, "get", "first-secret"),
redis.NewStringCmd(ctx, "get", "second-secret"),
redis.NewStatusCmd(ctx, "set", "third-secret", "value"),
}
if err := pipeline(ctx, commands); err != nil {
t.Fatal(err)
}
header := collector.HeaderValue(time.Now(), "bypass")
if !strings.Contains(header, `commands=4`) {
t.Fatalf("header %q does not report one command and a three-command pipeline", header)
}
if strings.Contains(header, "secret") || strings.Contains(header, "get") {
t.Fatalf("Redis command details leaked into header: %q", header)
}
}
func TestServerTimingRedisHookSkipsInactiveContext(t *testing.T) {
called := false
hook := serverTimingRedisHook{}
process := hook.ProcessHook(func(context.Context, redis.Cmder) error {
called = true
return nil
})
ctx := context.Background()
if err := process(ctx, redis.NewStringCmd(ctx, "ping")); err != nil {
t.Fatal(err)
}
if !called {
t.Fatal("inactive Redis command did not reach the next hook")
}
}
@@ -0,0 +1,311 @@
package repository
import (
"context"
"database/sql/driver"
"errors"
"io"
"reflect"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
)
type serverTimingConnector struct {
base driver.Connector
}
func newServerTimingConnector(base driver.Connector) driver.Connector {
return &serverTimingConnector{base: base}
}
func (c *serverTimingConnector) Connect(ctx context.Context) (driver.Conn, error) {
startedAt := time.Now()
conn, err := c.base.Connect(ctx)
servertiming.RecordInterval(ctx, servertiming.MetricDatabase, startedAt, time.Now())
if err != nil {
return nil, err
}
return &serverTimingConn{Conn: conn}, nil
}
func (c *serverTimingConnector) Driver() driver.Driver {
return c.base.Driver()
}
type serverTimingConn struct {
driver.Conn
}
func (c *serverTimingConn) Prepare(query string) (driver.Stmt, error) {
stmt, err := c.Conn.Prepare(query)
if err != nil {
return nil, err
}
return &serverTimingStmt{Stmt: stmt}, nil
}
func (c *serverTimingConn) PrepareContext(ctx context.Context, query string) (driver.Stmt, error) {
startedAt := time.Now()
var (
stmt driver.Stmt
err error
)
if preparer, ok := c.Conn.(driver.ConnPrepareContext); ok {
stmt, err = preparer.PrepareContext(ctx, query)
} else {
stmt, err = c.Conn.Prepare(query)
}
servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1)
if err != nil {
return nil, err
}
return &serverTimingStmt{Stmt: stmt}, nil
}
func (c *serverTimingConn) ExecContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Result, error) {
execer, ok := c.Conn.(driver.ExecerContext)
if !ok {
return nil, driver.ErrSkip
}
startedAt := time.Now()
result, err := execer.ExecContext(ctx, query, args)
servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1)
return result, err
}
func (c *serverTimingConn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
queryer, ok := c.Conn.(driver.QueryerContext)
if !ok {
return nil, driver.ErrSkip
}
startedAt := time.Now()
rows, err := queryer.QueryContext(ctx, query, args)
servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1)
if err != nil || rows == nil {
return rows, err
}
return newServerTimingRows(ctx, rows), nil
}
func (c *serverTimingConn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, error) {
startedAt := time.Now()
var (
tx driver.Tx
err error
)
if beginner, ok := c.Conn.(driver.ConnBeginTx); ok {
tx, err = beginner.BeginTx(ctx, opts)
} else {
if opts.Isolation != driver.IsolationLevel(0) {
return nil, errors.New("driver does not support non-default isolation")
}
if opts.ReadOnly {
return nil, errors.New("driver does not support read-only transactions")
}
// The wrapper exposes ConnBeginTx, so it must retain database/sql's
// legacy fallback for drivers that only implement Conn.Begin.
tx, err = c.Conn.Begin() //nolint:staticcheck // Required driver compatibility fallback.
}
servertiming.RecordInterval(ctx, servertiming.MetricDatabase, startedAt, time.Now())
if err != nil || tx == nil {
return tx, err
}
return &serverTimingTx{Tx: tx, ctx: ctx}, nil
}
func (c *serverTimingConn) Ping(ctx context.Context) error {
if pinger, ok := c.Conn.(driver.Pinger); ok {
startedAt := time.Now()
err := pinger.Ping(ctx)
servertiming.RecordInterval(ctx, servertiming.MetricDatabase, startedAt, time.Now())
return err
}
return nil
}
func (c *serverTimingConn) ResetSession(ctx context.Context) error {
if resetter, ok := c.Conn.(driver.SessionResetter); ok {
startedAt := time.Now()
err := resetter.ResetSession(ctx)
servertiming.RecordInterval(ctx, servertiming.MetricDatabase, startedAt, time.Now())
return err
}
return nil
}
func (c *serverTimingConn) IsValid() bool {
if validator, ok := c.Conn.(driver.Validator); ok {
return validator.IsValid()
}
return true
}
func (c *serverTimingConn) CheckNamedValue(value *driver.NamedValue) error {
if checker, ok := c.Conn.(driver.NamedValueChecker); ok {
return checker.CheckNamedValue(value)
}
return driver.ErrSkip
}
type serverTimingStmt struct {
driver.Stmt
}
func (s *serverTimingStmt) ExecContext(ctx context.Context, args []driver.NamedValue) (driver.Result, error) {
startedAt := time.Now()
var (
result driver.Result
err error
)
if execer, ok := s.Stmt.(driver.StmtExecContext); ok {
result, err = execer.ExecContext(ctx, args)
} else {
var values []driver.Value
values, err = namedValues(args)
if err == nil {
// The wrapper exposes StmtExecContext and must preserve the fallback
// database/sql would use for a legacy driver statement.
result, err = s.Stmt.Exec(values) //nolint:staticcheck // Required driver compatibility fallback.
}
}
servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1)
return result, err
}
func (s *serverTimingStmt) QueryContext(ctx context.Context, args []driver.NamedValue) (driver.Rows, error) {
startedAt := time.Now()
var (
rows driver.Rows
err error
)
if queryer, ok := s.Stmt.(driver.StmtQueryContext); ok {
rows, err = queryer.QueryContext(ctx, args)
} else {
var values []driver.Value
values, err = namedValues(args)
if err == nil {
// The wrapper exposes StmtQueryContext and must preserve the fallback
// database/sql would use for a legacy driver statement.
rows, err = s.Stmt.Query(values) //nolint:staticcheck // Required driver compatibility fallback.
}
}
servertiming.Record(ctx, servertiming.MetricDatabase, startedAt, time.Now(), 1)
if err != nil || rows == nil {
return rows, err
}
return newServerTimingRows(ctx, rows), nil
}
func (s *serverTimingStmt) CheckNamedValue(value *driver.NamedValue) error {
if checker, ok := s.Stmt.(driver.NamedValueChecker); ok {
return checker.CheckNamedValue(value)
}
return driver.ErrSkip
}
func namedValues(args []driver.NamedValue) ([]driver.Value, error) {
values := make([]driver.Value, len(args))
for i, arg := range args {
if arg.Name != "" {
return nil, errors.New("named parameters are not supported")
}
values[i] = arg.Value
}
return values, nil
}
type serverTimingRows struct {
driver.Rows
ctx context.Context
}
func newServerTimingRows(ctx context.Context, rows driver.Rows) *serverTimingRows {
return &serverTimingRows{Rows: rows, ctx: ctx}
}
func (r *serverTimingRows) Close() error {
startedAt := time.Now()
err := r.Rows.Close()
servertiming.RecordInterval(r.ctx, servertiming.MetricDatabase, startedAt, time.Now())
return err
}
func (r *serverTimingRows) Next(dest []driver.Value) error {
startedAt := time.Now()
err := r.Rows.Next(dest)
servertiming.RecordInterval(r.ctx, servertiming.MetricDatabase, startedAt, time.Now())
return err
}
func (r *serverTimingRows) HasNextResultSet() bool {
if rows, ok := r.Rows.(driver.RowsNextResultSet); ok {
return rows.HasNextResultSet()
}
return false
}
func (r *serverTimingRows) NextResultSet() error {
rows, ok := r.Rows.(driver.RowsNextResultSet)
if !ok {
return io.EOF
}
startedAt := time.Now()
err := rows.NextResultSet()
servertiming.RecordInterval(r.ctx, servertiming.MetricDatabase, startedAt, time.Now())
return err
}
func (r *serverTimingRows) ColumnTypeScanType(index int) reflect.Type {
if rows, ok := r.Rows.(driver.RowsColumnTypeScanType); ok {
return rows.ColumnTypeScanType(index)
}
return reflect.TypeOf(new(any)).Elem()
}
func (r *serverTimingRows) ColumnTypeDatabaseTypeName(index int) string {
if rows, ok := r.Rows.(driver.RowsColumnTypeDatabaseTypeName); ok {
return rows.ColumnTypeDatabaseTypeName(index)
}
return ""
}
func (r *serverTimingRows) ColumnTypeLength(index int) (int64, bool) {
if rows, ok := r.Rows.(driver.RowsColumnTypeLength); ok {
return rows.ColumnTypeLength(index)
}
return 0, false
}
func (r *serverTimingRows) ColumnTypeNullable(index int) (bool, bool) {
if rows, ok := r.Rows.(driver.RowsColumnTypeNullable); ok {
return rows.ColumnTypeNullable(index)
}
return false, false
}
func (r *serverTimingRows) ColumnTypePrecisionScale(index int) (int64, int64, bool) {
if rows, ok := r.Rows.(driver.RowsColumnTypePrecisionScale); ok {
return rows.ColumnTypePrecisionScale(index)
}
return 0, 0, false
}
type serverTimingTx struct {
driver.Tx
ctx context.Context
}
func (t *serverTimingTx) Commit() error {
startedAt := time.Now()
err := t.Tx.Commit()
servertiming.RecordInterval(t.ctx, servertiming.MetricDatabase, startedAt, time.Now())
return err
}
func (t *serverTimingTx) Rollback() error {
startedAt := time.Now()
err := t.Tx.Rollback()
servertiming.RecordInterval(t.ctx, servertiming.MetricDatabase, startedAt, time.Now())
return err
}
@@ -0,0 +1,258 @@
package repository
import (
"context"
"database/sql/driver"
"io"
"regexp"
"strconv"
"strings"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
)
const fakeDriverDelay = 2 * time.Millisecond
type timingFakeDriver struct{}
func (timingFakeDriver) Open(string) (driver.Conn, error) { return newTimingFakeConn(), nil }
type timingFakeConnector struct {
conn driver.Conn
}
func (c timingFakeConnector) Connect(context.Context) (driver.Conn, error) {
time.Sleep(fakeDriverDelay)
return c.conn, nil
}
func (timingFakeConnector) Driver() driver.Driver { return timingFakeDriver{} }
type timingFakeConn struct{}
func newTimingFakeConn() *timingFakeConn { return &timingFakeConn{} }
func (c *timingFakeConn) Prepare(string) (driver.Stmt, error) {
time.Sleep(fakeDriverDelay)
return &timingFakeStmt{}, nil
}
func (c *timingFakeConn) PrepareContext(context.Context, string) (driver.Stmt, error) {
time.Sleep(fakeDriverDelay)
return &timingFakeStmt{}, nil
}
func (c *timingFakeConn) Close() error { return nil }
func (c *timingFakeConn) Begin() (driver.Tx, error) {
time.Sleep(fakeDriverDelay)
return &timingFakeTx{}, nil
}
func (c *timingFakeConn) BeginTx(context.Context, driver.TxOptions) (driver.Tx, error) {
time.Sleep(fakeDriverDelay)
return &timingFakeTx{}, nil
}
func (c *timingFakeConn) ExecContext(context.Context, string, []driver.NamedValue) (driver.Result, error) {
time.Sleep(fakeDriverDelay)
return driver.RowsAffected(1), nil
}
func (c *timingFakeConn) QueryContext(context.Context, string, []driver.NamedValue) (driver.Rows, error) {
time.Sleep(fakeDriverDelay)
return &timingFakeRows{values: [][]driver.Value{{"value"}}}, nil
}
func (c *timingFakeConn) Ping(context.Context) error {
time.Sleep(fakeDriverDelay)
return nil
}
func (c *timingFakeConn) ResetSession(context.Context) error {
time.Sleep(fakeDriverDelay)
return nil
}
type timingFakeStmt struct{}
func (s *timingFakeStmt) Close() error { return nil }
func (s *timingFakeStmt) NumInput() int { return -1 }
func (s *timingFakeStmt) Exec([]driver.Value) (driver.Result, error) {
time.Sleep(fakeDriverDelay)
return driver.RowsAffected(1), nil
}
func (s *timingFakeStmt) Query([]driver.Value) (driver.Rows, error) {
time.Sleep(fakeDriverDelay)
return &timingFakeRows{values: [][]driver.Value{{"value"}}}, nil
}
func (s *timingFakeStmt) ExecContext(context.Context, []driver.NamedValue) (driver.Result, error) {
time.Sleep(fakeDriverDelay)
return driver.RowsAffected(1), nil
}
func (s *timingFakeStmt) QueryContext(context.Context, []driver.NamedValue) (driver.Rows, error) {
time.Sleep(fakeDriverDelay)
return &timingFakeRows{values: [][]driver.Value{{"value"}}}, nil
}
type timingFakeRows struct {
values [][]driver.Value
index int
}
func (r *timingFakeRows) Columns() []string { return []string{"value"} }
func (r *timingFakeRows) Close() error {
time.Sleep(fakeDriverDelay)
return nil
}
func (r *timingFakeRows) Next(dest []driver.Value) error {
time.Sleep(fakeDriverDelay)
if r.index >= len(r.values) {
return io.EOF
}
copy(dest, r.values[r.index])
r.index++
return nil
}
type timingFakeTx struct{}
func (t *timingFakeTx) Commit() error {
time.Sleep(fakeDriverDelay)
return nil
}
func (t *timingFakeTx) Rollback() error {
time.Sleep(fakeDriverDelay)
return nil
}
func metricDuration(t *testing.T, header, metric string) float64 {
t.Helper()
re := regexp.MustCompile(`(?:^|, )` + regexp.QuoteMeta(metric) + `;dur=([0-9]+(?:\.[0-9]+)?)`)
match := re.FindStringSubmatch(header)
if len(match) != 2 {
t.Fatalf("metric %q missing from header %q", metric, header)
}
value, err := strconv.ParseFloat(match[1], 64)
if err != nil {
t.Fatalf("parse %s duration: %v", metric, err)
}
return value
}
func TestServerTimingConnectorRecordsDriverCallsWithoutRowLifetime(t *testing.T) {
startedAt := time.Now()
collector := servertiming.New(startedAt)
ctx := servertiming.WithCollector(context.Background(), collector)
wrapped := newServerTimingConnector(timingFakeConnector{conn: newTimingFakeConn()})
rawConn, err := wrapped.Connect(ctx)
if err != nil {
t.Fatal(err)
}
conn, ok := rawConn.(*serverTimingConn)
if !ok {
t.Fatalf("Connect() returned %T, want *serverTimingConn", rawConn)
}
if _, err := conn.ExecContext(ctx, "sensitive update", nil); err != nil {
t.Fatal(err)
}
rows, err := conn.QueryContext(ctx, "sensitive select", nil)
if err != nil {
t.Fatal(err)
}
values := make([]driver.Value, 1)
if err := rows.Next(values); err != nil {
t.Fatal(err)
}
// Application work between row reads must remain app time.
time.Sleep(30 * time.Millisecond)
if err := rows.Next(values); err != io.EOF {
t.Fatalf("rows.Next() = %v, want EOF", err)
}
if err := rows.Close(); err != nil {
t.Fatal(err)
}
header := collector.HeaderValue(time.Now(), "bypass")
if !strings.Contains(header, `queries=2`) {
t.Fatalf("header %q does not report two SQL operations", header)
}
if strings.Contains(header, "sensitive") {
t.Fatalf("SQL text leaked into header: %q", header)
}
if app, db := metricDuration(t, header, "app"), metricDuration(t, header, "db"); app <= db {
t.Fatalf("row processing gap was counted as DB time: app=%.1fms db=%.1fms header=%q", app, db, header)
}
}
func TestServerTimingPreparedStatementsAndTransactions(t *testing.T) {
collector := servertiming.New(time.Now())
ctx := servertiming.WithCollector(context.Background(), collector)
conn := &serverTimingConn{Conn: newTimingFakeConn()}
stmt, err := conn.PrepareContext(ctx, "prepare sensitive statement")
if err != nil {
t.Fatal(err)
}
timedStmt, ok := stmt.(*serverTimingStmt)
if !ok {
t.Fatalf("PrepareContext() returned %T, want *serverTimingStmt", stmt)
}
if _, err := timedStmt.ExecContext(ctx, nil); err != nil {
t.Fatal(err)
}
rows, err := timedStmt.QueryContext(ctx, nil)
if err != nil {
t.Fatal(err)
}
if err := rows.Close(); err != nil {
t.Fatal(err)
}
tx, err := conn.BeginTx(ctx, driver.TxOptions{})
if err != nil {
t.Fatal(err)
}
if err := tx.Commit(); err != nil {
t.Fatal(err)
}
if err := conn.Ping(ctx); err != nil {
t.Fatal(err)
}
if err := conn.ResetSession(ctx); err != nil {
t.Fatal(err)
}
header := collector.HeaderValue(time.Now(), "bypass")
if !strings.Contains(header, `queries=3`) {
t.Fatalf("header %q does not report prepare, exec, and query operations", header)
}
if metricDuration(t, header, "db") <= 0 {
t.Fatalf("DB duration was not recorded: %q", header)
}
}
func TestNamedValuesRejectNamedParameters(t *testing.T) {
if _, err := namedValues([]driver.NamedValue{{Name: "secret", Value: 1}}); err == nil {
t.Fatal("namedValues accepted a named parameter")
}
values, err := namedValues([]driver.NamedValue{{Ordinal: 1, Value: "value"}})
if err != nil {
t.Fatal(err)
}
if len(values) != 1 || values[0] != "value" {
t.Fatalf("namedValues() = %#v", values)
}
}
+2 -2
View File
@@ -52,7 +52,7 @@ func CORS(cfg config.CORSConfig) gin.HandlerFunc {
}
allowHeaders := []string{
"Content-Type", "Content-Length", "Accept-Encoding", "X-CSRF-Token", "Authorization",
"accept", "origin", "Cache-Control", "X-Requested-With", "X-API-Key",
"accept", "origin", "Cache-Control", "X-Requested-With", "X-API-Key", "X-Admin-UI-Request",
}
// OpenAI Node SDK 会发送 x-stainless-* 请求头,需在 CORS 中显式放行。
openAIProperties := []string{
@@ -83,7 +83,7 @@ func CORS(cfg config.CORSConfig) gin.HandlerFunc {
}
c.Writer.Header().Set("Access-Control-Allow-Headers", allowHeadersValue)
c.Writer.Header().Set("Access-Control-Allow-Methods", "POST, OPTIONS, GET, PUT, DELETE, PATCH")
c.Writer.Header().Set("Access-Control-Expose-Headers", "ETag")
c.Writer.Header().Set("Access-Control-Expose-Headers", "ETag, Server-Timing")
c.Writer.Header().Set("Access-Control-Max-Age", "86400")
}
// 处理预检请求
@@ -103,8 +103,10 @@ func TestCORS_AllowedOrigin_HasAllowHeaders(t *testing.T) {
// 应设置 Allow-Headers、Allow-Methods 和 Max-Age
assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Headers"),
"允许的 origin 应收到 Allow-Headers")
assert.Contains(t, w.Header().Get("Access-Control-Allow-Headers"), "X-Admin-UI-Request")
assert.NotEmpty(t, w.Header().Get("Access-Control-Allow-Methods"),
"允许的 origin 应收到 Allow-Methods")
assert.Contains(t, w.Header().Get("Access-Control-Expose-Headers"), "Server-Timing")
assert.Equal(t, "86400", w.Header().Get("Access-Control-Max-Age"),
"允许的 origin 应收到 Max-Age=86400")
assert.Equal(t, "https://allowed.example.com", w.Header().Get("Access-Control-Allow-Origin"),
@@ -0,0 +1,132 @@
package middleware
import (
"net/http"
"strings"
"sync"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
"github.com/gin-gonic/gin"
)
const (
snapshotCacheHeader = "X-Snapshot-Cache"
usageCacheHeader = "X-Usage-Stats-Cache"
)
type serverTimingResponseWriter struct {
gin.ResponseWriter
context *gin.Context
once sync.Once
}
func (w *serverTimingResponseWriter) Unwrap() http.ResponseWriter {
return w.ResponseWriter
}
// ServerTiming collects timing only for requests made by the Admin web UI.
func ServerTiming(enabled bool) gin.HandlerFunc {
if !enabled {
return func(c *gin.Context) {
c.Next()
}
}
return func(c *gin.Context) {
if !isAdminUIRequest(c) || c.Request == nil {
c.Next()
return
}
collector := servertiming.New(time.Now())
c.Request = c.Request.WithContext(servertiming.WithCollector(c.Request.Context(), collector))
writer := &serverTimingResponseWriter{
ResponseWriter: c.Writer,
context: c,
}
c.Writer = writer
c.Next()
writer.finalize()
}
}
func (w *serverTimingResponseWriter) WriteHeader(statusCode int) {
w.ResponseWriter.WriteHeader(statusCode)
}
func (w *serverTimingResponseWriter) WriteHeaderNow() {
w.finalize()
w.ResponseWriter.WriteHeaderNow()
}
func (w *serverTimingResponseWriter) Write(data []byte) (int, error) {
w.finalize()
return w.ResponseWriter.Write(data)
}
func (w *serverTimingResponseWriter) WriteString(data string) (int, error) {
w.finalize()
return w.ResponseWriter.WriteString(data)
}
func (w *serverTimingResponseWriter) Flush() {
w.finalize()
w.ResponseWriter.Flush()
}
func (w *serverTimingResponseWriter) finalize() {
if w == nil {
return
}
w.once.Do(func() {
if value := ServerTimingHeaderValue(w.context); value != "" {
w.ResponseWriter.Header().Set(servertiming.HeaderName, value)
}
})
}
// ServerTimingHeaderValue returns a timing value only for an authenticated admin.
func ServerTimingHeaderValue(c *gin.Context) string {
if c == nil || c.Request == nil {
return ""
}
role, ok := GetUserRoleFromContext(c)
if !ok || role != "admin" {
return ""
}
return servertiming.HeaderValue(c.Request.Context(), time.Now(), responseCacheStatus(c.Writer.Header()))
}
// ServerTimingResponseHeader builds the extra header map required by WebSocket upgrades.
func ServerTimingResponseHeader(c *gin.Context) http.Header {
value := ServerTimingHeaderValue(c)
if value == "" {
return nil
}
return http.Header{servertiming.HeaderName: []string{value}}
}
func isAdminUIRequest(c *gin.Context) bool {
if c == nil || c.Request == nil || c.Request.URL == nil {
return false
}
if strings.TrimSpace(c.GetHeader(servertiming.AdminUIHeader)) == "1" {
return true
}
path := strings.TrimSpace(c.Request.URL.Path)
return path == "/api/v1/admin" || strings.HasPrefix(path, "/api/v1/admin/")
}
func responseCacheStatus(header http.Header) string {
for _, name := range []string{snapshotCacheHeader, usageCacheHeader} {
switch strings.ToLower(strings.TrimSpace(header.Get(name))) {
case "hit":
return "hit"
case "miss":
return "miss"
case "bypass":
return "bypass"
}
}
return "bypass"
}
@@ -0,0 +1,188 @@
package middleware
import (
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
"github.com/gin-gonic/gin"
)
func runServerTimingRequest(
t *testing.T,
enabled bool,
path string,
marker string,
role string,
handler gin.HandlerFunc,
) *httptest.ResponseRecorder {
t.Helper()
gin.SetMode(gin.TestMode)
engine := gin.New()
engine.Use(ServerTiming(enabled))
engine.Any("/*path", func(c *gin.Context) {
if role != "" {
c.Set(string(ContextKeyUserRole), role)
}
handler(c)
})
recorder := httptest.NewRecorder()
request := httptest.NewRequest(http.MethodGet, path, nil)
if marker != "" {
request.Header.Set(servertiming.AdminUIHeader, marker)
}
engine.ServeHTTP(recorder, request)
return recorder
}
func TestServerTimingScopesAndRoleGate(t *testing.T) {
tests := []struct {
name string
enabled bool
path string
marker string
role string
wantHeader bool
}{
{name: "disabled", enabled: false, path: "/api/v1/admin/users", role: "admin"},
{name: "admin API path", enabled: true, path: "/api/v1/admin/users", role: "admin", wantHeader: true},
{name: "shared API marked by admin UI", enabled: true, path: "/api/v1/groups/available", marker: "1", role: "admin", wantHeader: true},
{name: "non admin role", enabled: true, path: "/api/v1/groups/available", marker: "1", role: "user"},
{name: "unauthenticated public request", enabled: true, path: "/api/v1/settings/public", marker: "1"},
{name: "unmarked shared API", enabled: true, path: "/api/v1/groups/available", role: "admin"},
{name: "invalid marker", enabled: true, path: "/api/v1/groups/available", marker: "true", role: "admin"},
{name: "admin prefix boundary", enabled: true, path: "/api/v1/administrator", role: "admin"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
recorder := runServerTimingRequest(t, tt.enabled, tt.path, tt.marker, tt.role, func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
})
header := recorder.Header().Get(servertiming.HeaderName)
if tt.wantHeader && header == "" {
t.Fatalf("%s header missing", servertiming.HeaderName)
}
if !tt.wantHeader && header != "" {
t.Fatalf("unexpected %s header: %q", servertiming.HeaderName, header)
}
if header != "" && (!strings.Contains(header, "total;dur=") || !strings.Contains(header, `cache;desc="bypass"`)) {
t.Fatalf("incomplete timing header: %q", header)
}
})
}
}
func TestServerTimingCollectorIsRequestScoped(t *testing.T) {
active := false
recorder := runServerTimingRequest(t, true, "/api/v1/keys", "1", "admin", func(c *gin.Context) {
active = servertiming.Active(c.Request.Context())
c.Status(http.StatusNoContent)
})
if !active {
t.Fatal("collector was not attached to marked request context")
}
if recorder.Header().Get(servertiming.HeaderName) == "" {
t.Fatal("timing header missing from status-only response")
}
}
func TestServerTimingFinalizesBeforeEarlyCommit(t *testing.T) {
recorder := runServerTimingRequest(t, true, "/api/v1/admin/stream", "", "admin", func(c *gin.Context) {
c.Status(http.StatusAccepted)
c.Writer.WriteHeaderNow()
})
if got := recorder.Header().Get(servertiming.HeaderName); got == "" {
t.Fatal("timing header was not written before response commit")
}
}
func TestServerTimingFinalizesOnFlush(t *testing.T) {
recorder := runServerTimingRequest(t, true, "/api/v1/admin/export", "", "admin", func(c *gin.Context) {
c.Writer.Flush()
})
if got := recorder.Header().Get(servertiming.HeaderName); got == "" {
t.Fatal("timing header was not written before stream flush")
}
}
func TestServerTimingStatusResponses(t *testing.T) {
tests := []struct {
name string
status int
}{
{name: "not modified", status: http.StatusNotModified},
{name: "internal error", status: http.StatusInternalServerError},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
recorder := runServerTimingRequest(t, true, "/api/v1/admin/test", "", "admin", func(c *gin.Context) {
c.Status(tt.status)
})
if recorder.Code != tt.status {
t.Fatalf("status = %d, want %d", recorder.Code, tt.status)
}
if got := recorder.Header().Get(servertiming.HeaderName); got == "" {
t.Fatalf("timing header missing from status %d response", tt.status)
}
})
}
}
func TestServerTimingResponseWriterUnwraps(t *testing.T) {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
baseWriter := c.Writer
writer := &serverTimingResponseWriter{ResponseWriter: baseWriter}
if got := writer.Unwrap(); got != baseWriter {
t.Fatalf("Unwrap() = %T, want original Gin writer", got)
}
}
func TestServerTimingCacheOutcome(t *testing.T) {
tests := []struct {
name string
headerName string
value string
want string
}{
{name: "snapshot hit", headerName: snapshotCacheHeader, value: "hit", want: "hit"},
{name: "usage miss", headerName: usageCacheHeader, value: "MISS", want: "miss"},
{name: "invalid", headerName: snapshotCacheHeader, value: "stale", want: "bypass"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
recorder := runServerTimingRequest(t, true, "/api/v1/admin/dashboard", "", "admin", func(c *gin.Context) {
c.Header(tt.headerName, tt.value)
c.JSON(http.StatusOK, gin.H{"ok": true})
})
want := `cache;desc="` + tt.want + `"`
if got := recorder.Header().Get(servertiming.HeaderName); !strings.Contains(got, want) {
t.Fatalf("timing header %q does not contain %q", got, want)
}
})
}
}
func TestServerTimingResponseHeaderForWebSocket(t *testing.T) {
recorder := httptest.NewRecorder()
c, _ := gin.CreateTestContext(recorder)
c.Request = httptest.NewRequest(http.MethodGet, "/api/v1/admin/ops/ws/qps", nil)
collector := servertiming.New(time.Now())
c.Request = c.Request.WithContext(servertiming.WithCollector(c.Request.Context(), collector))
c.Set(string(ContextKeyUserRole), "admin")
header := ServerTimingResponseHeader(c)
if header.Get(servertiming.HeaderName) == "" {
t.Fatal("WebSocket response header missing timing value")
}
c.Set(string(ContextKeyUserRole), "user")
if got := ServerTimingResponseHeader(c); got != nil {
t.Fatalf("non-admin WebSocket received timing header: %#v", got)
}
}
+1
View File
@@ -60,6 +60,7 @@ func SetupRouter(
}
return nil
}))
r.Use(middleware2.ServerTiming(cfg.Server.EnableServerTiming))
// Serve embedded frontend with settings injection if available
if web.HasEmbeddedFrontend() {
@@ -13,6 +13,7 @@ import (
"strings"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
"github.com/tidwall/gjson"
)
@@ -34,7 +35,7 @@ func newSSRFSafeHTTPClient(timeout time.Duration) *http.Client {
TLSHandshakeTimeout: monitorTLSHandshakeTimeout,
ResponseHeaderTimeout: monitorResponseHeaderTimeout,
}
return &http.Client{Timeout: timeout, Transport: tr}
return &http.Client{Timeout: timeout, Transport: servertiming.WrapRoundTripper(tr)}
}
// CheckOptions 承载一次检测的自定义入参。
@@ -22,6 +22,7 @@ import (
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
)
const (
@@ -561,7 +562,7 @@ func NewContentModerationService(
userRepo: userRepo,
authCacheInvalidator: authCacheInvalidator,
emailService: emailService,
httpClient: &http.Client{},
httpClient: servertiming.InstrumentClient(nil),
workerCount: maxContentModerationWorkerCount,
asyncQueue: make(chan contentModerationTask, maxContentModerationQueueSize),
keyHealth: make(map[string]*contentModerationKeyHealth),
@@ -16,6 +16,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/payment"
"github.com/Wei-Shaw/sub2api/internal/payment/provider"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
"github.com/shopspring/decimal"
)
@@ -445,7 +446,9 @@ func (s *PaymentService) invokeProvider(ctx context.Context, order *dbent.Paymen
IsMobile: req.IsMobile,
ReturnURL: providerReturnURL,
}, sel, outTradeNo, payAmountStr, subject)
finishProviderCall := servertiming.ObserveDependency(ctx, "payment")
pr, err := prov.CreatePayment(ctx, providerReq)
finishProviderCall()
if err != nil {
slog.Error("[PaymentService] CreatePayment failed", "provider", sel.ProviderKey, "instance", sel.InstanceID, "error", err)
if appErr := new(infraerrors.ApplicationError); errors.As(err, &appErr) {
@@ -13,6 +13,7 @@ import (
"github.com/Wei-Shaw/sub2api/ent/paymentorder"
"github.com/Wei-Shaw/sub2api/internal/payment"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
)
// --- Cancel & Expire ---
@@ -157,7 +158,9 @@ func (s *PaymentService) checkPaidWithOptions(ctx context.Context, o *dbent.Paym
if queryRef == "" {
return ""
}
finishProviderCall := servertiming.ObserveDependency(ctx, "payment")
resp, err := prov.QueryOrder(ctx, queryRef)
finishProviderCall()
if err != nil {
slog.Warn("query upstream failed", "orderID", o.ID, "error", err)
return ""
@@ -199,7 +202,9 @@ func (s *PaymentService) checkPaidWithOptions(ctx context.Context, o *dbent.Paym
return ""
}
if cp, ok := prov.(payment.CancelableProvider); ok {
finishProviderCall := servertiming.ObserveDependency(ctx, "payment")
_ = cp.CancelPayment(ctx, queryRef)
finishProviderCall()
}
return ""
}
@@ -208,7 +213,9 @@ func requeryPaidOrderOnce(ctx context.Context, prov payment.Provider, queryRef s
if prov == nil || strings.TrimSpace(queryRef) == "" {
return nil, false
}
finishProviderCall := servertiming.ObserveDependency(ctx, "payment")
resp, err := prov.QueryOrder(ctx, queryRef)
finishProviderCall()
if err != nil {
slog.Warn("query upstream retry failed", "queryRef", queryRef, "error", err)
return nil, false
@@ -19,6 +19,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/payment"
"github.com/Wei-Shaw/sub2api/internal/payment/provider"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
)
// --- Refund Flow ---
@@ -347,12 +348,14 @@ func (s *PaymentService) gwRefund(ctx context.Context, p *RefundPlan) (*payment.
})
return nil, err
}
finishProviderCall := servertiming.ObserveDependency(ctx, "payment")
resp, err := prov.Refund(ctx, payment.RefundRequest{
TradeNo: p.Order.PaymentTradeNo,
OrderID: p.Order.OutTradeNo,
Amount: formatGatewayRefundAmount(p.GatewayAmount, p.Order),
Reason: p.Reason,
})
finishProviderCall()
if err != nil {
if resp != nil && strings.TrimSpace(resp.Status) == payment.ProviderStatusPending {
return resp, nil
@@ -417,12 +420,14 @@ func (s *PaymentService) QueryAndFinalizeRefund(ctx context.Context, oid int64)
}
pendingDetail := s.latestRefundPendingDetail(ctx, oid)
finishProviderCall := servertiming.ObserveDependency(ctx, "payment")
resp, err := queryProvider.QueryRefund(ctx, payment.RefundQueryRequest{
TradeNo: o.PaymentTradeNo,
OrderID: o.OutTradeNo,
RefundID: pendingDetail.RefundID,
Amount: formatGatewayRefundAmount(o.RefundAmount, o),
})
finishProviderCall()
if err != nil {
return nil, fmt.Errorf("query refund: %w", err)
}
@@ -18,6 +18,7 @@ import (
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyurl"
"github.com/Wei-Shaw/sub2api/internal/pkg/proxyutil"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
"github.com/golang-jwt/jwt/v5"
)
@@ -195,7 +196,7 @@ func vertexServiceAccountProxyURL(account *Account) string {
func newVertexServiceAccountHTTPClient(proxyURL string) (*http.Client, error) {
proxyURL = strings.TrimSpace(proxyURL)
if proxyURL == "" {
return &http.Client{Timeout: 15 * time.Second}, nil
return servertiming.InstrumentClient(&http.Client{Timeout: 15 * time.Second}), nil
}
_, parsedProxy, err := proxyurl.Parse(proxyURL)
@@ -211,7 +212,7 @@ func newVertexServiceAccountHTTPClient(proxyURL string) (*http.Client, error) {
if err := proxyutil.ConfigureTransportProxy(transport, parsedProxy); err != nil {
return nil, err
}
return &http.Client{Timeout: 15 * time.Second, Transport: transport}, nil
return servertiming.InstrumentClient(&http.Client{Timeout: 15 * time.Second, Transport: transport}), nil
}
func exchangeVertexServiceAccountToken(ctx context.Context, key *vertexServiceAccountKey, proxyURL string) (string, time.Duration, error) {
@@ -13,6 +13,7 @@ import (
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
"github.com/stretchr/testify/require"
"github.com/tidwall/gjson"
)
@@ -101,6 +102,24 @@ func TestVertexServiceAccountProxyURL(t *testing.T) {
require.Empty(t, vertexServiceAccountProxyURL(&Account{ProxyID: &proxyID}))
}
func TestVertexServiceAccountHTTPClientRecordsDependency(t *testing.T) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusNoContent)
}))
defer server.Close()
client, err := newVertexServiceAccountHTTPClient("")
require.NoError(t, err)
collector := servertiming.New(time.Now())
ctx := servertiming.WithCollector(context.Background(), collector)
request, err := http.NewRequestWithContext(ctx, http.MethodGet, server.URL, nil)
require.NoError(t, err)
response, err := client.Do(request)
require.NoError(t, err)
require.NoError(t, response.Body.Close())
require.Contains(t, collector.HeaderValue(time.Now(), "bypass"), "dep_http;dur=")
}
func TestExchangeVertexServiceAccountTokenUsesProxy(t *testing.T) {
privateKey, err := rsa.GenerateKey(rand.Reader, 2048)
require.NoError(t, err)
+3
View File
@@ -23,6 +23,9 @@ SERVER_PORT=8080
# Server mode: release or debug
SERVER_MODE=release
# Return Server-Timing for authenticated requests made by the Admin web UI
ENABLE_SERVER_TIMING=false
# Apple container image overrides (ignored by Docker Compose). Pin release tags
# or digests for repeatable operator-managed deployments.
APPLE_CONTAINER_SUB2API_IMAGE=weishaw/sub2api:latest
+3
View File
@@ -20,6 +20,9 @@ server:
# Mode: "debug" for development, "release" for production
# 运行模式:"debug" 用于开发,"release" 用于生产环境
mode: "release"
# Return Server-Timing for authenticated requests made by the Admin web UI
# 为管理端 Web 页面发出的已认证请求返回 Server-Timing
enable_server_timing: false
# Frontend base URL used to generate external links in emails (e.g. password reset)
# 用于生成邮件中的外部链接(例如:重置密码链接)的前端基础地址
# Example: "https://example.com"
+1
View File
@@ -26,6 +26,7 @@ services:
- SERVER_HOST=0.0.0.0
- SERVER_PORT=8080
- SERVER_MODE=debug
- ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false}
- RUN_MODE=${RUN_MODE:-standard}
- DATABASE_HOST=postgres
- DATABASE_PORT=5432
+1
View File
@@ -51,6 +51,7 @@ services:
- SERVER_HOST=0.0.0.0
- SERVER_PORT=8080
- SERVER_MODE=${SERVER_MODE:-release}
- ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false}
- RUN_MODE=${RUN_MODE:-standard}
# =======================================================================
+1
View File
@@ -37,6 +37,7 @@ services:
- SERVER_HOST=0.0.0.0
- SERVER_PORT=8080
- SERVER_MODE=${SERVER_MODE:-release}
- ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false}
- RUN_MODE=${RUN_MODE:-standard}
# =======================================================================
+1
View File
@@ -47,6 +47,7 @@ services:
- SERVER_HOST=0.0.0.0
- SERVER_PORT=8080
- SERVER_MODE=${SERVER_MODE:-release}
- ENABLE_SERVER_TIMING=${ENABLE_SERVER_TIMING:-false}
- RUN_MODE=${RUN_MODE:-standard}
# =======================================================================
@@ -0,0 +1,38 @@
import { describe, expect, it } from 'vitest'
import {
ADMIN_UI_REQUEST_HEADER,
shouldMarkAdminUIRequest,
} from '@/api/adminUIRequest'
describe('Admin UI request marker', () => {
it('uses the stable request header name', () => {
expect(ADMIN_UI_REQUEST_HEADER).toBe('X-Admin-UI-Request')
})
it.each([
'/admin',
'/admin/users',
'/api/v1/admin',
'/api/v1/admin/accounts?status=active',
'https://api.example.test/api/v1/admin/dashboard',
])('marks Admin API request %s before page navigation', (requestURL) => {
expect(shouldMarkAdminUIRequest(requestURL, '/login')).toBe(true)
})
it.each(['/keys', '/groups/available', '/auth/me', '/announcements'])(
'marks shared request %s while an Admin page is active',
(requestURL) => {
expect(shouldMarkAdminUIRequest(requestURL, '/admin/dashboard')).toBe(true)
}
)
it.each([
['/keys', '/dashboard'],
['/api/v1/administer', '/dashboard'],
['/keys', '/administrator'],
['', '/'],
])('does not mark request %s on page %s', (requestURL, pagePath) => {
expect(shouldMarkAdminUIRequest(requestURL, pagePath)).toBe(false)
})
})
+50
View File
@@ -12,6 +12,7 @@ describe('API Client', () => {
beforeEach(async () => {
localStorage.clear()
window.history.replaceState({}, '', '/')
// 每次测试重新导入以获取干净的模块状态
vi.resetModules()
const mod = await import('@/api/client')
@@ -120,6 +121,55 @@ describe('API Client', () => {
const config = adapter.mock.calls[0][0]
expect(config.withCredentials).toBe(true)
})
it('Admin API 在进入管理页面前也带 Admin UI 标记', async () => {
const adapter = vi.fn().mockResolvedValue({
status: 200,
data: { code: 0, data: {} },
headers: {},
config: {},
statusText: 'OK',
})
apiClient.defaults.adapter = adapter
await apiClient.get('/admin/users')
const config = adapter.mock.calls[0][0]
expect(config.headers.get('X-Admin-UI-Request')).toBe('1')
})
it('管理页面调用共享 API 时带 Admin UI 标记', async () => {
window.history.replaceState({}, '', '/admin/dashboard')
const adapter = vi.fn().mockResolvedValue({
status: 200,
data: { code: 0, data: {} },
headers: {},
config: {},
statusText: 'OK',
})
apiClient.defaults.adapter = adapter
await apiClient.get('/groups/available')
const config = adapter.mock.calls[0][0]
expect(config.headers.get('X-Admin-UI-Request')).toBe('1')
})
it('普通用户页面调用共享 API 时不带 Admin UI 标记', async () => {
const adapter = vi.fn().mockResolvedValue({
status: 200,
data: { code: 0, data: {} },
headers: {},
config: {},
statusText: 'OK',
})
apiClient.defaults.adapter = adapter
await apiClient.get('/groups/available')
const config = adapter.mock.calls[0][0]
expect(config.headers.get('X-Admin-UI-Request')).toBeFalsy()
})
})
// --- 响应拦截器 ---
+27
View File
@@ -0,0 +1,27 @@
export const ADMIN_UI_REQUEST_HEADER = 'X-Admin-UI-Request'
function isAdminPath(path: string): boolean {
return (
path === '/admin' ||
path.startsWith('/admin/') ||
path === '/api/v1/admin' ||
path.startsWith('/api/v1/admin/')
)
}
function requestPath(rawURL: string): string {
const value = rawURL.trim()
if (!value) return ''
try {
const origin = typeof window !== 'undefined' ? window.location.origin : 'http://localhost'
return new URL(value, origin).pathname
} catch {
return value.split(/[?#]/, 1)[0]
}
}
export function shouldMarkAdminUIRequest(requestURL: string, pagePath?: string): boolean {
const currentPath =
pagePath ?? (typeof window !== 'undefined' ? window.location.pathname : '')
return isAdminPath(requestPath(requestURL)) || isAdminPath(currentPath)
}
+5
View File
@@ -6,6 +6,7 @@
import axios, { AxiosInstance, AxiosError, InternalAxiosRequestConfig, AxiosResponse } from 'axios'
import type { ApiResponse } from '@/types'
import { getLocale } from '@/i18n'
import { ADMIN_UI_REQUEST_HEADER, shouldMarkAdminUIRequest } from './adminUIRequest'
import { getAPIBaseURL } from './url'
export { buildApiUrl, buildGatewayUrl } from './url'
@@ -74,6 +75,10 @@ apiClient.interceptors.request.use(
config.params.timezone = getUserTimezone()
}
if (config.headers && shouldMarkAdminUIRequest(String(config.url || ''))) {
config.headers[ADMIN_UI_REQUEST_HEADER] = '1'
}
return config
},
(error) => {
@@ -250,6 +250,7 @@ import TextArea from '@/components/common/TextArea.vue'
import { Icon } from '@/components/icons'
import { useClipboard } from '@/composables/useClipboard'
import { buildApiUrl } from '@/api/client'
import { ADMIN_UI_REQUEST_HEADER } from '@/api/adminUIRequest'
import { adminAPI } from '@/api/admin'
import type { Account, ClaudeModel } from '@/types'
@@ -438,7 +439,8 @@ const startTest = async () => {
method: 'POST',
headers: {
Authorization: `Bearer ${localStorage.getItem('auth_token')}`,
'Content-Type': 'application/json'
'Content-Type': 'application/json',
[ADMIN_UI_REQUEST_HEADER]: '1'
},
body: JSON.stringify(requestBody),
signal: abortController.signal