mirror of
https://github.com/buger/goreplay.git
synced 2024-04-21 12:32:02 +00:00
129 lines
2.4 KiB
Go
129 lines
2.4 KiB
Go
package gor
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
es "github.com/buger/gor/elasticsearch"
|
|
"io"
|
|
"log"
|
|
"net/http"
|
|
"net/url"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
)
|
|
|
|
type RedirectNotAllowed struct{}
|
|
|
|
func (e *RedirectNotAllowed) Error() string {
|
|
return "Redirects not allowed"
|
|
}
|
|
|
|
// customCheckRedirect disables redirects https://github.com/buger/gor/pull/15
|
|
func customCheckRedirect(req *http.Request, via []*http.Request) error {
|
|
if len(via) >= 0 {
|
|
return new(RedirectNotAllowed)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ParseRequest in []byte returns a http request or an error
|
|
func ParseRequest(data []byte) (request *http.Request, err error) {
|
|
buf := bytes.NewBuffer(data)
|
|
reader := bufio.NewReader(buf)
|
|
|
|
request, err = http.ReadRequest(reader)
|
|
|
|
return
|
|
}
|
|
|
|
type HTTPOutput struct {
|
|
address string
|
|
limit int
|
|
|
|
headers HTTPHeaders
|
|
|
|
elasticSearch *es.ESPlugin
|
|
}
|
|
|
|
func NewHTTPOutput(options string, headers HTTPHeaders, elasticSearchAddr string) io.Writer {
|
|
o := new(HTTPOutput)
|
|
|
|
optionsArr := strings.Split(options, "|")
|
|
address := optionsArr[0]
|
|
|
|
if !strings.HasPrefix(address, "http") {
|
|
address = "http://" + address
|
|
}
|
|
|
|
o.address = address
|
|
o.headers = headers
|
|
|
|
if elasticSearchAddr != "" {
|
|
o.elasticSearch = new(es.ESPlugin)
|
|
o.elasticSearch.Init(elasticSearchAddr)
|
|
}
|
|
|
|
if len(optionsArr) > 1 {
|
|
o.limit, _ = strconv.Atoi(optionsArr[1])
|
|
}
|
|
|
|
if o.limit > 0 {
|
|
return NewLimiter(o, o.limit)
|
|
} else {
|
|
return o
|
|
}
|
|
}
|
|
|
|
func (o *HTTPOutput) Write(data []byte) (n int, err error) {
|
|
go o.sendRequest(data)
|
|
|
|
return len(data), nil
|
|
}
|
|
|
|
func (o *HTTPOutput) sendRequest(data []byte) {
|
|
request, err := ParseRequest(data)
|
|
|
|
if err != nil {
|
|
log.Println("Can not parse request", string(data), err)
|
|
return
|
|
}
|
|
|
|
client := &http.Client{
|
|
CheckRedirect: customCheckRedirect,
|
|
}
|
|
|
|
// Change HOST of original request
|
|
URL := o.address + request.URL.Path + "?" + request.URL.RawQuery
|
|
|
|
request.RequestURI = ""
|
|
request.URL, _ = url.ParseRequestURI(URL)
|
|
|
|
for _, header := range o.headers {
|
|
request.Header.Set(header.Name, header.Value)
|
|
}
|
|
|
|
start := time.Now()
|
|
resp, err := client.Do(request)
|
|
stop := time.Now()
|
|
|
|
// We should not count Redirect as errors
|
|
if _, ok := err.(*RedirectNotAllowed); ok {
|
|
err = nil
|
|
}
|
|
|
|
if err == nil {
|
|
defer resp.Body.Close()
|
|
} else {
|
|
log.Println("Request error:", err)
|
|
}
|
|
|
|
if o.elasticSearch != nil {
|
|
o.elasticSearch.ResponseAnalyze(request, resp, start, stop)
|
|
}
|
|
}
|
|
|
|
func (o *HTTPOutput) String() string {
|
|
return "HTTP output: " + o.address
|
|
}
|