LokiStack/backend/services/log_query.go

package services

// 日志查询代理 — 后端代表前端访问 Loki HTTP API。
// cluster-svc 型经 kube-apiserver service proxy(不依赖 ClusterIP 直通与集群 DNS);
// url 型直连。契约参照 aiops 日志中心(query/labels/values/histogram + cursor 翻页)。
import (
	"context"
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os/exec"
	"regexp"
	"sort"
	"strconv"
	"strings"
	"time"

	"opsxloki/models"
	"opsxloki/pkg/crypto"
	"opsxloki/repository"
)

const maxLogQueryLimit = 1000

// ==================== 数据源 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
}

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-binary→url;k8s→cluster-svc)
func EnsureLogSourceForDeployment(d *models.LokiDeployment) {
	if d == nil || d.Endpoint == "" {
		return
	}
	db := repository.GetDB()
	db.RLock()
	for i := range db.LogSources {
		if db.LogSources[i].Name == d.Name {
			db.RUnlock()
			return
		}
	}
	db.RUnlock()

	svc := "loki"
	port := 3100
	src := &models.LogSource{
		Name: d.Name, TargetType: "cluster-svc",
		ClusterID: d.ClusterID, Namespace: d.Namespace, Service: svc, Port: port,
		AuthType: "none",
	}
	if d.Mode == "host-binary" {
		src.TargetType = "url"
		src.URL = d.Endpoint
		src.ClusterID = nil
	}
	_, _ = CreateLogSource(src)
}

// ==================== Loki HTTP 访问 ====================

// lokiProxyGET 对数据源执行 GET;cluster-svc 走 kubectl --raw service proxy。
func lokiProxyGET(ctx context.Context, src *models.LogSource, path string, query url.Values) ([]byte, int, error) {
	// cluster-svc:kubectl get --raw /api/v1/namespaces/{ns}/services/{svc}:{port}/proxy{path}?q
	if src.TargetType == "cluster-svc" {
		db := repository.GetDB()
		db.RLock()
		var cluster *models.Cluster
		for i := range db.Clusters {
			if src.ClusterID != nil && db.Clusters[i].ID == *src.ClusterID {
				c := db.Clusters[i]
				cluster = &c
				break
			}
		}
		db.RUnlock()
		if cluster == nil {
			return nil, 500, fmt.Errorf("数据源关联的集群不存在")
		}
		env, cleanup, err := kubeconfigEnv(cluster)
		if err != nil {
			return nil, 500, err
		}
		defer cleanup()

		proxyPath := fmt.Sprintf("/api/v1/namespaces/%s/services/%s:%d/proxy%s",
			src.Namespace, src.Service, src.Port, path)
		full := proxyPath
		if qs := query.Encode(); qs != "" {
			full += "?" + qs
		}
		cmd := exec.CommandContext(ctx, "kubectl", "get", "--raw", full)
		cmd.Env = env
		out, err := cmd.CombinedOutput()
		if err != nil {
			return nil, 502, fmt.Errorf("kubectl 代理查询失败: %v\n%s", err, truncateStr(string(out), 300))
		}
		return out, 200, nil
	}

	// url 直连
	if src.URL == "" {
		return nil, 500, fmt.Errorf("数据源未配置访问地址")
	}
	full := strings.TrimSuffix(src.URL, "/") + path
	if qs := query.Encode(); qs != "" {
		full += "?" + qs
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, full, nil)
	if err != nil {
		return nil, 500, err
	}
	applyLokiAuth(req, src)
	client := &http.Client{Timeout: 60 * time.Second, Transport: &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: src.SkipTLSVerify},
	}}
	resp, err := client.Do(req)
	if err != nil {
		return nil, 502, err
	}
	defer resp.Body.Close()
	body, err := io.ReadAll(io.LimitReader(resp.Body, 32<<20))
	if err != nil {
		return nil, 502, err
	}
	if resp.StatusCode >= 400 {
		return body, resp.StatusCode, fmt.Errorf("Loki 返回 HTTP %d: %s", resp.StatusCode, truncateStr(string(body), 300))
	}
	return body, resp.StatusCode, nil
}

func applyLokiAuth(req *http.Request, src *models.LogSource) {
	switch src.AuthType {
	case "basic":
		if pass, err := DecryptCred(src.BasicPass); err == nil {
			req.Header.Set("Authorization", "Basic "+
				base64.StdEncoding.EncodeToString([]byte(src.BasicUser+":"+pass)))
		}
	case "bearer":
		if token, err := DecryptCred(src.BearerToken); err == nil {
			req.Header.Set("Authorization", "Bearer "+token)
		}
	}
	for name, encrypted := range src.Headers {
		if validSourceHeader(name) {
			if value, err := DecryptCred(encrypted); err == nil {
				req.Header.Set(name, value)
			}
		}
	}
}

func truncateStr(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "..."
}

// ==================== 查询实现(契约对齐 aiops) ====================

type LogEntryRow struct {
	Timestamp   string            `json:"timestamp"`
	TimestampNs string            `json:"timestampNs"`
	Level       string            `json:"level"`
	Message     string            `json:"message"`
	Labels      map[string]string `json:"labels"`
}

type LogQueryResult struct {
	Entries    []LogEntryRow `json:"entries"`
	HasMore    bool          `json:"hasMore"`
	NextCursor string        `json:"nextCursor"`
}

// QueryLokiLogs 日志查询:flatten streams → 统一行;cursor(最后一条纳秒时间戳)向后翻页
func QueryLokiLogs(ctx context.Context, src *models.LogSource, query string, start, end int64, limit int, direction, cursor string) (*LogQueryResult, error) {
	if limit <= 0 {
		limit = 200
	}
	if limit > maxLogQueryLimit {
		limit = maxLogQueryLimit
	}
	if direction == "" {
		direction = "backward"
	}
	if end <= 0 {
		end = time.Now().UnixNano()
	}
	if start <= 0 {
		start = end - int64(time.Hour) // 默认最近 1 小时
	}
	if cursor != "" {
		if c, err := strconv.ParseInt(cursor, 10, 64); err == nil {
			if direction == "forward" {
				start = c + 1
			} else {
				end = c - 1
			}
		}
	}

	q := url.Values{}
	q.Set("query", query)
	q.Set("start", strconv.FormatInt(start, 10))
	q.Set("end", strconv.FormatInt(end, 10))
	q.Set("limit", strconv.Itoa(limit))
	q.Set("direction", direction)

	body, _, err := lokiProxyGET(ctx, src, "/loki/api/v1/query_range", q)
	if err != nil {
		return nil, err
	}
	var resp struct {
		Status string `json:"status"`
		Data   struct {
			ResultType string `json:"resultType"`
			Result     []struct {
				Stream map[string]string `json:"stream"`
				Values [][2]string       `json:"values"`
			} `json:"result"`
		} `json:"data"`
	}
	if err := json.Unmarshal(body, &resp); err != nil {
		return nil, fmt.Errorf("解析 Loki 响应失败: %w", err)
	}

	entries := make([]LogEntryRow, 0, limit)
	var lastTs string
	for _, stream := range resp.Data.Result {
		for _, v := range stream.Values {
			entries = append(entries, LogEntryRow{
				Timestamp:   formatLokiTimestamp(v[0]),
				TimestampNs: v[0],
				Level:       detectLevel(stream.Stream, v[1]),
				Message:     v[1],
				Labels:      stream.Stream,
			})
		}
	}
	sort.SliceStable(entries, func(i, j int) bool {
		a, _ := strconv.ParseInt(entries[i].TimestampNs, 10, 64)
		b, _ := strconv.ParseInt(entries[j].TimestampNs, 10, 64)
		if direction == "forward" {
			return a < b
		}
		return a > b
	})
	if len(entries) > limit {
		entries = entries[:limit]
	}
	if len(entries) > 0 {
		lastTs = entries[len(entries)-1].TimestampNs
	}

	nextCursor := ""
	if len(entries) >= limit && lastTs != "" {
		if ts, err := strconv.ParseInt(strings.Split(lastTs, ".")[0], 10, 64); err == nil {
			nextCursor = strconv.FormatInt(ts, 10)
		}
	}
	return &LogQueryResult{Entries: entries, HasMore: nextCursor != "", NextCursor: nextCursor}, nil
}

// formatLokiTimestamp Loki 时间串(纳秒字符串/科学计数法)→ RFC3339
func formatLokiTimestamp(ts string) string {
	if i := strings.IndexAny(ts, ".eE"); i > 0 {
		ts = ts[:i]
	}
	n, err := strconv.ParseInt(ts, 10, 64)
	if err != nil {
		return ts
	}
	// 自适应纳秒/微秒/毫秒/秒
	switch {
	case n > 1e18:
		return time.Unix(0, n).Format(time.RFC3339)
	case n > 1e15:
		return time.Unix(0, n*1e3).Format(time.RFC3339)
	case n > 1e12:
		return time.Unix(0, n*1e6).Format(time.RFC3339)
	case n > 1e9:
		return time.Unix(0, n*1e9).Format(time.RFC3339)
	default:
		return time.Unix(n, 0).Format(time.RFC3339)
	}
}

// detectLevel 优先取 level 标签;否则按行内容推断
func detectLevel(labels map[string]string, line string) string {
	if lv, ok := labels["level"]; ok && lv != "" {
		return strings.ToLower(lv)
	}
	l := strings.ToLower(line)
	for _, p := range []string{"error", "err", "fatal", "panic"} {
		if strings.Contains(l, p) {
			return "error"
		}
	}
	if strings.Contains(l, "warn") {
		return "warn"
	}
	return "info"
}

// LokiLabels 标签名列表
func LokiLabels(ctx context.Context, src *models.LogSource, start, end int64) ([]string, error) {
	q := url.Values{}
	if end > 0 {
		q.Set("end", strconv.FormatInt(end, 10))
	}
	if start > 0 {
		q.Set("start", strconv.FormatInt(start, 10))
	}
	body, _, err := lokiProxyGET(ctx, src, "/loki/api/v1/labels", q)
	if err != nil {
		return nil, err
	}
	var resp struct {
		Data []string `json:"data"`
	}
	if err := json.Unmarshal(body, &resp); err != nil {
		return nil, err
	}
	return resp.Data, nil
}

// LokiLabelValues 标签值列表(query 只取 {...} selector 部分)
func LokiLabelValues(ctx context.Context, src *models.LogSource, name, selector string, start, end int64) ([]string, error) {
	q := url.Values{}
	if s := extractStreamSelector(selector); s != "" && s != "{}" {
		q.Set("query", s)
	}
	if end > 0 {
		q.Set("end", strconv.FormatInt(end, 10))
	}
	if start > 0 {
		q.Set("start", strconv.FormatInt(start, 10))
	}
	body, _, err := lokiProxyGET(ctx, src, "/loki/api/v1/label/"+url.PathEscape(name)+"/values", q)
	if err != nil {
		return nil, err
	}
	var resp struct {
		Data []string `json:"data"`
	}
	if err := json.Unmarshal(body, &resp); err != nil {
		return nil, err
	}
	return resp.Data, nil
}

// extractStreamSelector 从 LogQL 中提取 {...} 部分(labels API 仅接受 stream selector)
func extractStreamSelector(query string) string {
	i := strings.Index(query, "{")
	if i < 0 {
		return ""
	}
	j := strings.Index(query[i:], "}")
	if j < 0 {
		return ""
	}
	return query[i : i+j+1]
}

// LokiHistogram 日志量直方图:服务端包 sum(count_over_time(q[step]))
func LokiHistogram(ctx context.Context, src *models.LogSource, query string, start, end int64) ([]map[string]interface{}, error) {
	if end <= 0 {
		end = time.Now().UnixNano()
	}
	if start <= 0 {
		start = end - int64(time.Hour)
	}
	span := end - start
	step := pickHistogramStep(span)
	wrapped := fmt.Sprintf("sum(count_over_time(%s[%s]))", query, step)

	q := url.Values{}
	q.Set("query", wrapped)
	q.Set("start", strconv.FormatInt(start, 10))
	q.Set("end", strconv.FormatInt(end, 10))
	q.Set("step", step)
	body, _, err := lokiProxyGET(ctx, src, "/loki/api/v1/query_range", q)
	if err != nil {
		return nil, err
	}
	var resp struct {
		Data struct {
			Result []struct {
				Values [][2]interface{} `json:"values"`
			} `json:"result"`
		} `json:"data"`
	}
	if err := json.Unmarshal(body, &resp); err != nil {
		return nil, err
	}
	buckets := []map[string]interface{}{}
	if len(resp.Data.Result) > 0 {
		for _, v := range resp.Data.Result[0].Values {
			// metric API 的 ts 为秒级浮点数(如 1.79e+09),query 日志的才是纳秒字符串
			tsf, _ := strconv.ParseFloat(fmt.Sprintf("%v", v[0]), 64)
			cnt, _ := strconv.ParseFloat(fmt.Sprintf("%v", v[1]), 64)
			buckets = append(buckets, map[string]interface{}{
				"timestamp": time.Unix(int64(tsf), 0).Format(time.RFC3339),
				"count":     int64(cnt),
			})
		}
	}
	return buckets, nil
}

// pickHistogramStep 按时间跨度自适应步长(aiops 同款)
func pickHistogramStep(spanNano int64) string {
	span := time.Duration(spanNano)
	switch {
	case span <= 15*time.Minute:
		return "30s"
	case span <= time.Hour:
		return "1m"
	case span <= 6*time.Hour:
		return "5m"
	case span <= 24*time.Hour:
		return "15m"
	case span <= 7*24*time.Hour:
		return "1h"
	default:
		return "6h"
	}
}

// DeploymentReady 部署级就绪探测:按部署名匹配自动注册的数据源,经其访问 Loki /ready
func DeploymentReady(ctx context.Context, deploymentName string) (bool, string) {
	if deploymentName == "" {
		return false, ""
	}
	db := repository.GetDB()
	db.RLock()
	var src *models.LogSource
	for i := range db.LogSources {
		if db.LogSources[i].Name == deploymentName {
			s := db.LogSources[i]
			src = &s
			break
		}
	}
	db.RUnlock()
	if src == nil {
		return false, "未注册数据源"
	}
	if LokiReady(ctx, src) {
		return true, ""
	}
	return false, "Loki 未就绪"
}

// LokiReady Loki /ready 探活(部署卡片/查询页状态点)
func LokiReady(ctx context.Context, src *models.LogSource) bool {
	body, code, err := lokiProxyGET(ctx, src, "/ready", nil)
	if err != nil || code != 200 {
		return false
	}
	return strings.Contains(strings.ToLower(string(body)), "ready")
}

// TestLogSource 数据源连通性(/loki/api/v1/labels 探活)
func TestLogSource(ctx context.Context, src *models.LogSource) map[string]interface{} {
	labels, err := LokiLabels(ctx, src, 0, 0)
	if err != nil {
		return map[string]interface{}{"success": false, "message": err.Error()}
	}
	return map[string]interface{}{
		"success": true,
		"message": fmt.Sprintf("连接成功,发现 %d 个标签", len(labels)),
		"labels":  labels,
	}
}