LokiStack/backend/services/storage_check.go

package services

import (
	"crypto/hmac"
	"crypto/sha256"
	"encoding/hex"
	"encoding/xml"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"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
}

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

// TestStorageConnection 逐级测试对象存储:可达性 → 凭证 → 桶存在性(可选自动创建)。
// OSS 走其 S3 兼容端点;不覆盖 OSS 原生协议的独有能力。
func TestStorageConnection(p StorageTestParams) StorageTestResult {
	// ① ListBuckets:同时验证 网络可达 + AK/SK 有效
	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()

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

	// 解析 ListBuckets,顺带确认目标桶是否在列表中
	body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	var listResp struct {
		Buckets struct {
			Bucket []struct {
				Name string `xml:"Name"`
			} `xml:"Bucket"`
		} `xml:"Buckets"`
	}
	_ = xml.Unmarshal(body, &listResp)
	found := false
	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, "/"+url.PathEscape(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 {
		result.BucketExists = true
		result.BucketCreated = true
		result.Message = fmt.Sprintf("连接成功,桶「%s」不存在,已自动创建 ✓", p.Bucket)
		return result
	}

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