mirror of
https://github.com/chainreactors/spray.git
synced 2025-09-15 11:40:13 +00:00
现在同时支持http.net与fasthttp两个库, 适用不同的场景
This commit is contained in:
parent
029a83faa8
commit
9582a32586
@ -2,46 +2,48 @@ package internal
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"github.com/chainreactors/parsers"
|
||||
"github.com/chainreactors/spray/pkg"
|
||||
"github.com/valyala/fasthttp"
|
||||
"github.com/chainreactors/spray/pkg/ihttp"
|
||||
"strconv"
|
||||
"strings"
|
||||
)
|
||||
|
||||
func NewBaseline(u *fasthttp.URI, resp *fasthttp.Response) *baseline {
|
||||
func NewBaseline(u, host string, resp *ihttp.Response) *baseline {
|
||||
bl := &baseline{
|
||||
Url: u,
|
||||
UrlString: u.String(),
|
||||
//Url: u,
|
||||
UrlString: u,
|
||||
Host: host,
|
||||
Status: resp.StatusCode(),
|
||||
IsValid: true,
|
||||
}
|
||||
|
||||
bl.Body = resp.Body()
|
||||
bl.BodyLength = resp.Header.ContentLength()
|
||||
bl.Header = resp.Header.Header()
|
||||
bl.HeaderLength = resp.Header.Len()
|
||||
bl.RedirectURL = string(resp.Header.Peek("Location"))
|
||||
bl.BodyLength = resp.ContentLength()
|
||||
bl.Header = resp.Header()
|
||||
bl.HeaderLength = len(bl.Header)
|
||||
bl.RedirectURL = resp.GetHeader("Location")
|
||||
bl.Raw = append(bl.Header, bl.Body...)
|
||||
return bl
|
||||
}
|
||||
|
||||
func NewInvalidBaseline(u *fasthttp.URI, resp *fasthttp.Response) *baseline {
|
||||
func NewInvalidBaseline(u, host string, resp *ihttp.Response) *baseline {
|
||||
bl := &baseline{
|
||||
Url: u,
|
||||
UrlString: u.String(),
|
||||
//Url: u,
|
||||
UrlString: u,
|
||||
Host: host,
|
||||
Status: resp.StatusCode(),
|
||||
IsValid: false,
|
||||
}
|
||||
|
||||
bl.RedirectURL = string(resp.Header.Peek("Location"))
|
||||
bl.RedirectURL = string(resp.GetHeader("Location"))
|
||||
|
||||
return bl
|
||||
}
|
||||
|
||||
type baseline struct {
|
||||
Url *fasthttp.URI `json:"-"`
|
||||
UrlString string `json:"url_string"`
|
||||
UrlString string `json:"url"`
|
||||
Host string `json:"host"`
|
||||
Body []byte `json:"-"`
|
||||
BodyLength int `json:"body_length"`
|
||||
Header []byte `json:"-"`
|
||||
@ -51,6 +53,7 @@ type baseline struct {
|
||||
Status int `json:"status"`
|
||||
IsDynamicUrl bool `json:"is_dynamic_url"` // 判断是否存在动态的url
|
||||
Spended int `json:"spended"` // 耗时, 毫秒
|
||||
Title string `json:"title"`
|
||||
Frameworks pkg.Frameworks `json:"frameworks"`
|
||||
Extracteds pkg.Extracteds `json:"extracts"`
|
||||
Err error `json:"-"`
|
||||
@ -61,6 +64,7 @@ type baseline struct {
|
||||
// Collect 深度收集信息
|
||||
func (bl *baseline) Collect() {
|
||||
bl.Hashes = parsers.NewHashes(bl.Raw)
|
||||
bl.Title = parsers.MatchTitle(string(bl.Body))
|
||||
// todo extract
|
||||
bl.Extracteds = pkg.Extractors.Extract(string(bl.Raw))
|
||||
// todo 指纹识别
|
||||
@ -95,11 +99,21 @@ func (bl *baseline) String() string {
|
||||
var line strings.Builder
|
||||
//line.WriteString("[+] ")
|
||||
line.WriteString(bl.UrlString)
|
||||
line.WriteString(fmt.Sprintf(" - %d - %d ", bl.Status, bl.BodyLength))
|
||||
line.WriteString(" (" + bl.Host + ")")
|
||||
line.WriteString(" - ")
|
||||
line.WriteString(strconv.Itoa(bl.Status))
|
||||
line.WriteString(" - ")
|
||||
line.WriteString(strconv.Itoa(bl.BodyLength))
|
||||
if bl.RedirectURL != "" {
|
||||
line.WriteString(" -> ")
|
||||
line.WriteString(bl.RedirectURL)
|
||||
line.WriteString(" ")
|
||||
}
|
||||
line.WriteString(" [" + bl.Title + "]")
|
||||
if bl.Hashes != nil {
|
||||
line.WriteString(" [" + bl.Hashes.BodyMd5 + "]")
|
||||
}
|
||||
|
||||
line.WriteString(bl.Frameworks.ToString())
|
||||
//line.WriteString(bl.Extracteds)
|
||||
//line.WriteString("\n")
|
||||
|
@ -4,9 +4,11 @@ import (
|
||||
"context"
|
||||
"github.com/chainreactors/logs"
|
||||
"github.com/chainreactors/spray/pkg"
|
||||
"github.com/chainreactors/spray/pkg/ihttp"
|
||||
"github.com/chainreactors/words"
|
||||
"github.com/panjf2000/ants/v2"
|
||||
"github.com/valyala/fasthttp"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
@ -25,7 +27,7 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
|
||||
Config: config,
|
||||
ctx: pctx,
|
||||
cancel: cancel,
|
||||
client: pkg.NewClient(config.Thread, 2),
|
||||
client: ihttp.NewClient(config.Thread, 2, config.ClientType),
|
||||
worder: words.NewWorder(config.Wordlist),
|
||||
outputCh: outputCh,
|
||||
tempCh: make(chan *baseline, config.Thread),
|
||||
@ -37,13 +39,32 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
|
||||
|
||||
switch config.Mod {
|
||||
case pkg.PathSpray:
|
||||
pool.genReq = func(s string) (*fasthttp.Request, error) {
|
||||
pool.genReq = func(s string) (*ihttp.Request, error) {
|
||||
return pool.buildPathRequest(s)
|
||||
}
|
||||
pool.check = func() {
|
||||
pool.wg.Add(1)
|
||||
_ = pool.pool.Invoke(newUnit(pkg.RandPath(), CheckSource))
|
||||
|
||||
if pool.failedCount > breakThreshold {
|
||||
// 当报错次数超过上限是, 结束任务
|
||||
pool.cancel()
|
||||
}
|
||||
}
|
||||
case pkg.HostSpray:
|
||||
pool.genReq = func(s string) (*fasthttp.Request, error) {
|
||||
pool.genReq = func(s string) (*ihttp.Request, error) {
|
||||
return pool.buildHostRequest(s)
|
||||
}
|
||||
|
||||
pool.check = func() {
|
||||
pool.wg.Add(1)
|
||||
_ = pool.pool.Invoke(newUnit(pkg.RandHost(), CheckSource))
|
||||
|
||||
if pool.failedCount > breakThreshold {
|
||||
// 当报错次数超过上限是, 结束任务
|
||||
pool.cancel()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
p, _ := ants.NewPoolWithFunc(config.Thread, func(i interface{}) {
|
||||
@ -56,8 +77,11 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
|
||||
|
||||
var bl *baseline
|
||||
resp, reqerr := pool.client.Do(pctx, req)
|
||||
defer fasthttp.ReleaseResponse(resp)
|
||||
defer fasthttp.ReleaseRequest(req)
|
||||
if pool.ClientType == ihttp.FAST {
|
||||
defer fasthttp.ReleaseResponse(resp.FastResponse)
|
||||
defer fasthttp.ReleaseRequest(req.FastRequest)
|
||||
}
|
||||
|
||||
if reqerr != nil && reqerr != fasthttp.ErrBodyTooLarge {
|
||||
pool.failedCount++
|
||||
bl = &baseline{UrlString: pool.BaseURL + unit.path, Err: reqerr}
|
||||
@ -65,9 +89,9 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
|
||||
pool.failedCount = 0
|
||||
if err = pool.PreCompare(resp); err == nil || unit.source == CheckSource {
|
||||
// 通过预对比跳过一些无用数据, 减少性能消耗
|
||||
bl = NewBaseline(req.URI(), resp)
|
||||
bl = NewBaseline(req.URI(), req.Host(), resp)
|
||||
} else {
|
||||
bl = NewInvalidBaseline(req.URI(), resp)
|
||||
bl = NewInvalidBaseline(req.URI(), req.Host(), resp)
|
||||
}
|
||||
bl.Err = reqerr
|
||||
}
|
||||
@ -108,7 +132,7 @@ func NewPool(ctx context.Context, config *pkg.Config, outputCh chan *baseline) (
|
||||
|
||||
type Pool struct {
|
||||
*pkg.Config
|
||||
client *pkg.Client
|
||||
client *ihttp.Client
|
||||
pool *ants.PoolWithFunc
|
||||
bar *pkg.Bar
|
||||
ctx context.Context
|
||||
@ -122,22 +146,13 @@ type Pool struct {
|
||||
checkPeriod int
|
||||
errPeriod int
|
||||
analyzeDone bool
|
||||
genReq func(s string) (*fasthttp.Request, error)
|
||||
genReq func(s string) (*ihttp.Request, error)
|
||||
check func()
|
||||
worder *words.Worder
|
||||
wg sync.WaitGroup
|
||||
initwg sync.WaitGroup // 初始化用, 之后改成锁
|
||||
}
|
||||
|
||||
func (p *Pool) check() {
|
||||
p.wg.Add(1)
|
||||
_ = p.pool.Invoke(newUnit(pkg.RandPath(), CheckSource))
|
||||
|
||||
if p.failedCount > breakThreshold {
|
||||
// 当报错次数超过上限是, 结束任务
|
||||
p.cancel()
|
||||
}
|
||||
}
|
||||
|
||||
func (p *Pool) Init() error {
|
||||
p.initwg.Add(1)
|
||||
p.check()
|
||||
@ -189,12 +204,12 @@ Loop:
|
||||
p.Close()
|
||||
}
|
||||
|
||||
func (p *Pool) PreCompare(resp *fasthttp.Response) error {
|
||||
func (p *Pool) PreCompare(resp *ihttp.Response) error {
|
||||
if !CheckStatusCode(resp.StatusCode()) {
|
||||
return ErrBadStatus
|
||||
}
|
||||
|
||||
if CheckRedirect != nil && !CheckRedirect(string(resp.Header.Peek("Location"))) {
|
||||
if CheckRedirect != nil && !CheckRedirect(string(resp.GetHeader("Location"))) {
|
||||
return ErrRedirect
|
||||
}
|
||||
|
||||
@ -238,15 +253,26 @@ func (p *Pool) Close() {
|
||||
}
|
||||
}
|
||||
|
||||
func (p *Pool) buildPathRequest(path string) (*fasthttp.Request, error) {
|
||||
func (p *Pool) buildPathRequest(path string) (*ihttp.Request, error) {
|
||||
if p.Config.ClientType == ihttp.FAST {
|
||||
req := fasthttp.AcquireRequest()
|
||||
req.SetRequestURI(p.BaseURL + path)
|
||||
return req, nil
|
||||
return &ihttp.Request{FastRequest: req}, nil
|
||||
} else {
|
||||
req, err := http.NewRequest("GET", p.BaseURL+path, nil)
|
||||
return &ihttp.Request{StandardRequest: req}, err
|
||||
}
|
||||
}
|
||||
|
||||
func (p *Pool) buildHostRequest(host string) (*fasthttp.Request, error) {
|
||||
func (p *Pool) buildHostRequest(host string) (*ihttp.Request, error) {
|
||||
if p.Config.ClientType == ihttp.FAST {
|
||||
req := fasthttp.AcquireRequest()
|
||||
req.SetRequestURI(p.BaseURL)
|
||||
req.SetHost(host)
|
||||
return req, nil
|
||||
return &ihttp.Request{FastRequest: req}, nil
|
||||
} else {
|
||||
req, err := http.NewRequest("GET", p.BaseURL, nil)
|
||||
req.Host = host
|
||||
return &ihttp.Request{StandardRequest: req}, err
|
||||
}
|
||||
}
|
||||
|
@ -5,7 +5,9 @@ import (
|
||||
"fmt"
|
||||
"github.com/chainreactors/logs"
|
||||
"github.com/chainreactors/spray/pkg"
|
||||
"github.com/chainreactors/spray/pkg/ihttp"
|
||||
"github.com/gosuri/uiprogress"
|
||||
"github.com/panjf2000/ants/v2"
|
||||
"io/ioutil"
|
||||
"net/http"
|
||||
"os"
|
||||
@ -27,8 +29,9 @@ type Runner struct {
|
||||
Offset int `long:"offset"`
|
||||
Limit int `long:"limit"`
|
||||
Threads int `short:"t" long:"thread" default:"20"`
|
||||
PoolSize int `short:"p" long:"pool"`
|
||||
Pools map[string]*Pool
|
||||
PoolSize int `short:"p" long:"pool" default:"5"`
|
||||
Pools *ants.PoolWithFunc
|
||||
poolwg sync.WaitGroup
|
||||
Deadline int `long:"deadline" default:"600"` // todo 总的超时时间,适配云函数的deadline
|
||||
Debug bool `long:"debug"`
|
||||
Quiet bool `short:"q" long:"quiet"`
|
||||
@ -108,19 +111,10 @@ func (r *Runner) Prepare() error {
|
||||
}
|
||||
|
||||
r.OutputCh = make(chan *baseline, 100)
|
||||
r.Pools = make(map[string]*Pool)
|
||||
go r.Outputting()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *Runner) Run() {
|
||||
// todo pool 结束与并发控制
|
||||
ctx := context.Background()
|
||||
var wg sync.WaitGroup
|
||||
for _, u := range r.URLList {
|
||||
wg.Add(1)
|
||||
u := u
|
||||
go func() {
|
||||
|
||||
r.Pools, err = ants.NewPoolWithFunc(r.PoolSize, func(i interface{}) {
|
||||
u := i.(string)
|
||||
config := &pkg.Config{
|
||||
BaseURL: u,
|
||||
Wordlist: r.Wordlist,
|
||||
@ -130,6 +124,13 @@ func (r *Runner) Run() {
|
||||
Mod: pkg.ModMap[r.Mod],
|
||||
DeadlineTime: r.Deadline,
|
||||
}
|
||||
|
||||
if config.Mod == pkg.PathSpray {
|
||||
config.ClientType = ihttp.FAST
|
||||
} else if config.Mod == pkg.HostSpray {
|
||||
config.ClientType = ihttp.STANDARD
|
||||
}
|
||||
|
||||
pool, err := NewPool(ctx, config, r.OutputCh)
|
||||
if err != nil {
|
||||
logs.Log.Error(err.Error())
|
||||
@ -141,13 +142,21 @@ func (r *Runner) Run() {
|
||||
logs.Log.Error(err.Error())
|
||||
return
|
||||
}
|
||||
r.Pools[u] = pool
|
||||
// todo pool 总超时时间
|
||||
pool.Run(ctx)
|
||||
wg.Done()
|
||||
}()
|
||||
r.poolwg.Done()
|
||||
})
|
||||
go r.Outputting()
|
||||
return nil
|
||||
}
|
||||
wg.Wait()
|
||||
|
||||
func (r *Runner) Run() {
|
||||
// todo pool 结束与并发控制
|
||||
for _, u := range r.URLList {
|
||||
r.poolwg.Add(1)
|
||||
r.Pools.Invoke(u)
|
||||
}
|
||||
r.poolwg.Wait()
|
||||
for {
|
||||
if len(r.OutputCh) == 0 {
|
||||
close(r.OutputCh)
|
||||
|
@ -1,67 +0,0 @@
|
||||
package pkg
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"github.com/valyala/fasthttp"
|
||||
"net/http"
|
||||
"time"
|
||||
)
|
||||
|
||||
var (
|
||||
DefaultMaxBodySize = 1024 * 100 // 100k
|
||||
)
|
||||
|
||||
func NewClient(thread int, timeout int) *Client {
|
||||
c := &Client{
|
||||
client: &fasthttp.Client{
|
||||
TLSConfig: &tls.Config{
|
||||
Renegotiation: tls.RenegotiateOnceAsClient,
|
||||
InsecureSkipVerify: true,
|
||||
},
|
||||
//ReadBufferSize: 20480,
|
||||
MaxConnsPerHost: thread * 2,
|
||||
MaxIdleConnDuration: time.Duration(timeout) * time.Second,
|
||||
MaxConnWaitTimeout: time.Duration(timeout) * time.Second,
|
||||
ReadTimeout: time.Duration(timeout) * time.Second,
|
||||
WriteTimeout: time.Duration(timeout) * time.Second,
|
||||
MaxResponseBodySize: DefaultMaxBodySize,
|
||||
},
|
||||
timeout: time.Duration(timeout) * time.Second,
|
||||
}
|
||||
//c := &Client{
|
||||
// client: &http.Client{
|
||||
// Transport: tr,
|
||||
// Timeout: time.Second * time.Duration(timeout),
|
||||
// CheckRedirect: checkRedirect,
|
||||
// },
|
||||
//}
|
||||
|
||||
//c.Method = method
|
||||
//c.Headers = Opt.Headers
|
||||
//c.Mod = Opt.Mod
|
||||
|
||||
return c
|
||||
}
|
||||
|
||||
type Client struct {
|
||||
client *fasthttp.Client
|
||||
timeout time.Duration
|
||||
}
|
||||
|
||||
func (c *Client) Do(ctx context.Context, req *fasthttp.Request) (*fasthttp.Response, error) {
|
||||
//if req.Header == nil {
|
||||
// req.Header = c.Headers
|
||||
//}
|
||||
resp := fasthttp.AcquireResponse()
|
||||
return resp, c.client.Do(req, resp)
|
||||
}
|
||||
|
||||
var MaxRedirects = 0
|
||||
var checkRedirect = func(req *http.Request, via []*http.Request) error {
|
||||
if len(via) > MaxRedirects {
|
||||
return http.ErrUseLastResponse
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
@ -29,4 +29,5 @@ type Config struct {
|
||||
Headers http.Header
|
||||
DeadlineTime int
|
||||
EnableFuzzy bool
|
||||
ClientType int
|
||||
}
|
||||
|
94
pkg/ihttp/client.go
Normal file
94
pkg/ihttp/client.go
Normal file
@ -0,0 +1,94 @@
|
||||
package ihttp
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/tls"
|
||||
"fmt"
|
||||
"github.com/valyala/fasthttp"
|
||||
"net/http"
|
||||
"time"
|
||||
)
|
||||
|
||||
var (
|
||||
DefaultMaxBodySize = 1024 * 100 // 100k
|
||||
)
|
||||
|
||||
const (
|
||||
FAST = iota
|
||||
STANDARD
|
||||
)
|
||||
|
||||
func NewClient(thread int, timeout int, clientType int) *Client {
|
||||
if clientType == FAST {
|
||||
return &Client{
|
||||
fastClient: &fasthttp.Client{
|
||||
TLSConfig: &tls.Config{
|
||||
Renegotiation: tls.RenegotiateOnceAsClient,
|
||||
InsecureSkipVerify: true,
|
||||
},
|
||||
//ReadBufferSize: 20480,
|
||||
MaxConnsPerHost: thread * 2,
|
||||
MaxIdleConnDuration: time.Duration(timeout) * time.Second,
|
||||
MaxConnWaitTimeout: time.Duration(timeout) * time.Second,
|
||||
ReadTimeout: time.Duration(timeout) * time.Second,
|
||||
WriteTimeout: time.Duration(timeout) * time.Second,
|
||||
MaxResponseBodySize: DefaultMaxBodySize,
|
||||
},
|
||||
timeout: time.Duration(timeout) * time.Second,
|
||||
}
|
||||
} else {
|
||||
return &Client{
|
||||
standardClient: &http.Client{
|
||||
Transport: &http.Transport{
|
||||
//Proxy: Proxy,
|
||||
//TLSHandshakeTimeout : delay * time.Second,
|
||||
TLSClientConfig: &tls.Config{
|
||||
Renegotiation: tls.RenegotiateOnceAsClient,
|
||||
InsecureSkipVerify: true,
|
||||
},
|
||||
MaxConnsPerHost: thread,
|
||||
IdleConnTimeout: time.Duration(timeout) * time.Second,
|
||||
},
|
||||
Timeout: time.Second * time.Duration(timeout),
|
||||
CheckRedirect: checkRedirect,
|
||||
},
|
||||
timeout: time.Duration(timeout) * time.Second,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
type Client struct {
|
||||
fastClient *fasthttp.Client
|
||||
standardClient *http.Client
|
||||
timeout time.Duration
|
||||
}
|
||||
|
||||
func (c *Client) FastDo(ctx context.Context, req *fasthttp.Request) (*fasthttp.Response, error) {
|
||||
resp := fasthttp.AcquireResponse()
|
||||
return resp, c.fastClient.Do(req, resp)
|
||||
}
|
||||
|
||||
func (c *Client) StandardDo(ctx context.Context, req *http.Request) (*http.Response, error) {
|
||||
return c.standardClient.Do(req)
|
||||
}
|
||||
|
||||
func (c *Client) Do(ctx context.Context, req *Request) (*Response, error) {
|
||||
if c.fastClient != nil {
|
||||
resp, err := c.FastDo(ctx, req.FastRequest)
|
||||
return &Response{FastResponse: resp}, err
|
||||
} else if c.standardClient != nil {
|
||||
resp, err := c.StandardDo(ctx, req.StandardRequest)
|
||||
return &Response{StandardResponse: resp}, err
|
||||
} else {
|
||||
return nil, fmt.Errorf("not found client")
|
||||
}
|
||||
}
|
||||
|
||||
var MaxRedirects = 0
|
||||
var checkRedirect = func(req *http.Request, via []*http.Request) error {
|
||||
if len(via) > MaxRedirects {
|
||||
return http.ErrUseLastResponse
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
31
pkg/ihttp/request.go
Normal file
31
pkg/ihttp/request.go
Normal file
@ -0,0 +1,31 @@
|
||||
package ihttp
|
||||
|
||||
import (
|
||||
"github.com/valyala/fasthttp"
|
||||
"net/http"
|
||||
)
|
||||
|
||||
type Request struct {
|
||||
StandardRequest *http.Request
|
||||
FastRequest *fasthttp.Request
|
||||
}
|
||||
|
||||
func (r *Request) URI() string {
|
||||
if r.FastRequest != nil {
|
||||
return r.FastRequest.URI().String()
|
||||
} else if r.StandardRequest != nil {
|
||||
return r.StandardRequest.URL.String()
|
||||
} else {
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Request) Host() string {
|
||||
if r.FastRequest != nil {
|
||||
return string(r.FastRequest.Host())
|
||||
} else if r.StandardRequest != nil {
|
||||
return r.StandardRequest.Host
|
||||
} else {
|
||||
return ""
|
||||
}
|
||||
}
|
83
pkg/ihttp/response.go
Normal file
83
pkg/ihttp/response.go
Normal file
@ -0,0 +1,83 @@
|
||||
package ihttp
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"github.com/chainreactors/logs"
|
||||
"github.com/valyala/fasthttp"
|
||||
"io"
|
||||
"net/http"
|
||||
)
|
||||
|
||||
type Response struct {
|
||||
StandardResponse *http.Response
|
||||
FastResponse *fasthttp.Response
|
||||
}
|
||||
|
||||
func (r *Response) StatusCode() int {
|
||||
if r.FastResponse != nil {
|
||||
return r.FastResponse.StatusCode()
|
||||
} else if r.StandardResponse != nil {
|
||||
return r.StandardResponse.StatusCode
|
||||
} else {
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Response) Body() []byte {
|
||||
if r.FastResponse != nil {
|
||||
return r.FastResponse.Body()
|
||||
} else if r.StandardResponse != nil {
|
||||
body := make([]byte, 20480)
|
||||
if r.StandardResponse.ContentLength > 0 {
|
||||
n, err := io.ReadFull(r.StandardResponse.Body, body)
|
||||
_ = r.StandardResponse.Body.Close()
|
||||
if err == nil {
|
||||
return body
|
||||
} else if err == io.ErrUnexpectedEOF {
|
||||
return body[:n]
|
||||
} else {
|
||||
logs.Log.Error("readfull failed" + err.Error())
|
||||
return nil
|
||||
}
|
||||
}
|
||||
return nil
|
||||
} else {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Response) ContentLength() int {
|
||||
if r.FastResponse != nil {
|
||||
return r.FastResponse.Header.ContentLength()
|
||||
} else if r.StandardResponse != nil {
|
||||
return int(r.StandardResponse.ContentLength)
|
||||
} else {
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Response) Header() []byte {
|
||||
if r.FastResponse != nil {
|
||||
return r.FastResponse.Header.Header()
|
||||
} else if r.StandardResponse != nil {
|
||||
var header bytes.Buffer
|
||||
for k, v := range r.StandardResponse.Header {
|
||||
for _, i := range v {
|
||||
header.WriteString(k + ": " + i + "\r\n")
|
||||
}
|
||||
}
|
||||
return header.Bytes()
|
||||
} else {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func (r *Response) GetHeader(key string) string {
|
||||
if r.FastResponse != nil {
|
||||
return string(r.FastResponse.Header.Peek(key))
|
||||
} else if r.StandardResponse != nil {
|
||||
return r.StandardResponse.Header.Get(key)
|
||||
} else {
|
||||
return ""
|
||||
}
|
||||
}
|
20
pkg/utils.go
20
pkg/utils.go
@ -56,3 +56,23 @@ func RandPath() string {
|
||||
}
|
||||
return *(*string)(unsafe.Pointer(&b))
|
||||
}
|
||||
|
||||
func RandHost() string {
|
||||
n := 8
|
||||
b := make([]byte, n)
|
||||
// A rand.Int63() generates 63 random bits, enough for letterIdMax letters!
|
||||
for i, cache, remain := n-1, src.Int63(), letterIdMax; i >= 1; {
|
||||
if remain == 0 {
|
||||
cache, remain = src.Int63(), letterIdMax
|
||||
}
|
||||
if idx := int(cache & letterIdMask); idx < len(letters) {
|
||||
b[i] = letters[idx]
|
||||
i--
|
||||
}
|
||||
cache >>= letterIdBits
|
||||
remain--
|
||||
}
|
||||
|
||||
b[5] = byte(0x2e)
|
||||
return *(*string)(unsafe.Pointer(&b))
|
||||
}
|
||||
|
Loading…
x
Reference in New Issue
Block a user