162 lines
2.7 KiB
Go
162 lines
2.7 KiB
Go
package starnet
|
||
|
||
import (
|
||
"bytes"
|
||
"encoding/json"
|
||
"io"
|
||
"net/http"
|
||
"sync"
|
||
)
|
||
|
||
// Response HTTP 响应
|
||
type Response struct {
|
||
*http.Response
|
||
request *Request
|
||
httpClient *http.Client
|
||
body *Body
|
||
}
|
||
|
||
// Body 响应体
|
||
type Body struct {
|
||
raw io.ReadCloser
|
||
data []byte
|
||
consumed bool
|
||
mu sync.Mutex
|
||
}
|
||
|
||
// Request 获取原始请求
|
||
func (r *Response) Request() *Request {
|
||
return r.request
|
||
}
|
||
|
||
// Body 获取响应体
|
||
func (r *Response) Body() *Body {
|
||
return r.body
|
||
}
|
||
|
||
// Close 关闭响应体
|
||
func (r *Response) Close() error {
|
||
if r.body != nil && r.body.raw != nil {
|
||
return r.body.raw.Close()
|
||
}
|
||
return nil
|
||
}
|
||
|
||
// CloseWithClient 关闭响应体并关闭空闲连接
|
||
func (r *Response) CloseWithClient() error {
|
||
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
|
||
}
|
||
|
||
data, err := io.ReadAll(b.raw)
|
||
if err != nil {
|
||
return wrapError(err, "read response body")
|
||
}
|
||
|
||
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)
|
||
}
|