LokiStack/backend/services/storage_check.go

package services

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"encoding/xml"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"regexp"
	"strings"
	"time"
)

// StorageTestResult 存储连接测试结果(逐级:可达 → 认证 → 桶)
type StorageTestResult struct {
	Reachable     bool   `json:"reachable"`
	Authenticated bool   `json:"authenticated"`
	BucketExists  bool   `json:"bucketExists"`
	BucketCreated bool   `json:"bucketCreated"`
	Message       string `json:"message"`
}

// StorageTestParams 测试参数(与创建存储配置一致,额外支持自动建桶)
type StorageTestParams struct {
	StorageType     string
	Endpoint        string
	Bucket          string
	Region          string
	AccessKey       string
	SecretKey       string
	Insecure        bool
	CreateIfMissing bool
}

// ==================== S3 SigV4(零依赖实现,兼容 MinIO/RustFS/S3/OSS S3 兼容端点) ====================

func hmacSHA256(key []byte, data string) []byte {
	h := hmac.New(sha256.New, key)
	h.Write([]byte(data))
	return h.Sum(nil)
}

func sha256Hex(s string) string {
	sum := sha256.Sum256([]byte(s))
	return hex.EncodeToString(sum[:])
}

// s3Request 发起一次 SigV4 签名的 S3 请求(路径风格),path 为 "/" 或 "/<bucket>"
func s3Request(method, endpoint, region, accessKey, secretKey, path string, insecure bool) (*http.Response, error) {
	// endpoint 解析:带 scheme 用之;裸主机名默认 https,insecure 时 http
	base := strings.TrimSpace(endpoint)
	scheme := ""
	if strings.HasPrefix(base, "http://") {
		scheme, base = "http", strings.TrimPrefix(base, "http://")
	} else if strings.HasPrefix(base, "https://") {
		scheme, base = "https", strings.TrimPrefix(base, "https://")
	} else {
		if insecure {
			scheme = "http"
		} else {
			scheme = "https"
		}
	}
	base = strings.TrimSuffix(base, "/")
	if region == "" {
		region = "us-east-1"
	}

	now := time.Now().UTC()
	amzDate := now.Format("20060102T150405Z")
	dateStamp := now.Format("20060102")
	payloadHash := sha256Hex("") // 空 body

	// 规范请求(SigV4: host/x-amz-content-sha256/x-amz-date 参与签名)
	canonicalURI := path
	if canonicalURI == "" {
		canonicalURI = "/"
	}
	canonicalRequest := strings.Join([]string{
		method,
		canonicalURI,
		"", // query
		"host:" + base + "\n" +
			"x-amz-content-sha256:" + payloadHash + "\n" +
			"x-amz-date:" + amzDate + "\n",
		"host;x-amz-content-sha256;x-amz-date",
		payloadHash,
	}, "\n")

	scope := dateStamp + "/" + region + "/s3/aws4_request"
	stringToSign := strings.Join([]string{
		"AWS4-HMAC-SHA256",
		amzDate,
		scope,
		sha256Hex(canonicalRequest),
	}, "\n")

	kDate := hmacSHA256([]byte("AWS4"+secretKey), dateStamp)
	kRegion := hmacSHA256(kDate, region)
	kService := hmacSHA256(kRegion, "s3")
	kSigning := hmacSHA256(kService, "aws4_request")
	signature := hex.EncodeToString(hmacSHA256(kSigning, stringToSign))

	req, err := http.NewRequest(method, scheme+"://"+base+canonicalURI, nil)
	if err != nil {
		return nil, err
	}
	req.Host = base
	req.Header.Set("x-amz-date", amzDate)
	req.Header.Set("x-amz-content-sha256", payloadHash)
	req.Header.Set("Authorization", fmt.Sprintf(
		"AWS4-HMAC-SHA256 Credential=%s/%s, SignedHeaders=host;x-amz-content-sha256;x-amz-date, Signature=%s",
		accessKey, scope, signature))

	client := &http.Client{Timeout: 10 * time.Second}
	return client.Do(req)
}

// s3ErrBrief 从 S3 错误响应 XML 中提取 Code/Message,避免把整段 XML 暴露给用户
func s3ErrBrief(body []byte) string {
	var e struct {
		Code    string `xml:"Code"`
		Message string `xml:"Message"`
	}
	if err := xml.Unmarshal(body, &e); err == nil && (e.Code != "" || e.Message != "") {
		if e.Code != "" && e.Message != "" {
			return e.Code + ": " + e.Message
		}
		return e.Code + e.Message
	}
	brief := strings.TrimSpace(string(body))
	if len(brief) > 120 {
		brief = brief[:120] + "..."
	}
	return brief
}

// ==================== 测试连接 ====================

// endpointHasBucket 判断 endpoint 是否已包含桶名（虚拟主机式）
func endpointHasBucket(endpoint, bucket string) bool {
	host := strings.TrimSpace(endpoint)
	if i := strings.Index(host, "://"); i >= 0 {
		host = host[i+3:]
	}
	if i := strings.IndexByte(host, '/'); i >= 0 {
		host = host[:i]
	}
	return bucket != "" && strings.HasPrefix(strings.ToLower(host), strings.ToLower(bucket)+".")
}

// stripBucketPrefix 从 endpoint 里剥掉桶名前缀，得到区域端点。
//
//	mybucket-1250000000.cos.ap-hongkong.myqcloud.com → cos.ap-hongkong.myqcloud.com
//	cos.ap-hongkong.myqcloud.com                     → 原样返回
//
// Loki 的 S3 客户端（s3forcepathstyle=false）会自己把 bucketnames 拼到 endpoint 前，
// 所以配置里必须是区域端点，否则会拼出 <bucket>.<bucket>.cos... 触发 TLS 证书错误。
func stripBucketPrefix(endpoint, bucket string) string {
	if bucket == "" || endpoint == "" {
		return endpoint
	}
	scheme := ""
	host := strings.TrimSpace(endpoint)
	if i := strings.Index(host, "://"); i >= 0 {
		scheme, host = host[:i+3], host[i+3:]
	}
	prefix := strings.ToLower(bucket) + "."
	if strings.HasPrefix(strings.ToLower(host), prefix) {
		return scheme + host[len(prefix):]
	}
	return endpoint
}

// bucketPathFor 根据 endpoint 风格决定桶路径（用于建桶 PUT）。
//
//	1) 区域端点（推荐）:  cos.ap-hongkong.myqcloud.com       → 路径式  /<bucket>
//	2) 虚拟主机式:        mybucket-1250000000.cos.<region>…  → 路径为  /
func bucketPathFor(endpoint, bucket string) string {
	if endpointHasBucket(endpoint, bucket) {
		return "/"
	}
	return "/" + url.PathEscape(bucket)
}

// TestStorageConnection 逐级测试对象存储:可达性 → 凭证 → 桶存在性(可选自动创建)。
//
// 探测统一用 GET /，再按 endpoint 形态解释响应（不用 HEAD /<bucket>）：
//   - 区域端点      → GET / 是 ListBuckets，解析 <Buckets><Bucket><Name>
//   - 虚拟主机式端点 → GET / 是 ListObjects，返回 <ListBucketResult><Name>桶名</Name>
//
// 之所以不用 HEAD：腾讯云仅有 QcloudCOSDataFullControl（数据读写）的 AK
// 并没有 cos:HeadBucket 权限，HEAD 会 403，会被误判成"桶不存在/无权限"。
func TestStorageConnection(p StorageTestParams) StorageTestResult {
	vhost := endpointHasBucket(p.Endpoint, p.Bucket)

	resp, err := s3Request("GET", p.Endpoint, p.Region, p.AccessKey, p.SecretKey, "/", p.Insecure)
	if err != nil {
		return StorageTestResult{Message: fmt.Sprintf("无法连接对象存储:%v", err)}
	}
	defer resp.Body.Close()
	body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))

	if resp.StatusCode == http.StatusForbidden || resp.StatusCode == http.StatusUnauthorized {
		return StorageTestResult{
			Reachable: true,
			Message: fmt.Sprintf("连接成功但认证失败(HTTP %d),请检查 Access Key / Secret Key。%s",
				resp.StatusCode, s3ErrBrief(body)),
		}
	}
	if resp.StatusCode != http.StatusOK {
		return StorageTestResult{
			Reachable: true,
			Message: fmt.Sprintf("对象存储返回异常(HTTP %d):%s",
				resp.StatusCode, s3ErrBrief(body)),
		}
	}

	found := false
	if vhost {
		// ListObjects：<ListBucketResult><Name>桶名</Name>… 能返回即说明桶存在且可读
		if strings.Contains(string(body), "<ListBucketResult") {
			if m := regexp.MustCompile(`<Name>([^<]*)</Name>`).FindSubmatch(body); len(m) > 1 {
				found = strings.EqualFold(string(m[1]), p.Bucket)
			} else {
				found = true
			}
		}
	} else {
		// ListBuckets：<ListAllMyBucketsResult><Buckets><Bucket><Name>桶名</Name>
		var listResp struct {
			Buckets struct {
				Bucket []struct {
					Name string `xml:"Name"`
				} `xml:"Bucket"`
			} `xml:"Buckets"`
		}
		_ = xml.Unmarshal(body, &listResp)
		for _, b := range listResp.Buckets.Bucket {
			if b.Name == p.Bucket {
				found = true
				break
			}
		}
	}

	result := StorageTestResult{Reachable: true, Authenticated: true, BucketExists: found}

	if found {
		result.Message = "连接成功,凭证有效,桶已存在"
		return result
	}

	// ② 桶不存在:按需创建
	if !p.CreateIfMissing {
		result.Message = fmt.Sprintf("连接成功,凭证有效,但桶「%s」不存在;勾选「不存在时自动创建」可一键创建", p.Bucket)
		return result
	}

	putResp, err := s3Request("PUT", p.Endpoint, p.Region, p.AccessKey, p.SecretKey, bucketPathFor(p.Endpoint, p.Bucket), p.Insecure)
	if err != nil {
		result.Message = fmt.Sprintf("创建桶失败:%v", err)
		return result
	}
	defer putResp.Body.Close()

	// 200=创建成功;409 BucketAlreadyOwnedByYou 视为已存在
	if putResp.StatusCode == http.StatusOK || putResp.StatusCode == http.StatusConflict {
		// 二次校验：再用 GET / 确认桶真的存在（不用 HEAD —— 很多 AK 没有 cos:HeadBucket）
		verified := false
		if vresp, verr := s3Request("GET", p.Endpoint, p.Region, p.AccessKey, p.SecretKey, "/", p.Insecure); verr == nil {
			vbody, _ := io.ReadAll(io.LimitReader(vresp.Body, 1<<20))
			vresp.Body.Close()
			if vhost {
				verified = strings.Contains(string(vbody), "<ListBucketResult")
			} else {
				var vl struct {
					Buckets struct {
						Bucket []struct {
							Name string `xml:"Name"`
						} `xml:"Bucket"`
					} `xml:"Buckets"`
				}
				_ = xml.Unmarshal(vbody, &vl)
				for _, b := range vl.Buckets.Bucket {
					if b.Name == p.Bucket {
						verified = true
						break
					}
				}
			}
		}
		if verified {
			result.BucketExists = true
			result.BucketCreated = true
			result.Message = fmt.Sprintf("连接成功,桶「%s」不存在,已自动创建并校验通过 ✓", p.Bucket)
			return result
		}
		result.Message = fmt.Sprintf(
			"凭证有效,但桶「%s」创建后校验失败:重新校验仍未找到该桶。"+
				"请确认 Endpoint / Region / 桶名是否正确"+
				"(腾讯云 COS 桶名需带 APPID,例如 mybucket-1250000000)", p.Bucket)
		return result
	}

	putBody, _ := io.ReadAll(io.LimitReader(putResp.Body, 512))

	// 403：AK 通常只有数据读写权限，没有 CreateBucket 权限
	if putResp.StatusCode == http.StatusForbidden || putResp.StatusCode == http.StatusUnauthorized {
		result.Message = fmt.Sprintf(
			"凭证有效,但当前 Access Key 无建桶权限(HTTP %d)。"+
				"腾讯云 QcloudCOSDataFullControl 只授予数据读写/列出的权限，不含 CreateBucket。"+
				"请到 COS 控制台手动创建桶「%s」(区域需与 Endpoint 一致)，创建后重新测试；"+
				"或在 CAM 为该 AK 追加 cos:PutBucket 权限",
			putResp.StatusCode, p.Bucket)
		return result
	}

	result.Message = fmt.Sprintf("凭证有效,但创建桶失败(HTTP %d):%s", putResp.StatusCode, s3ErrBrief(putBody))
	return result
}