starnet/request_header.go

package starnet

import (
	"net/http"
)

func isHostHeaderKey(key string) bool {
	return http.CanonicalHeaderKey(key) == "Host"
}

func setRequestHostConfig(config *RequestConfig, host string) {
	if config == nil {
		return
	}
	if config.Headers == nil {
		config.Headers = make(http.Header)
	}
	config.Host = host
	if host == "" {
		config.Headers.Del("Host")
		return
	}
	config.Headers.Set("Host", host)
}

// SetHeader 设置 Header（覆盖）
func (r *Request) SetHeader(key, value string) *Request {
	if r.err != nil {
		return r
	}
	if r.doRaw {
		return r
	}
	if isHostHeaderKey(key) {
		return r.SetHost(value)
	}
	r.config.Headers.Set(key, value)
	r.invalidatePreparedState()
	return r
}

// AddHeader 添加 Header
func (r *Request) AddHeader(key, value string) *Request {
	if r.err != nil {
		return r
	}
	if r.doRaw {
		return r
	}
	if isHostHeaderKey(key) {
		return r.SetHost(value)
	}
	r.config.Headers.Add(key, value)
	r.invalidatePreparedState()
	return r
}

// SetHeaders 设置所有 Headers（覆盖）
func (r *Request) SetHeaders(headers http.Header) *Request {
	if r.err != nil {
		return r
	}
	if r.doRaw {
		return r
	}
	r.config.Headers = cloneHeader(headers)
	r.config.Host = r.config.Headers.Get("Host")
	r.syncRequestHost()
	r.invalidatePreparedState()
	return r
}

// AddHeaders 批量添加 Headers
func (r *Request) AddHeaders(headers map[string]string) *Request {
	if r.err != nil {
		return r
	}
	if r.doRaw {
		return r
	}
	for k, v := range headers {
		if isHostHeaderKey(k) {
			setRequestHostConfig(r.config, v)
			continue
		}
		r.config.Headers.Add(k, v)
	}
	r.syncRequestHost()
	r.invalidatePreparedState()
	return r
}

// DeleteHeader 删除 Header
func (r *Request) DeleteHeader(key string) *Request {
	if r.err != nil {
		return r
	}
	if r.doRaw {
		return r
	}
	if isHostHeaderKey(key) {
		setRequestHostConfig(r.config, "")
		r.syncRequestHost()
		r.invalidatePreparedState()
		return r
	}
	r.config.Headers.Del(key)
	r.invalidatePreparedState()
	return r
}

// GetHeader 获取 Header
func (r *Request) GetHeader(key string) string {
	if isHostHeaderKey(key) {
		return r.config.Host
	}
	return r.config.Headers.Get(key)
}

// Headers 获取所有 Headers
func (r *Request) Headers() http.Header {
	if r == nil || r.config == nil {
		return make(http.Header)
	}
	return cloneHeader(r.config.Headers)
}

// SetHost 设置请求 Host 头覆盖。
func (r *Request) SetHost(host string) *Request {
	if r.err != nil {
		return r
	}
	if r.doRaw {
		return r
	}
	setRequestHostConfig(r.config, host)
	r.syncRequestHost()
	r.invalidatePreparedState()
	return r
}

// Host 获取显式 Host 覆盖。
func (r *Request) Host() string {
	if r.config != nil && r.config.Host != "" {
		return r.config.Host
	}
	if r.httpReq != nil {
		return r.httpReq.Host
	}
	return ""
}

// SetContentType 设置 Content-Type
func (r *Request) SetContentType(contentType string) *Request {
	return r.SetHeader("Content-Type", contentType)
}

// SetUserAgent 设置 User-Agent
func (r *Request) SetUserAgent(userAgent string) *Request {
	return r.SetHeader("User-Agent", userAgent)
}

// SetReferer 设置 Referer
func (r *Request) SetReferer(referer string) *Request {
	return r.SetHeader("Referer", referer)
}

// SetBearerToken 设置 Bearer Token
func (r *Request) SetBearerToken(token string) *Request {
	return r.SetHeader("Authorization", "Bearer "+token)
}

// AddCookie 添加 Cookie
func (r *Request) AddCookie(cookie *http.Cookie) *Request {
	if r.err != nil {
		return r
	}
	if r.doRaw {
		return r
	}
	r.config.Cookies = append(r.config.Cookies, cloneCookie(cookie))
	r.invalidatePreparedState()
	return r
}

// AddSimpleCookie 添加简单 Cookie（path 为 /）
func (r *Request) AddSimpleCookie(name, value string) *Request {
	return r.AddCookie(&http.Cookie{
		Name:  name,
		Value: value,
		Path:  "/",
	})
}

// AddCookieKV 添加 Cookie（指定 path）
func (r *Request) AddCookieKV(name, value, path string) *Request {
	return r.AddCookie(&http.Cookie{
		Name:  name,
		Value: value,
		Path:  path,
	})
}

// SetCookies 设置所有 Cookies（覆盖）
func (r *Request) SetCookies(cookies []*http.Cookie) *Request {
	if r.err != nil {
		return r
	}
	if r.doRaw {
		return r
	}
	r.config.Cookies = cloneCookies(cookies)
	r.invalidatePreparedState()
	return r
}

// AddCookies 批量添加 Cookies
func (r *Request) AddCookies(cookies map[string]string) *Request {
	if r.err != nil {
		return r
	}
	if r.doRaw {
		return r
	}
	for name, value := range cookies {
		r.config.Cookies = append(r.config.Cookies, &http.Cookie{
			Name:  name,
			Value: value,
			Path:  "/",
		})
	}
	r.invalidatePreparedState()
	return r
}

// Cookies 获取所有 Cookies
func (r *Request) Cookies() []*http.Cookie {
	if r == nil || r.config == nil {
		return nil
	}
	return cloneCookies(r.config.Cookies)
}

// ResetHeaders 重置所有 Headers
func (r *Request) ResetHeaders() *Request {
	if r.err != nil {
		return r
	}
	r.config.Headers = make(http.Header)
	r.config.Host = ""
	r.syncRequestHost()
	r.invalidatePreparedState()
	return r
}

// ResetCookies 重置所有 Cookies
func (r *Request) ResetCookies() *Request {
	if r.err != nil {
		return r
	}
	r.config.Cookies = []*http.Cookie{}
	r.invalidatePreparedState()
	return r
}