middleware

cmd/s3-balance/main.go

@@ -70,8 +70,8 @@ func main() {
 
 	// 创建预签名URL生成器
 	signer := presigner.NewPresigner(
-		15*time.Minute,  // 上传URL有效期
+		15*time.Minute, // 上传URL有效期
-		60*time.Minute,  // 下载URL有效期
+		60*time.Minute, // 下载URL有效期
 	)
 
 	// 创建存储服务

internal/api/auth_middleware.go

@@ -1,57 +0,0 @@
-package api
-
-import (
-	"encoding/base64"
-	"net/http"
-	"strings"
-)
-
-// authMiddleware 处理 Basic Auth 校验
-func (h *S3Handler) authMiddleware(next http.Handler) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		if !h.authRequired {
-			next.ServeHTTP(w, r)
-			return
-		}
-
-		authHeader := r.Header.Get("Authorization")
-		if authHeader == "" {
-			h.requireAuth(w)
-			return
-		}
-
-		if strings.HasPrefix(authHeader, "Basic ") {
-			payload := strings.TrimPrefix(authHeader, "Basic ")
-			decoded, err := base64.StdEncoding.DecodeString(payload)
-			if err != nil {
-				h.requireAuth(w)
-				return
-			}
-
-			parts := strings.SplitN(string(decoded), ":", 2)
-			if len(parts) != 2 {
-				h.requireAuth(w)
-				return
-			}
-
-			if parts[0] != h.accessKey {
-				h.sendS3Error(w, "InvalidAccessKeyId", "The AWS Access Key Id you provided does not match the configured key.", "")
-				return
-			}
-			if parts[1] != h.secretKey {
-				h.sendS3Error(w, "SignatureDoesNotMatch", "The request signature we calculated does not match the signature you provided.", "")
-				return
-			}
-
-			next.ServeHTTP(w, r)
-			return
-		}
-
-		h.requireAuth(w)
-	})
-}
-
-func (h *S3Handler) requireAuth(w http.ResponseWriter) {
-	w.Header().Set("WWW-Authenticate", "Basic realm=\"s3-balance\"")
-	h.sendS3Error(w, "AccessDenied", "Access Denied", "")
-}

internal/api/s3_handler.go

@@ -4,6 +4,7 @@ import (
 	"github.com/DullJZ/s3-balance/internal/balancer"
 	"github.com/DullJZ/s3-balance/internal/bucket"
 	"github.com/DullJZ/s3-balance/internal/metrics"
+	"github.com/DullJZ/s3-balance/internal/middleware"
 	"github.com/DullJZ/s3-balance/internal/storage"
 	"github.com/DullJZ/s3-balance/pkg/presigner"
 	"github.com/gorilla/mux"
@@ -75,7 +76,18 @@ func (h *S3Handler) RegisterS3Routes(router *mux.Router) {
 	// Object operations - must be registered after multipart operations to avoid conflicts
 	router.HandleFunc("/{bucket}/{key:.*}", h.handleObjectOperations).Methods("GET", "HEAD", "PUT", "DELETE")
 
-	// 添加认证中间件
+	// 添加中间件
-	router.Use(h.virtualHostMiddleware)
+	router.Use(middleware.VirtualHost(middleware.VirtualHostConfig{
-	router.Use(h.authMiddleware)
+		Enabled: h.virtualHost,
+		BucketExists: func(name string) bool {
+			_, ok := h.bucketManager.GetBucket(name)
+			return ok
+		},
+	}))
+	router.Use(middleware.BasicAuth(middleware.AuthConfig{
+		Required:  h.authRequired,
+		AccessKey: h.accessKey,
+		SecretKey: h.secretKey,
+		OnError:   h.sendS3Error,
+	}))
 }

internal/api/virtual_host_middleware.go

@@ -1,72 +0,0 @@
-package api
-
-import (
-	"net"
-	"net/http"
-	"strings"
-)
-
-// virtualHostMiddleware 支持根据 Host 头推断存储桶名称
-func (h *S3Handler) virtualHostMiddleware(next http.Handler) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		if !h.virtualHost {
-			next.ServeHTTP(w, r)
-			return
-		}
-
-		bucketName := h.bucketFromHost(r.Host)
-		if bucketName == "" {
-			next.ServeHTTP(w, r)
-			return
-		}
-
-		// 若路由中已包含桶名称则无需改写
-		if strings.HasPrefix(r.URL.Path, "/"+bucketName) {
-			next.ServeHTTP(w, r)
-			return
-		}
-
-		// 确保桶存在
-		if _, ok := h.bucketManager.GetBucket(bucketName); !ok {
-			next.ServeHTTP(w, r)
-			return
-		}
-
-		newPath := "/" + bucketName
-		if r.URL.Path != "/" {
-			newPath += r.URL.Path
-		}
-
-		clone := r.Clone(r.Context())
-		clone.URL.Path = newPath
-		clone.RequestURI = newPath
-
-		next.ServeHTTP(w, clone)
-	})
-}
-
-func (h *S3Handler) bucketFromHost(host string) string {
-	if host == "" {
-		return ""
-	}
-
-	cleanHost := host
-	if strings.Contains(host, ":") {
-		hostname, _, err := net.SplitHostPort(host)
-		if err == nil {
-			cleanHost = hostname
-		}
-	}
-
-	parts := strings.Split(cleanHost, ".")
-	if len(parts) == 0 {
-		return ""
-	}
-
-	candidate := parts[0]
-	if candidate == "" {
-		return ""
-	}
-
-	return candidate
-}

internal/middleware/auth.go

@@ -0,0 +1,70 @@
+package middleware
+
+import (
+	"encoding/base64"
+	"net/http"
+	"strings"
+)
+
+// AuthConfig controls Basic Auth validation.
+type AuthConfig struct {
+	Required  bool
+	AccessKey string
+	SecretKey string
+	OnError   func(http.ResponseWriter, string, string, string)
+}
+
+// BasicAuth enforces static access/secret key authentication when Required is true.
+func BasicAuth(cfg AuthConfig) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if !cfg.Required {
+				next.ServeHTTP(w, r)
+				return
+			}
+
+			authHeader := r.Header.Get("Authorization")
+			if !strings.HasPrefix(authHeader, "Basic ") {
+				requireAuth(w, cfg)
+				return
+			}
+
+			payload := strings.TrimPrefix(authHeader, "Basic ")
+			decoded, err := base64.StdEncoding.DecodeString(payload)
+			if err != nil {
+				requireAuth(w, cfg)
+				return
+			}
+
+			parts := strings.SplitN(string(decoded), ":", 2)
+			if len(parts) != 2 {
+				requireAuth(w, cfg)
+				return
+			}
+
+			if parts[0] != cfg.AccessKey {
+				invokeOnError(w, cfg, "InvalidAccessKeyId", "The AWS Access Key Id you provided does not match the configured key.")
+				return
+			}
+			if parts[1] != cfg.SecretKey {
+				invokeOnError(w, cfg, "SignatureDoesNotMatch", "The request signature we calculated does not match the signature you provided.")
+				return
+			}
+
+			next.ServeHTTP(w, r)
+		})
+	}
+}
+
+func requireAuth(w http.ResponseWriter, cfg AuthConfig) {
+	w.Header().Set("WWW-Authenticate", "Basic realm=\"s3-balance\"")
+	invokeOnError(w, cfg, "AccessDenied", "Access Denied")
+}
+
+func invokeOnError(w http.ResponseWriter, cfg AuthConfig, code, message string) {
+	if cfg.OnError != nil {
+		cfg.OnError(w, code, message, "")
+		return
+	}
+	http.Error(w, message, http.StatusForbidden)
+}

internal/middleware/virtual_host.go

@@ -0,0 +1,73 @@
+package middleware
+
+import (
+	"net"
+	"net/http"
+	"strings"
+)
+
+// VirtualHostConfig controls host-style bucket resolution.
+type VirtualHostConfig struct {
+	Enabled      bool
+	BucketExists func(string) bool
+}
+
+// VirtualHost rewrites host-style requests (bucket.example.com) into path-style paths.
+func VirtualHost(cfg VirtualHostConfig) func(http.Handler) http.Handler {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if !cfg.Enabled {
+				next.ServeHTTP(w, r)
+				return
+			}
+
+			bucket := bucketFromHost(r.Host)
+			if bucket == "" {
+				next.ServeHTTP(w, r)
+				return
+			}
+
+			if cfg.BucketExists != nil && !cfg.BucketExists(bucket) {
+				next.ServeHTTP(w, r)
+				return
+			}
+
+			if strings.HasPrefix(r.URL.Path, "/"+bucket) {
+				next.ServeHTTP(w, r)
+				return
+			}
+
+			newPath := "/" + bucket
+			if r.URL.Path != "/" {
+				newPath += r.URL.Path
+			}
+
+			clone := r.Clone(r.Context())
+			clone.URL.Path = newPath
+			clone.RequestURI = newPath
+
+			next.ServeHTTP(w, clone)
+		})
+	}
+}
+
+func bucketFromHost(host string) string {
+	if host == "" {
+		return ""
+	}
+
+	hostname := host
+	if strings.Contains(host, ":") {
+		h, _, err := net.SplitHostPort(host)
+		if err == nil {
+			hostname = h
+		}
+	}
+
+	parts := strings.Split(hostname, ".")
+	if len(parts) == 0 {
+		return ""
+	}
+
+	return parts[0]
+}