276 lines
7.1 KiB
Go
276 lines
7.1 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"mime"
|
|
"mime/multipart"
|
|
"net/http"
|
|
"os"
|
|
"path/filepath"
|
|
"sort"
|
|
"strings"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
func initTrace(logDir string, cfg TraceConfig) (*traceRecorder, func()) {
|
|
if !cfg.Enabled {
|
|
return &traceRecorder{enabled: false}, func() {}
|
|
}
|
|
if logDir == "" {
|
|
logDir = defaultLogDir
|
|
}
|
|
if err := os.MkdirAll(logDir, 0755); err != nil {
|
|
logEvent("Trace", "error", "trace 目录创建失败，已关闭 trace: %v", err)
|
|
return &traceRecorder{enabled: false}, func() {}
|
|
}
|
|
|
|
path := filepath.Join(logDir, "douyu_trace_"+time.Now().In(cst).Format("20060102_150405")+".log")
|
|
file, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0600)
|
|
if err != nil {
|
|
logEvent("Trace", "error", "trace 文件创建失败，已关闭 trace: %v", err)
|
|
return &traceRecorder{enabled: false}, func() {}
|
|
}
|
|
|
|
recorder := &traceRecorder{
|
|
enabled: true,
|
|
maxBodyBytes: cfg.MaxBodyBytes,
|
|
writer: newAsyncLogWriter(file, 2048),
|
|
filePath: path,
|
|
}
|
|
if cfg.MaxBodyBytes > 0 {
|
|
logEvent("Trace", "warn", "trace 已启用: %s，单段正文最多保留 %d 字节，不做脱敏", path, cfg.MaxBodyBytes)
|
|
} else {
|
|
logEvent("Trace", "warn", "trace 已启用: %s，全量保留请求和响应，不做脱敏", path)
|
|
}
|
|
return recorder, func() {
|
|
_ = recorder.Close()
|
|
_ = file.Close()
|
|
}
|
|
}
|
|
|
|
func (t *traceRecorder) Close() error {
|
|
if t == nil || t.writer == nil {
|
|
return nil
|
|
}
|
|
return t.writer.Close()
|
|
}
|
|
|
|
func (t *traceRecorder) record(label string, req *http.Request, status int, cost time.Duration, reqBody, respBody []byte, errText string, apiError int, apiMsg string) {
|
|
if t == nil || !t.enabled || req == nil || req.URL == nil {
|
|
return
|
|
}
|
|
if label == "服务时间" && status == http.StatusOK && errText == "" && apiError == 0 {
|
|
return
|
|
}
|
|
entry := traceEntry{
|
|
Seq: atomic.AddUint64(&t.seq, 1),
|
|
Time: time.Now().In(cst).Format("2006-01-02 15:04:05.000"),
|
|
Label: label,
|
|
Method: req.Method,
|
|
URL: req.URL.Redacted(),
|
|
Path: req.URL.Path,
|
|
Status: status,
|
|
Cost: cost,
|
|
ReqHeader: sanitizeHeader(req.Header),
|
|
ReqBody: sanitizeBodyForTrace(reqBody, req.Header.Get("Content-Type"), t.maxBodyBytes),
|
|
RespBody: sanitizeResponseBody(respBody, t.maxBodyBytes),
|
|
APIError: apiError,
|
|
APIMessage: apiMsg,
|
|
ErrorText: errText,
|
|
MaxBodyBytes: t.maxBodyBytes,
|
|
}
|
|
_, _ = t.writer.Write([]byte(formatTraceEntry(entry)))
|
|
}
|
|
|
|
func traceRequestBody(t *traceRecorder, req *http.Request) []byte {
|
|
if t == nil || !t.enabled || req == nil || req.GetBody == nil {
|
|
return nil
|
|
}
|
|
body, err := req.GetBody()
|
|
if err != nil {
|
|
return nil
|
|
}
|
|
defer body.Close()
|
|
if t.maxBodyBytes > 0 {
|
|
data, _ := io.ReadAll(io.LimitReader(body, int64(t.maxBodyBytes)+1))
|
|
return data
|
|
}
|
|
data, _ := io.ReadAll(body)
|
|
return data
|
|
}
|
|
|
|
func sanitizeHeader(header http.Header) map[string][]string {
|
|
out := make(map[string][]string, len(header))
|
|
for key, values := range header {
|
|
canonical := http.CanonicalHeaderKey(key)
|
|
copied := make([]string, len(values))
|
|
copy(copied, values)
|
|
out[canonical] = copied
|
|
}
|
|
return out
|
|
}
|
|
|
|
func sanitizeBodyForTrace(body []byte, contentType string, max int) string {
|
|
if len(body) == 0 {
|
|
return ""
|
|
}
|
|
mediaType, params, _ := mime.ParseMediaType(contentType)
|
|
if strings.HasPrefix(mediaType, "multipart/") {
|
|
boundary := params["boundary"]
|
|
if boundary != "" {
|
|
return sanitizeMultipartBody(body, boundary, max)
|
|
}
|
|
}
|
|
return sanitizeTextBody(body, max)
|
|
}
|
|
|
|
func sanitizeResponseBody(body []byte, max int) string {
|
|
return sanitizeTextBody(body, max)
|
|
}
|
|
|
|
func sanitizeTextBody(body []byte, max int) string {
|
|
raw := string(body)
|
|
if pretty, ok := prettyJSON(raw); ok {
|
|
return truncateTraceText(pretty, max)
|
|
}
|
|
return truncateTraceText(raw, max)
|
|
}
|
|
|
|
func sanitizeMultipartBody(body []byte, boundary string, max int) string {
|
|
reader := multipartReader(bytes.NewReader(body), boundary)
|
|
fields := make(map[string][]string)
|
|
for {
|
|
part, err := reader.NextPart()
|
|
if err == io.EOF {
|
|
break
|
|
}
|
|
if err != nil {
|
|
return truncateTraceText(string(body), max)
|
|
}
|
|
name := part.FormName()
|
|
if name == "" {
|
|
continue
|
|
}
|
|
var valueBytes []byte
|
|
if max > 0 {
|
|
valueBytes, _ = io.ReadAll(io.LimitReader(part, int64(max)+1))
|
|
} else {
|
|
valueBytes, _ = io.ReadAll(part)
|
|
}
|
|
value := truncateTraceText(string(valueBytes), max)
|
|
fields[name] = append(fields[name], value)
|
|
}
|
|
encoded, err := json.Marshal(fields)
|
|
if err != nil {
|
|
return ""
|
|
}
|
|
return string(encoded)
|
|
}
|
|
|
|
func multipartReader(r io.Reader, boundary string) *multipart.Reader {
|
|
return multipart.NewReader(r, boundary)
|
|
}
|
|
|
|
func truncateTraceText(s string, max int) string {
|
|
if max <= 0 || len(s) <= max {
|
|
return s
|
|
}
|
|
return s[:max] + fmt.Sprintf("...<truncated %d bytes>", len(s)-max)
|
|
}
|
|
|
|
type traceEntry struct {
|
|
Seq uint64
|
|
Time string
|
|
Label string
|
|
Method string
|
|
URL string
|
|
Path string
|
|
Status int
|
|
Cost time.Duration
|
|
ReqHeader map[string][]string
|
|
ReqBody string
|
|
RespBody string
|
|
APIError int
|
|
APIMessage string
|
|
ErrorText string
|
|
MaxBodyBytes int
|
|
}
|
|
|
|
func formatTraceEntry(e traceEntry) string {
|
|
var b strings.Builder
|
|
statusText := "成功"
|
|
if e.Status < 200 || e.Status >= 300 || e.APIError != 0 || e.ErrorText != "" {
|
|
statusText = "异常"
|
|
}
|
|
|
|
b.WriteString(strings.Repeat("=", 96))
|
|
b.WriteByte('\n')
|
|
fmt.Fprintf(&b, "#%06d %s [%s] %s %s\n", e.Seq, e.Time, statusText, e.Method, e.URL)
|
|
fmt.Fprintf(&b, "接口: %s\n", emptyDefault(e.Label, "-"))
|
|
fmt.Fprintf(&b, "路径: %s\n", emptyDefault(e.Path, "-"))
|
|
fmt.Fprintf(&b, "结果: HTTP %d | 耗时 %s | apiError=%d", e.Status, roundDuration(e.Cost), e.APIError)
|
|
if e.APIMessage != "" {
|
|
fmt.Fprintf(&b, " | msg=%s", e.APIMessage)
|
|
}
|
|
b.WriteByte('\n')
|
|
if e.ErrorText != "" {
|
|
fmt.Fprintf(&b, "错误: %s\n", e.ErrorText)
|
|
}
|
|
b.WriteString("\n[请求头]\n")
|
|
b.WriteString(formatTraceHeaders(e.ReqHeader))
|
|
b.WriteString("\n[请求体]\n")
|
|
b.WriteString(formatTraceBody(e.ReqBody))
|
|
b.WriteString("\n[响应体]\n")
|
|
b.WriteString(formatTraceBody(e.RespBody))
|
|
if e.MaxBodyBytes > 0 {
|
|
fmt.Fprintf(&b, "\n提示: 单段正文最多保留 %d 字节；trace 不做脱敏。\n", e.MaxBodyBytes)
|
|
} else {
|
|
b.WriteString("\n提示: trace 全量保留请求和响应，不做脱敏。\n")
|
|
}
|
|
b.WriteByte('\n')
|
|
return b.String()
|
|
}
|
|
|
|
func formatTraceHeaders(headers map[string][]string) string {
|
|
if len(headers) == 0 {
|
|
return "<empty>\n"
|
|
}
|
|
keys := make([]string, 0, len(headers))
|
|
for key := range headers {
|
|
keys = append(keys, key)
|
|
}
|
|
sort.Strings(keys)
|
|
var b strings.Builder
|
|
for _, key := range keys {
|
|
fmt.Fprintf(&b, "%s: %s\n", key, strings.Join(headers[key], ", "))
|
|
}
|
|
return b.String()
|
|
}
|
|
|
|
func formatTraceBody(body string) string {
|
|
body = strings.TrimSpace(body)
|
|
if body == "" {
|
|
return "<empty>\n"
|
|
}
|
|
if pretty, ok := prettyJSON(body); ok {
|
|
return pretty + "\n"
|
|
}
|
|
return body + "\n"
|
|
}
|
|
|
|
func prettyJSON(s string) (string, bool) {
|
|
var value any
|
|
if err := json.Unmarshal([]byte(s), &value); err != nil {
|
|
return "", false
|
|
}
|
|
out, err := json.MarshalIndent(value, "", " ")
|
|
if err != nil {
|
|
return "", false
|
|
}
|
|
return string(out), true
|
|
}
|