requests/session.go

package requests

import (
	"crypto/tls"
	"errors"
	"net/http"
	"net/http/cookiejar"
	"net/url"
	"reflect"
	"runtime"
	"time"

	"golang.org/x/net/publicsuffix"
)

// Body 相关参数结构
type Body struct {
	// Query       map[string][]string
	IOBody interface{}
	// Files       []UploadFile
	ContentType string
}

// BasicAuth 帐号认真结构
type BasicAuth struct {
	// User 帐号
	User string
	// Password 密码
	Password string
}

// Session 的基本方法
type Session struct {
	client    *http.Client
	transport *http.Transport
	cookiejar http.CookieJar
	body      *Body
	auth      *BasicAuth

	Header http.Header
	Query  url.Values
}

const (
	// TypeJSON 类型
	TypeJSON = "application/json"
	// TypeXML 类型
	TypeXML = "text/xml"
	// TypeURLENCODED 类型
	TypeURLENCODED = "application/x-www-form-urlencoded"
	// TypeFormData 类型
	TypeFormData = "multipart/form-data"
	// HeaderKeyHost Host
	HeaderKeyHost = "Host"
	// HeaderKeyUA User-Agent
	HeaderKeyUA = "User-Agent"
	// HeaderKeyContentType Content-Type
	HeaderKeyContentType = "Content-Type"
)

// TypeConfig 配置类型
type TypeConfig int

const (
	_ TypeConfig = iota
	// ConfigRequestTimeout request 包括 dial request redirect 总时间超时
	ConfigRequestTimeout // 支持time.Duration 和 int(秒为单位)

	// ConfigDialTimeout 一个Connect过程的Timeout
	ConfigDialTimeout // 支持time.Duration 和 int(秒为单位)

	// ConfigProxy 代理链接
	ConfigProxy // http, https, socks5

	// ConfigInsecure InsecureSkipVerify
	ConfigInsecure // true, false

	// ConfigBasicAuth 帐号认证
	ConfigBasicAuth // user pwd

	// ConfigTLS 帐号认证
	ConfigTLS // user pwd

	// ConfigCookiejar 持久化 CookieJar
	ConfigCookiejar // true or false ; default = true
)

// NewSession 创建Session
func NewSession() *Session {
	client := &http.Client{}
	transport := &http.Transport{DisableCompression: true}

	client.Transport = transport
	cjar, err := cookiejar.New(&cookiejar.Options{PublicSuffixList: publicsuffix.List})
	if err != nil {
		panic(err)
	}

	client.Jar = cjar
	return &Session{client: client, body: &Body{}, transport: transport, auth: nil, cookiejar: client.Jar, Header: make(http.Header)}
}

// SetConfig 设置配置
func (ses *Session) SetConfig(typeConfig TypeConfig, values interface{}) {

	switch typeConfig {
	case ConfigRequestTimeout:
		switch v := values.(type) {
		case time.Duration:
			ses.client.Timeout = v
		case int:
			ses.client.Timeout = time.Duration(v * int(time.Second))
		case int64:
			ses.client.Timeout = time.Duration(v * int64(time.Second))
		case float32:
			ses.client.Timeout = time.Duration(v * float32(time.Second))
		case float64:
			ses.client.Timeout = time.Duration(v * float64(time.Second))
		default:
			panic(errors.New("error type " + reflect.TypeOf(v).String()))
		}
	case ConfigDialTimeout:
		// 没时间实现这些小细节
	case ConfigCookiejar:
		v := values.(bool)
		if v {
			if ses.cookiejar == nil {
				j, err := cookiejar.New(nil)
				if err != nil {
					panic(err)
				}
				ses.cookiejar = j
			}
		} else {
			ses.cookiejar = nil
		}
	case ConfigProxy:
		switch v := values.(type) {
		case string:
			purl, err := (url.Parse(v))
			if err != nil {
				panic(err)
			}
			ses.transport.Proxy = http.ProxyURL(purl)
		case *url.URL:
			ses.transport.Proxy = http.ProxyURL(v)
		}
	case ConfigInsecure:
		ses.transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: !values.(bool)}
	case ConfigTLS:
		ses.transport.TLSClientConfig = values.(*tls.Config)
	case ConfigBasicAuth:
		if ses.auth == nil {
			ses.auth = &BasicAuth{}
		}

		switch v := values.(type) {
		case *BasicAuth:
			ses.auth.User = v.User
			ses.auth.User = v.Password
		case BasicAuth:
			ses.auth.User = v.User
			ses.auth.User = v.Password
		case []string:
			ses.auth.User = v[0]
			ses.auth.User = v[1]
		case nil:
			ses.auth = nil
		}
	default:
		panic(errors.New("unknown typeConfig " + reflect.TypeOf(typeConfig).String()))
	}
	return
}

// SetCookies 设置Cookies 或者添加Cookies Del
func (ses *Session) SetCookies(u *url.URL, cookies []*http.Cookie) {
	ses.cookiejar.SetCookies(u, cookies)
}

// Cookies 返回 Cookies
func (ses *Session) Cookies(u *url.URL) []*http.Cookie {
	return ses.cookiejar.Cookies(u)
}

// DelCookies 删除 Cookies
func (ses *Session) DelCookies(u *url.URL, name string) {
	cookies := ses.cookiejar.Cookies(u)
	for _, c := range cookies {
		if c.Name == name {
			c.MaxAge = -1
		}
	}
	ses.SetCookies(u, cookies)
}

// Get 请求
func (ses *Session) Get(url string) *Workflow {
	wf := NewWorkflow(ses)
	wf.Method = "GET"
	wf.SetURL(url)
	return wf
}

// Post 请求
func (ses *Session) Post(url string) *Workflow {
	wf := NewWorkflow(ses)
	wf.Method = "POST"
	wf.SetURL(url)
	return wf
}

// Put 请求
func (ses *Session) Put(url string) *Workflow {
	wf := NewWorkflow(ses)
	wf.Method = "PUT"
	wf.SetURL(url)
	return wf
}

// Patch 请求
func (ses *Session) Patch(url string) *Workflow {
	wf := NewWorkflow(ses)
	wf.Method = "PATCH"
	wf.SetURL(url)
	return wf
}

// Delete 请求
func (ses *Session) Delete(url string) *Workflow {
	wf := NewWorkflow(ses)
	wf.Method = "DELETE"
	wf.SetURL(url)
	return wf
}

// Head 请求
func (ses *Session) Head(url string) *Workflow {
	wf := NewWorkflow(ses)
	wf.Method = "HEAD"
	wf.SetURL(url)
	return wf
}

// Options 请求
func (ses *Session) Options(url string) *Workflow {
	wf := NewWorkflow(ses)
	wf.Method = "OPTIONS"
	wf.SetURL(url)
	return wf
}

// CloseIdleConnections  closes the idle connections that a session client may make use of
// 从levigross/grequests 借鉴
func (ses *Session) CloseIdleConnections() {
	ses.client.Transport.(*http.Transport).CloseIdleConnections()
}

// EnsureTransporterFinalized will ensure that when the HTTP client is GCed
// the runtime will close the idle connections (so that they won't leak)
// this function was adopted from Hashicorp's go-cleanhttp package
// 暂时不用, 标记到以后是否起作用
func EnsureTransporterFinalized(httpTransport *http.Transport) {
	runtime.SetFinalizer(&httpTransport, func(transportInt **http.Transport) {
		(*transportInt).CloseIdleConnections()
	})
}