starnet/response.go

package starnet

import (
	"bytes"
	"encoding/json"
	"io"
	"net/http"
	"sync"
)

// Response HTTP 响应
type Response struct {
	*http.Response
	request      *Request
	httpClient   *http.Client
	cancel       func()
	body         *Body
	traceSummary *TraceSummary
}

// Body 响应体
type Body struct {
	raw      io.ReadCloser
	data     []byte
	consumed bool
	maxBytes int64
	mu       sync.Mutex
}

type cancelReadCloser struct {
	io.ReadCloser
	cancel func()
	once   sync.Once
}

func (c *cancelReadCloser) Close() error {
	err := c.ReadCloser.Close()
	c.once.Do(func() {
		if c.cancel != nil {
			c.cancel()
		}
	})
	return err
}

// Request 获取原始请求
func (r *Response) Request() *Request {
	return r.request
}

// TraceSummary 获取当前响应对应的 trace 摘要快照。
func (r *Response) TraceSummary() *TraceSummary {
	if r == nil || r.traceSummary == nil {
		return nil
	}
	summary := cloneTraceSummary(*r.traceSummary)
	return &summary
}

// Body 获取响应体
func (r *Response) Body() *Body {
	return r.body
}

// Close 关闭响应体
func (r *Response) Close() error {
	if r == nil {
		return nil
	}
	if r.body != nil && r.body.raw != nil {
		return r.body.raw.Close()
	}
	if r.cancel != nil {
		r.cancel()
		r.cancel = nil
	}
	return nil
}

// CloseWithClient 关闭响应体并关闭空闲连接
func (r *Response) CloseWithClient() error {
	if r == nil {
		return nil
	}
	if r.httpClient != nil {
		r.httpClient.CloseIdleConnections()
	}
	return r.Close()
}

// readAll 读取所有数据
func (b *Body) readAll() error {
	b.mu.Lock()
	defer b.mu.Unlock()

	if b.consumed {
		return nil
	}

	if b.raw == nil {
		b.consumed = true
		return nil
	}

	reader := io.Reader(b.raw)
	if b.maxBytes > 0 {
		reader = io.LimitReader(b.raw, b.maxBytes+1)
	}

	data, err := io.ReadAll(reader)
	if err != nil {
		return wrapError(err, "read response body")
	}
	if b.maxBytes > 0 && int64(len(data)) > b.maxBytes {
		b.consumed = true
		_ = b.raw.Close()
		return wrapError(ErrRespBodyTooLarge, "response body exceeds max bytes: %d > %d", len(data), b.maxBytes)
	}

	b.data = data
	b.consumed = true
	_ = b.raw.Close()

	return nil
}

// Bytes 获取响应体字节
func (b *Body) Bytes() ([]byte, error) {
	if err := b.readAll(); err != nil {
		return nil, err
	}
	return b.data, nil
}

// String 获取响应体字符串
func (b *Body) String() (string, error) {
	data, err := b.Bytes()
	if err != nil {
		return "", err
	}
	return string(data), nil
}

// JSON 解析 JSON 响应
func (b *Body) JSON(v interface{}) error {
	data, err := b.Bytes()
	if err != nil {
		return err
	}
	return json.Unmarshal(data, v)
}

// Reader 获取 Reader（只能调用一次）
func (b *Body) Reader() (io.ReadCloser, error) {
	b.mu.Lock()
	defer b.mu.Unlock()

	if b.consumed {
		if b.data != nil {
			// 已读取，返回缓存数据的 Reader
			return io.NopCloser(bytes.NewReader(b.data)), nil
		}
		return nil, ErrBodyAlreadyConsumed
	}

	b.consumed = true
	return b.raw, nil
}

// IsConsumed 检查是否已消费
func (b *Body) IsConsumed() bool {
	b.mu.Lock()
	defer b.mu.Unlock()
	return b.consumed
}

// Close 关闭 Body
func (b *Body) Close() error {
	b.mu.Lock()
	defer b.mu.Unlock()

	if b.raw != nil {
		return b.raw.Close()
	}
	return nil
}

// MustBytes 获取响应体字节（忽略错误，失败返回 nil）
func (b *Body) MustBytes() []byte {
	data, err := b.Bytes()
	if err != nil {
		return nil
	}
	return data
}

// MustString 获取响应体字符串（忽略错误，失败返回空串）
func (b *Body) MustString() string {
	s, err := b.String()
	if err != nil {
		return ""
	}
	return s
}

// Unmarshal 解析 JSON 响应（兼容旧 API）
func (b *Body) Unmarshal(v interface{}) error {
	return b.JSON(v)
}