LokiStack/backend/services/log_source.go

// 日志数据源管理 — CRUD、默认源、部署自动注册(Managed 标记)。
package services

import (
	"encoding/json"
	"fmt"
	"regexp"
	"sort"
	"strings"
	"time"

	"lokistack/models"
	"lokistack/pkg/crypto"
	"lokistack/repository"
)

// ==================== 数据源 CRUD ====================

func GetLogSource(id uint) *models.LogSource {
	db := repository.GetDB()
	db.RLock()
	defer db.RUnlock()
	for i := range db.LogSources {
		if db.LogSources[i].ID == id {
			s := db.LogSources[i]
			return &s
		}
	}
	return nil
}

func ListLogSources() []map[string]interface{} {
	db := repository.GetDB()
	db.RLock()
	defer db.RUnlock()
	out := make([]map[string]interface{}, 0, len(db.LogSources))
	for i := range db.LogSources {
		out = append(out, db.LogSources[i].Sanitized())
	}
	return out
}

func ListVisibleLogSources(userID uint, admin bool) []map[string]interface{} {
	db := repository.GetDB()
	db.RLock()
	defer db.RUnlock()
	out := make([]map[string]interface{}, 0)
	for i := range db.LogSources {
		s := &db.LogSources[i]
		if admin || policyGroupAccess(db.LogPolicyGroups, s.ID, userID) != nil || sourceAccess(s, userID) != nil {
			item := s.Sanitized()
			if !admin {
				item = map[string]interface{}{"id": s.ID, "name": s.Name, "isDefault": s.IsDefault}
			}
			out = append(out, item)
		}
	}
	return out
}

// UserAccessLabels 返回某用户在指定数据源上的授权 labels。
//
// 返回值语义：
//   - nil      → 该用户对此数据源完全无权限
//   - 空 map   → 有权限但不限定范围（可查全部）
//   - 非空 map → 只能查匹配这些 label 的日志（如 {"namespace":"astralx"}）
func UserAccessLabels(src *models.LogSource, userID uint) map[string]string {
	db := repository.GetDB()
	var access *models.LogSourceAccess
	if db != nil {
		db.RLock()
		access = policyGroupAccess(db.LogPolicyGroups, src.ID, userID)
		db.RUnlock()
	}
	if access == nil {
		access = sourceAccess(src, userID)
	}
	if access == nil {
		return nil
	}
	if access.Labels == nil {
		return map[string]string{}
	}
	return access.Labels
}

func sourceAccess(src *models.LogSource, userID uint) *models.LogSourceAccess {
	for i := range src.Access {
		if src.Access[i].UserID == userID {
			return &src.Access[i]
		}
	}
	return nil
}

var labelNamePattern = regexp.MustCompile(`^[a-zA-Z_][a-zA-Z0-9_]*$`)

func ValidLogLabelName(name string) bool { return labelNamePattern.MatchString(name) }

// AuthorizedLogQuery constrains the Loki stream selector server-side. Non-admins must be explicitly granted access.
func AuthorizedLogQuery(src *models.LogSource, userID uint, admin bool, query string) (string, error) {
	if admin {
		return query, nil
	}
	db := repository.GetDB()
	var access *models.LogSourceAccess
	if db != nil {
		db.RLock()
		access = policyGroupAccess(db.LogPolicyGroups, src.ID, userID)
		db.RUnlock()
	}
	if access == nil {
		access = sourceAccess(src, userID)
	}
	if access == nil {
		return "", fmt.Errorf("无权查询此数据源")
	}
	q := strings.TrimSpace(query)
	if !strings.HasPrefix(q, "{") {
		return "", fmt.Errorf("普通用户查询必须以日志流选择器开头")
	}
	close := strings.IndexByte(q, '}')
	if close < 0 || strings.Contains(q[close+1:], "{") {
		return "", fmt.Errorf("查询只能包含一个日志流选择器")
	}
	if len(access.Labels) == 0 {
		return q, nil
	}
	keys := make([]string, 0, len(access.Labels))
	for k := range access.Labels {
		if !labelNamePattern.MatchString(k) {
			return "", fmt.Errorf("访问策略包含无效标签名")
		}
		keys = append(keys, k)
	}
	sort.Strings(keys)
	matchers := make([]string, 0, len(keys))
	for _, k := range keys {
		v, _ := json.Marshal(access.Labels[k])
		matchers = append(matchers, k+"="+string(v))
	}
	inner := strings.TrimSpace(q[1:close])
	if inner != "" {
		inner += ","
	}
	return "{" + inner + strings.Join(matchers, ",") + "}" + q[close+1:], nil
}

func CreateLogSource(s *models.LogSource) (*models.LogSource, error) {
	if err := encryptSourceCredentials(s); err != nil {
		return nil, err
	}
	db := repository.GetDB()
	newID := db.NextID("logSources")
	db.Lock()
	defer db.Unlock()
	s.ID = newID
	if s.IsDefault {
		for i := range db.LogSources {
			db.LogSources[i].IsDefault = false
		}
	}
	// 首个数据源自动设默认
	if len(db.LogSources) == 0 {
		s.IsDefault = true
	}
	db.LogSources = append(db.LogSources, *s)
	db.RequestSave()
	return s, nil
}

func encryptSourceCredentials(s *models.LogSource) error {
	if s.BasicPass != "" {
		v, err := crypto.EncryptField(s.BasicPass)
		if err != nil {
			return err
		}
		s.BasicPass = v
	}
	if s.BearerToken != "" {
		v, err := crypto.EncryptField(s.BearerToken)
		if err != nil {
			return err
		}
		s.BearerToken = v
	}
	for name, value := range s.Headers {
		if !validSourceHeader(name) {
			return fmt.Errorf("自定义 Header 名称无效: %s", name)
		}
		v, err := crypto.EncryptField(value)
		if err != nil {
			return err
		}
		s.Headers[name] = v
	}
	return nil
}

func validSourceHeader(name string) bool {
	if name == "" || strings.EqualFold(name, "Host") || strings.EqualFold(name, "Authorization") {
		return false
	}
	for _, c := range name {
		if !(c >= 'a' && c <= 'z' || c >= 'A' && c <= 'Z' || c >= '0' && c <= '9' || c == '-') {
			return false
		}
	}
	return true
}

func UpdateLogSource(id uint, patch map[string]interface{}) error {
	db := repository.GetDB()
	db.Lock()
	defer db.Unlock()
	for i := range db.LogSources {
		if db.LogSources[i].ID != id {
			continue
		}
		if v, ok := patch["name"].(string); ok && v != "" {
			db.LogSources[i].Name = v
		}
		if v, ok := patch["namespace"].(string); ok {
			db.LogSources[i].Namespace = v
		}
		if v, ok := patch["service"].(string); ok {
			db.LogSources[i].Service = v
		}
		if v, ok := patch["port"].(float64); ok {
			db.LogSources[i].Port = int(v)
		}
		if v, ok := patch["url"].(string); ok {
			db.LogSources[i].URL = v
		}
		if v, ok := patch["authType"].(string); ok {
			db.LogSources[i].AuthType = v
		}
		if v, ok := patch["basicUser"].(string); ok {
			db.LogSources[i].BasicUser = v
		}
		if v, ok := patch["bearerToken"].(string); ok {
			if v != "" {
				enc, err := crypto.EncryptField(v)
				if err != nil {
					return err
				}
				db.LogSources[i].BearerToken = enc
			}
		}
		if v, ok := patch["skipTLSVerify"].(bool); ok {
			db.LogSources[i].SkipTLSVerify = v
		}
		if raw, ok := patch["headers"]; ok {
			data, err := json.Marshal(raw)
			if err != nil {
				return err
			}
			var headers map[string]string
			if err := json.Unmarshal(data, &headers); err != nil {
				return err
			}
			for name, value := range headers {
				if !validSourceHeader(name) {
					return fmt.Errorf("自定义 Header 名称无效: %s", name)
				}
				if value == "" {
					if old, ok := db.LogSources[i].Headers[name]; ok {
						headers[name] = old
					}
					continue
				}
				enc, err := crypto.EncryptField(value)
				if err != nil {
					return err
				}
				headers[name] = enc
			}
			db.LogSources[i].Headers = headers
		}
		if v, ok := patch["basicPass"].(string); ok && v != "" {
			if enc, err := crypto.EncryptField(v); err == nil {
				db.LogSources[i].BasicPass = enc
			}
		}
		if v, ok := patch["isDefault"].(bool); ok && v {
			for j := range db.LogSources {
				db.LogSources[j].IsDefault = db.LogSources[j].ID == id
			}
		}
		if raw, ok := patch["access"]; ok {
			data, err := json.Marshal(raw)
			if err != nil {
				return err
			}
			var access []models.LogSourceAccess
			if err := json.Unmarshal(data, &access); err != nil {
				return err
			}
			if err := ValidateSourceAccess(access); err != nil {
				return err
			}
			db.LogSources[i].Access = access
		}
		db.LogSources[i].UpdatedAt = time.Now()
		db.RequestSave()
		return nil
	}
	return fmt.Errorf("数据源不存在")
}

func ValidateSourceAccess(access []models.LogSourceAccess) error {
	seen := map[uint]bool{}
	for _, rule := range access {
		if rule.UserID == 0 || seen[rule.UserID] {
			return fmt.Errorf("访问策略的用户 ID 无效或重复")
		}
		seen[rule.UserID] = true
		for k := range rule.Labels {
			if !labelNamePattern.MatchString(k) {
				return fmt.Errorf("访问策略包含无效标签名: %s", k)
			}
		}
	}
	return nil
}

func DeleteLogSource(id uint) error {
	db := repository.GetDB()
	db.Lock()
	defer db.Unlock()
	for i := range db.LogSources {
		if db.LogSources[i].ID == id {
			db.LogSources = append(db.LogSources[:i], db.LogSources[i+1:]...)
			groups := db.LogPolicyGroups[:0]
			for _, group := range db.LogPolicyGroups {
				if group.SourceID != id {
					groups = append(groups, group)
				}
			}
			db.LogPolicyGroups = groups
			db.RequestSave()
			return nil
		}
	}
	return fmt.Errorf("数据源不存在")
}

// EnsureLogSourceForDeployment 部署成功后自动注册查询数据源(host→url;k8s→cluster-svc)。
// 已存在同名数据源时：平台托管的自动刷新连接信息，用户手填的不覆盖。
func EnsureLogSourceForDeployment(d *models.LokiDeployment) {
	if d == nil || d.Endpoint == "" {
		return
	}
	svc := "loki"
	port := 3100
	target := &models.LogSource{
		Name: d.Name, TargetType: "cluster-svc",
		ClusterID: d.ClusterID, Namespace: d.Namespace, Service: svc, Port: port,
		AuthType: "none", Managed: true,
	}
	if d.Mode == "host-binary" || d.Mode == "host-docker" {
		target.TargetType = "url"
		target.URL = d.Endpoint
		target.ClusterID = nil
	}

	db := repository.GetDB()
	db.Lock()
	for i := range db.LogSources {
		if db.LogSources[i].Name != target.Name {
			continue
		}
		if !db.LogSources[i].Managed {
			db.Unlock()
			return
		}
		db.LogSources[i].TargetType = target.TargetType
		db.LogSources[i].ClusterID = target.ClusterID
		db.LogSources[i].Namespace = target.Namespace
		db.LogSources[i].Service = target.Service
		db.LogSources[i].Port = target.Port
		db.LogSources[i].URL = target.URL
		db.LogSources[i].UpdatedAt = time.Now()
		db.Unlock()
		db.RequestSave()
		return
	}
	db.Unlock()
	_, _ = CreateLogSource(target)
}

// EnsureLogSourcesForRunningDeployments 补齐运行中部署的数据源，避免查询页默认没有数据源。
func EnsureLogSourcesForRunningDeployments() {
	db := repository.GetDB()
	db.RLock()
	running := make([]models.LokiDeployment, 0, len(db.Deployments))
	for i := range db.Deployments {
		d := db.Deployments[i]
		if d.Status == "running" && d.Endpoint != "" {
			running = append(running, d)
		}
	}
	db.RUnlock()
	for i := range running {
		EnsureLogSourceForDeployment(&running[i])
	}
}