From 322d41fbb20bc24eda8ad05796f3faadebb05a58 Mon Sep 17 00:00:00 2001
From: DullJZ <79080562+DullJZ@users.noreply.github.com>
Date: Sun, 21 Dec 2025 14:37:06 +0800
Subject: [PATCH] Feat: Token auth for metrics

---
 README.md                  |  4 ++--
 README_EN.md               |  6 +++---
 VERSION                    |  2 +-
 cmd/s3-balance/main.go     | 12 ++++++++++--
 config/config.example.yaml |  2 ++
 internal/config/config.go  |  1 +
 6 files changed, 19 insertions(+), 8 deletions(-)

diff --git a/README.md b/README.md
index a73fe67..416b5ea 100644
--- a/README.md
+++ b/README.md
@@ -88,12 +88,12 @@ go build -o s3-balance cmd/s3-balance/main.go
 - `database`：GORM 支持 sqlite/mysql/postgres，并存储对象元数据、分片会话等。
 - `buckets`：列出真实与虚拟桶。`virtual: true` 的条目会对外暴露，真实桶为 `virtual: false`。可设置 `path_style` 与 `max_size`。
 - `balancer`：策略 (`round-robin`|`least-space`|`weighted`)、健康检查周期、重试次数与延迟。
-- `metrics`：是否启用 Prometheus 指标及路径。
+- `metrics`：是否启用 Prometheus 指标、路径与可选抓取 Token。
 - `s3api`：Access/Secret Key、`proxy_mode`（true=服务代理，false=重定向）、`auth_required`（SigV4 校验）、`virtual_host`（Host-style 路由）。
 
 ## API & 测试
 
-- 默认监听 `http://localhost:8080`，支持 `GET /health` 健康检查、`GET /metrics` 指标。
+- 默认监听 `http://localhost:8080`，支持 `GET /health` 健康检查、`GET /metrics` 指标（配置 `metrics.token` 后需携带 `Authorization: Bearer <token>`）。
 - 可使用 AWS CLI、s3cmd、MinIO Client 或 `python3 test_virtual_bucket_s3.py` 验证兼容性；脚本运行前需修改 endpoint 与凭据。
 
 ## 项目结构
diff --git a/README_EN.md b/README_EN.md
index 43042ff..67e91f9 100644
--- a/README_EN.md
+++ b/README_EN.md
@@ -86,12 +86,12 @@ go build -o s3-balance cmd/s3-balance/main.go
 - `database`: GORM supports sqlite/mysql/postgres and stores object metadata, multipart sessions, etc.
 - `buckets`: Lists real and virtual buckets. Entries with `virtual: true` are exposed externally, while real buckets are marked as `virtual: false`. Supports `path_style` and `max_size` settings.
 - `balancer`: Strategy (`round-robin`|`least-space`|`weighted`), health check intervals, retry counts, and delays.
-- `metrics`: Whether to enable Prometheus metrics and their path.
+- `metrics`: Whether to enable Prometheus metrics, the path, and optional scrape token.
 - `s3api`: Access/Secret Key, `proxy_mode` (true=proxy, false=redirect), `auth_required` (SigV4 validation), `virtual_host` (host-style routing).
 
 ## API & Testing
 
-- Default listening at `http://localhost:8080`, supports `GET /health` for health checks and `GET /metrics` for metrics.
+- Default listening at `http://localhost:8080`, supports `GET /health` for health checks and `GET /metrics` for metrics (if `metrics.token` is set, send `Authorization: Bearer <token>`).
 - Compatibility can be verified using AWS CLI, s3cmd, MinIO Client, or `python3 test_virtual_bucket_s3.py`. Modify the endpoint and credentials in the script before running.
 
 ## Project Structure
@@ -106,4 +106,4 @@ internal/storage/    # GORM models and services
 pkg/presigner/       # Pre-signed URL utilities
 config/              # Example configurations and deployment manifests
 deploy/              # Docker/Kubernetes/Helm manifests
-```
\ No newline at end of file
+```
diff --git a/VERSION b/VERSION
index eac0a14..eb9e63b 100644
--- a/VERSION
+++ b/VERSION
@@ -1 +1 @@
-v0.2.1
\ No newline at end of file
+v0.2.2
\ No newline at end of file
diff --git a/cmd/s3-balance/main.go b/cmd/s3-balance/main.go
index e1de95a..df2d3c7 100644
--- a/cmd/s3-balance/main.go
+++ b/cmd/s3-balance/main.go
@@ -137,8 +137,16 @@ func main() {
 
 	// 添加指标端点
 	if cfg.Metrics.Enabled {
-		router.Path(cfg.Metrics.Path).Handler(promhttp.Handler())
-		log.Printf("Metrics server enabled at %s", cfg.Metrics.Path)
+		metricsHandler := promhttp.Handler()
+		if cfg.Metrics.Token != "" {
+			metricsHandler = middleware.TokenAuthMiddleware(cfg.Metrics.Token)(metricsHandler)
+		}
+		router.Path(cfg.Metrics.Path).Handler(metricsHandler)
+		log.Printf(
+			"Metrics server enabled at %s (auth required: %t)",
+			cfg.Metrics.Path,
+			cfg.Metrics.Token != "",
+		)
 	}
 
 	// 注册管理API路由（如果启用）
diff --git a/config/config.example.yaml b/config/config.example.yaml
index 6ee70d6..b7e02ed 100644
--- a/config/config.example.yaml
+++ b/config/config.example.yaml
@@ -128,6 +128,8 @@ balancer:
 metrics:
   enabled: true
   path: "/metrics"
+  # Prometheus 抓取时需要携带的 Token（可选，配置后需使用 Authorization: Bearer <token> 访问）
+  token: ""
 
 # S3兼容API配置
 s3api:
diff --git a/internal/config/config.go b/internal/config/config.go
index 2f3d8b7..3190020 100644
--- a/internal/config/config.go
+++ b/internal/config/config.go
@@ -63,6 +63,7 @@ type BalancerConfig struct {
 type MetricsConfig struct {
 	Enabled bool   `yaml:"enabled"`
 	Path    string `yaml:"path"`
+	Token   string `yaml:"token"` // 可选Token，保护Prometheus端点
 	// Port    int    `yaml:"port"`  // 目前未使用，与主服务共享端口
 }