From 908881574d5f1bc8e9674210d8662f91b298cc29 Mon Sep 17 00:00:00 2001 From: DullJZ <79080562+DullJZ@users.noreply.github.com> Date: Sun, 5 Oct 2025 01:58:54 +0800 Subject: [PATCH] Support setting host & recording error req --- go.mod | 2 +- go.sum | 4 ++-- internal/api/access_log.go | 23 +++++++++++++++++++++-- internal/api/s3_handler.go | 14 +++++++------- internal/middleware/auth.go | 8 ++++---- internal/storage/models.go | 3 ++- internal/storage/service.go | 3 ++- 7 files changed, 39 insertions(+), 18 deletions(-) diff --git a/go.mod b/go.mod index 325d777..4d635e6 100644 --- a/go.mod +++ b/go.mod @@ -3,7 +3,7 @@ module github.com/DullJZ/s3-balance go 1.24.5 require ( - github.com/DullJZ/s3-validate v0.0.0-20250930120412-fc4ea70939f6 + github.com/DullJZ/s3-validate v0.0.0-20251004111253-b3ec227d3796 github.com/aws/aws-sdk-go-v2 v1.39.2 github.com/aws/aws-sdk-go-v2/config v1.31.1 github.com/aws/aws-sdk-go-v2/credentials v1.18.5 diff --git a/go.sum b/go.sum index 1a958ca..7895e0d 100644 --- a/go.sum +++ b/go.sum @@ -1,7 +1,7 @@ filippo.io/edwards25519 v1.1.0 h1:FNf4tywRC1HmFuKW5xopWpigGjJKiJSV0Cqo0cJWDaA= filippo.io/edwards25519 v1.1.0/go.mod h1:BxyFTGdWcka3PhytdK4V28tE5sGfRvvvRV7EaN4VDT4= -github.com/DullJZ/s3-validate v0.0.0-20250930120412-fc4ea70939f6 h1:UZ4i/MFU0ttINUch3GYyJayq6Y2ODm+RPawLgPna5L8= -github.com/DullJZ/s3-validate v0.0.0-20250930120412-fc4ea70939f6/go.mod h1:OEx+/bRlDdI0oj/Bb1Plsq+1+qU1qal3/g9phixhU6Y= +github.com/DullJZ/s3-validate v0.0.0-20251004111253-b3ec227d3796 h1:0Lipgc3EHF2QOKpCziXApbVocdyzZ/3a52xluuWraXg= +github.com/DullJZ/s3-validate v0.0.0-20251004111253-b3ec227d3796/go.mod h1:OEx+/bRlDdI0oj/Bb1Plsq+1+qU1qal3/g9phixhU6Y= github.com/aws/aws-sdk-go-v2 v1.39.2 h1:EJLg8IdbzgeD7xgvZ+I8M1e0fL0ptn/M47lianzth0I= github.com/aws/aws-sdk-go-v2 v1.39.2/go.mod h1:sDioUELIUO9Znk23YVmIk86/9DOpkbyyVb1i/gUNFXY= github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.0 h1:6GMWV6CNpA/6fbFHnoAjrv4+LGfyTqZz2LtCHnspgDg= diff --git a/internal/api/access_log.go b/internal/api/access_log.go index 8f8c3c7..285331a 100644 --- a/internal/api/access_log.go +++ b/internal/api/access_log.go @@ -2,6 +2,7 @@ package api import ( "bufio" + "context" "errors" "log" "net" @@ -12,6 +13,12 @@ import ( "github.com/gorilla/mux" ) +type accessLogContextKey string + +const ( + errorCodeKey accessLogContextKey = "errorCode" +) + func (h *S3Handler) accessLogMiddleware(next http.Handler) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if h.storage == nil { @@ -34,7 +41,10 @@ func (h *S3Handler) accessLogMiddleware(next http.Handler) http.Handler { success := lrw.statusCode < 400 errMsg := "" if !success { - if code := lrw.Header().Get("X-Amz-Error-Code"); code != "" { + // 优先使用context中的错误码（auth错误设置的） + if code, ok := r.Context().Value(errorCodeKey).(string); ok && code != "" { + errMsg = code + } else if code := lrw.Header().Get("X-Amz-Error-Code"); code != "" { errMsg = code } else { errMsg = http.StatusText(lrw.statusCode) @@ -50,14 +60,23 @@ func (h *S3Handler) accessLogMiddleware(next http.Handler) http.Handler { func (h *S3Handler) recordAccessLog(r *http.Request, action, bucket, key string, size int64, success bool, errMsg string, duration time.Duration) { clientIP := extractClientIP(r) userAgent := r.UserAgent() + host := r.Host // 异步记录日志，避免阻塞请求响应 go func() { - if err := h.storage.RecordAccessLog(action, key, bucket, clientIP, userAgent, size, success, errMsg, duration.Milliseconds()); err != nil { + if err := h.storage.RecordAccessLog(action, key, bucket, clientIP, userAgent, host, size, success, errMsg, duration.Milliseconds()); err != nil { log.Printf("failed to record access log: %v", err) } }() } +func (h *S3Handler) handleAuthError(w http.ResponseWriter, r *http.Request, code, message, resource string) { + // 将错误码存入context，供accessLogMiddleware使用 + ctx := context.WithValue(r.Context(), errorCodeKey, code) + *r = *r.WithContext(ctx) + + h.sendS3Error(w, code, message, resource) +} + type loggingResponseWriter struct { http.ResponseWriter statusCode int diff --git a/internal/api/s3_handler.go b/internal/api/s3_handler.go index 32b2e10..8a8c7b0 100644 --- a/internal/api/s3_handler.go +++ b/internal/api/s3_handler.go @@ -24,11 +24,11 @@ type S3Handler struct { } type handlerSettings struct { - accessKey string - secretKey string - proxyMode bool - authRequired bool - virtualHost bool + accessKey string + secretKey string + proxyMode bool + authRequired bool + virtualHost bool signatureHost string } @@ -94,6 +94,7 @@ func (h *S3Handler) RegisterS3Routes(router *mux.Router) { protected.HandleFunc("/{key:.*}", h.handleObjectOperations).Methods("GET", "HEAD", "PUT", "DELETE") // 添加中间件 + protected.Use(h.accessLogMiddleware) protected.Use(middleware.VirtualHost(middleware.VirtualHostConfig{ Enabled: h.virtualHostEnabled, BucketExists: func(name string) bool { @@ -104,10 +105,9 @@ func (h *S3Handler) RegisterS3Routes(router *mux.Router) { protected.Use(middleware.S3Signature(middleware.S3SignatureConfig{ Required: h.authRequired, Credentials: h.credentials, - OnError: h.sendS3Error, + OnError: h.handleAuthError, SignatureHost: h.signatureHost, })) - protected.Use(h.accessLogMiddleware) } func (h *S3Handler) loadSettings() handlerSettings { diff --git a/internal/middleware/auth.go b/internal/middleware/auth.go index e4f9e24..c29ffc6 100644 --- a/internal/middleware/auth.go +++ b/internal/middleware/auth.go @@ -12,7 +12,7 @@ import ( type S3SignatureConfig struct { Required func() bool Credentials func() (string, string) - OnError func(http.ResponseWriter, string, string, string) + OnError func(http.ResponseWriter, *http.Request, string, string, string) SignatureHost func() string // 用于签名验证的Host（为空则使用请求的Host） } @@ -57,7 +57,7 @@ func S3Signature(cfg S3SignatureConfig) func(http.Handler) http.Handler { result, err := verifier.Verify(r.Context(), r) if err != nil { - invokeOnError(w, cfg, "SignatureDoesNotMatch", err.Error()) + invokeOnError(w, r, cfg, "SignatureDoesNotMatch", err.Error()) return } @@ -69,9 +69,9 @@ func S3Signature(cfg S3SignatureConfig) func(http.Handler) http.Handler { } } -func invokeOnError(w http.ResponseWriter, cfg S3SignatureConfig, code, message string) { +func invokeOnError(w http.ResponseWriter, r *http.Request, cfg S3SignatureConfig, code, message string) { if cfg.OnError != nil { - cfg.OnError(w, code, message, "") + cfg.OnError(w, r, code, message, "") return } http.Error(w, message, http.StatusForbidden) diff --git a/internal/storage/models.go b/internal/storage/models.go index 9f5c24a..08a114f 100644 --- a/internal/storage/models.go +++ b/internal/storage/models.go @@ -87,7 +87,8 @@ type AccessLog struct { Size int64 `gorm:"default:0" json:"size"` ClientIP string `gorm:"size:64" json:"client_ip"` UserAgent string `gorm:"size:512" json:"user_agent"` - Success bool `gorm:"default:true" json:"success"` + Host string `gorm:"size:255" json:"host"` + Success bool `gorm:"not null" json:"success"` ErrorMsg string `gorm:"type:text" json:"error_msg,omitempty"` ResponseTime int64 `gorm:"default:0" json:"response_time"` // 响应时间（毫秒） CreatedAt time.Time `gorm:"index" json:"created_at"` diff --git a/internal/storage/service.go b/internal/storage/service.go index 1170cd8..20ad733 100644 --- a/internal/storage/service.go +++ b/internal/storage/service.go @@ -367,13 +367,14 @@ func (s *Service) CleanExpiredSessions() error { } // RecordAccessLog 记录访问日志 -func (s *Service) RecordAccessLog(action, key, bucketName, clientIP, userAgent string, size int64, success bool, errorMsg string, responseTime int64) error { +func (s *Service) RecordAccessLog(action, key, bucketName, clientIP, userAgent, host string, size int64, success bool, errorMsg string, responseTime int64) error { log := &AccessLog{ Action: action, Key: key, BucketName: bucketName, ClientIP: clientIP, UserAgent: userAgent, + Host: host, Size: size, Success: success, ErrorMsg: errorMsg,