pkg/rest/transport/curl.go

/*
 * TencentBlueKing is pleased to support the open source community by making
 * 蓝鲸智云 - 配置平台 (BlueKing - CMDB) available.
 * Copyright (C) 2025 Tencent. All rights reserved.
 * Licensed under the MIT License (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at http://opensource.org/licenses/MIT
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on
 * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
 * either express or implied. See the License for the
 * specific language governing permissions and limitations under the License.
 * We undertake not to change the open source license (MIT license) applicable
 * to the current version of the project delivered to anyone in the future.
 */

package transport

import (
	"bytes"
	"fmt"
	"io"
	"log/slog"
	"net/http"
	"strings"
	"time"
)

// CurlLog make a http curl log transport
func CurlLog(maskKeys ...string) func(http.RoundTripper) http.RoundTripper {
	keys := make(map[string]struct{}, len(maskKeys))
	for _, key := range maskKeys {
		keys[key] = struct{}{}
	}

	return func(base http.RoundTripper) http.RoundTripper {
		t := &curlLogTransport{baseTransport: base, maskKeys: keys}
		return t
	}
}

// reqToCurl curl 格式的请求日志
func (c *curlLogTransport) reqToCurl(r *http.Request) (string, error) {
	headers := strings.Builder{}

	// 过滤掉敏感信息, header 和 query
	for key, values := range r.Header {
		for _, value := range values {
			if _, ok := c.maskKeys[key]; ok {
				value = "***"
			}
			headers.WriteString(fmt.Sprintf(" -H %q", fmt.Sprintf("%s: %s", key, value)))
		}
	}

	rawURL := *r.URL
	queryValue := rawURL.Query()
	for key := range queryValue {
		if _, ok := c.maskKeys[key]; ok {
			queryValue.Set(key, "<masked>")
		}
	}
	rawURL.RawQuery = queryValue.Encode()

	reqMsg := fmt.Sprintf("curl -X %s '%s'%s", r.Method, rawURL.String(), headers.String())
	if r.Body != nil {
		bodyBytes, err := io.ReadAll(r.Body)
		if err != nil {
			return "", err
		}
		r.Body.Close()

		r.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
		if len(bodyBytes) > 1024 {
			reqMsg += fmt.Sprintf(" -d '%s...(total %dB)'", bodyBytes[:1024], len(bodyBytes))
		} else {
			reqMsg += fmt.Sprintf(" -d '%s'", bodyBytes)
		}

	}

	return reqMsg, nil
}

// respToCurl 返回日志
func (c *curlLogTransport) respToCurl(resp *http.Response, st time.Time) (string, error) {
	var (
		bodyBytes []byte
		err       error
	)
	if resp.Body != nil {
		bodyBytes, err = io.ReadAll(resp.Body)
		if err != nil {
			return "", err
		}
		resp.Body.Close()

		resp.Body = io.NopCloser(bytes.NewBuffer(bodyBytes))
	}

	if len(bodyBytes) > 1024 {
		respMsg := fmt.Sprintf("[%d] %s %s...(Total %dB)", resp.StatusCode, time.Since(st), bodyBytes[:1024],
			len(bodyBytes))
		return respMsg, nil
	}

	if len(bodyBytes) > 0 {
		respMsg := fmt.Sprintf("[%d] %s %s", resp.StatusCode, time.Since(st), bodyBytes)
		return respMsg, nil
	}

	respMsg := fmt.Sprintf("[%d] %s", resp.StatusCode, time.Since(st))
	return respMsg, nil
}

// curlLogTransport print curl log transport
type curlLogTransport struct {
	baseTransport http.RoundTripper
	maskKeys      map[string]struct{}
}

// RoundTrip curlLog Transport
func (c *curlLogTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	ctx := req.Context()
	st := time.Now()

	// 记录请求
	rbody, err := c.reqToCurl(req)
	if err != nil {
		return nil, err
	}
	slog.InfoContext(ctx, "curl REQ: "+rbody)

	resp, err := c.baseTransport.RoundTrip(req)
	if err != nil {
		slog.ErrorContext(ctx, "curl RESP: [err] "+err.Error())
		return nil, err
	}

	// 记录返回
	respBody, err := c.respToCurl(resp, st)
	if err != nil {
		return nil, err
	}
	slog.InfoContext(ctx, "curl RESP: "+respBody)

	return resp, nil
}