Logging-System/server/server.go

// Package server 是日志采集服务：批量写入、条件查询、SSE 尾随。默认监听 :8092。
package server

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"strings"

	"git-ssh.u1680832.nyat.app/weichen/Logging-System.git/logsys"
	"git-ssh.u1680832.nyat.app/weichen/Logging-System.git/redact"
)

// DefaultAddr 是 logcollector 的默认监听地址。
const DefaultAddr = ":8092"

// TraceHeader 与错误码系统使用同一个追踪头。
const TraceHeader = "X-Trace-Id"

// Server 提供日志采集 HTTP API。
type Server struct {
	Store *Store
	Token string
}

// New 创建服务。redact 为 nil 时仍做内置脱敏。
func New(path, token string, limit int) (*Server, error) {
	rd, err := redact.New(nil)
	if err != nil {
		return nil, err
	}
	st, err := OpenStore(path, limit, rd)
	if err != nil {
		return nil, err
	}
	return &Server{Store: st, Token: token}, nil
}

// Handler 返回路由。
func (s *Server) Handler() http.Handler {
	mux := http.NewServeMux()
	mux.HandleFunc("GET /v1/health", s.handleHealth)
	mux.HandleFunc("POST /v1/logs", s.handleIngest)
	mux.HandleFunc("GET /v1/logs", s.handleQuery)
	mux.HandleFunc("GET /v1/logs/tail", s.handleTail)
	return mux
}

// ListenAndServe 在 addr 上服务。
func (s *Server) ListenAndServe(addr string) error {
	if strings.TrimSpace(addr) == "" {
		addr = DefaultAddr
	}
	return http.ListenAndServe(addr, s.Handler())
}

// Close 关闭存储。
func (s *Server) Close() error {
	if s.Store == nil {
		return nil
	}
	return s.Store.Close()
}

func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) {
	writeJSON(w, http.StatusOK, map[string]any{"ok": true})
}

func (s *Server) handleIngest(w http.ResponseWriter, r *http.Request) {
	if !s.authorized(r) {
		writeJSON(w, http.StatusUnauthorized, map[string]string{"error": "unauthorized"})
		return
	}
	body, err := io.ReadAll(http.MaxBytesReader(w, r.Body, 8<<20))
	if err != nil {
		writeJSON(w, http.StatusBadRequest, map[string]string{"error": "body too large or unreadable"})
		return
	}
	var req struct {
		Entries []logsys.Entry `json:"entries"`
	}
	if err := json.Unmarshal(body, &req); err != nil {
		writeJSON(w, http.StatusBadRequest, map[string]string{"error": "invalid json"})
		return
	}
	traceID := sanitizeTrace(r.Header.Get(TraceHeader))
	for i := range req.Entries {
		if req.Entries[i].TraceID == "" && traceID != "" {
			req.Entries[i].TraceID = traceID
		}
	}
	recs := s.Store.Append(req.Entries)
	writeJSON(w, http.StatusOK, map[string]any{"accepted": len(recs)})
}

func (s *Server) handleQuery(w http.ResponseWriter, r *http.Request) {
	if !s.authorized(r) {
		writeJSON(w, http.StatusUnauthorized, map[string]string{"error": "unauthorized"})
		return
	}
	q := r.URL.Query()
	cursor, _ := strconv.ParseInt(q.Get("cursor"), 10, 64)
	limit, _ := strconv.Atoi(q.Get("limit"))
	recs, next := s.Store.Query(Query{
		TraceID: q.Get("trace_id"),
		Level:   q.Get("level"),
		Service: q.Get("service"),
		Text:    q.Get("q"),
		Cursor:  cursor,
		Limit:   limit,
	})
	if recs == nil {
		recs = []Record{}
	}
	writeJSON(w, http.StatusOK, map[string]any{"entries": recs, "next_cursor": next})
}

func (s *Server) handleTail(w http.ResponseWriter, r *http.Request) {
	if !s.authorized(r) {
		writeJSON(w, http.StatusUnauthorized, map[string]string{"error": "unauthorized"})
		return
	}
	flusher, ok := w.(http.Flusher)
	if !ok {
		writeJSON(w, http.StatusInternalServerError, map[string]string{"error": "streaming unsupported"})
		return
	}
	cursor, _ := strconv.ParseInt(r.URL.Query().Get("cursor"), 10, 64)
	ch, cancel := s.Store.Subscribe(cursor)
	defer cancel()
	w.Header().Set("Content-Type", "text/event-stream")
	w.Header().Set("Cache-Control", "no-cache")
	w.WriteHeader(http.StatusOK)
	for {
		select {
		case <-r.Context().Done():
			return
		case rec, ok := <-ch:
			if !ok {
				return
			}
			data, err := json.Marshal(rec)
			if err != nil {
				continue
			}
			fmt.Fprintf(w, "data: %s\n\n", data)
			flusher.Flush()
		}
	}
}

func (s *Server) authorized(r *http.Request) bool {
	if s.Token == "" {
		return true
	}
	if r.Header.Get("X-Log-Token") == s.Token {
		return true
	}
	auth := r.Header.Get("Authorization")
	return auth == "Bearer "+s.Token
}

func writeJSON(w http.ResponseWriter, code int, v any) {
	w.Header().Set("Content-Type", "application/json; charset=utf-8")
	w.WriteHeader(code)
	enc := json.NewEncoder(w)
	enc.SetEscapeHTML(false)
	_ = enc.Encode(v)
}

func sanitizeTrace(s string) string {
	s = strings.TrimSpace(s)
	if len(s) < 8 || len(s) > 128 {
		return ""
	}
	for _, r := range s {
		switch {
		case r >= 'a' && r <= 'z', r >= 'A' && r <= 'Z', r >= '0' && r <= '9', r == '-', r == '_', r == '.', r == ':':
			continue
		default:
			return ""
		}
	}
	return s
}