github/goreplay
4
0
Fork 0
mirror of https://github.com/buger/goreplay.git synced 2024-04-21 12:32:02 +00:00
Code Issues Releases Wiki Activity

changes plugins reader and writer method

Browse Source 
// PluginReader is an interface for input plugins
type PluginReader interface {
	PluginRead() (msg *Message, err error)
}

// PluginWriter is an interface for output plugins
type PluginWriter interface {
	PluginWrite(msg *Message) (n int, err error)
}
This commit is contained in:
Urban Ishimwe
2020-11-02 06:15:10 +02:00
parent 20435af11b
commit 6d812ceb7f
39 changed files with 716 additions and 781 deletions
Download Patch File Download Diff File Expand all files Collapse all files
elasticsearch.go
+1 -1
View File
@@ -124,6 +124,7 @@ func (p *ESPlugin) RttDurationToMs(d time.Duration) int64 {
return int64(fl)
}
// ResponseAnalyze send req and resp to ES
func (p *ESPlugin) ResponseAnalyze(req, resp []byte, start, stop time.Time) {
if len(resp) == 0 {
// nil http response - skipped elasticsearch export for this request
@@ -131,7 +132,6 @@ func (p *ESPlugin) ResponseAnalyze(req, resp []byte, start, stop time.Time) {
}
t := time.Now()
rtt := p.RttDurationToMs(stop.Sub(start))
req = payloadBody(req)
esResp := ESRequestResponse{
ReqURL: string(proto.Path(req)),
emitter.go
+60 -116
View File
@@ -1,30 +1,30 @@
package main
import (
"bytes"
"fmt"
"hash/fnv"
"io"
"log"
"sync"
"time"
"github.com/buger/goreplay/byteutils"
)
type emitter struct {
// Emitter represents an abject to manage plugins communication
type Emitter struct {
sync.Mutex
sync.WaitGroup
quit chan int
plugins *InOutPlugins
}
// NewEmitter creates and initializes new `emitter` object.
func NewEmitter(quit chan int) *emitter {
return &emitter{
quit: quit,
}
// NewEmitter creates and initializes new Emitter object.
func NewEmitter() *Emitter {
return &Emitter{}
}
// Start initialize loop for sending data from inputs to outputs
func (e *emitter) Start(plugins *InOutPlugins, middlewareCmd string) {
func (e *Emitter) Start(plugins *InOutPlugins, middlewareCmd string) {
defer e.Wait()
if Settings.CopyBufferSize < 1 {
Settings.CopyBufferSize = 5 << 20
@@ -38,147 +38,92 @@ func (e *emitter) Start(plugins *InOutPlugins, middlewareCmd string) {
middleware.ReadFrom(in)
}
// We are going only to read responses, so using same ReadFrom method
for _, out := range plugins.Outputs {
if r, ok := out.(io.Reader); ok {
middleware.ReadFrom(r)
}
}
e.plugins.Inputs = append(e.plugins.Inputs, middleware)
e.plugins.All = append(e.plugins.All, middleware)
e.Add(1)
go func() {
defer e.Done()
if err := CopyMulty(e.quit, middleware, plugins.Outputs...); err != nil {
Debug(2, "Error during copy: ", err)
e.Close()
}
}()
go func() {
for {
select {
case <-e.quit:
middleware.Close()
return
}
if err := CopyMulty(middleware, plugins.Outputs...); err != nil {
Debug(2, fmt.Sprintf("[EMITTER] error during copy: %q", err))
}
}()
} else {
for _, in := range plugins.Inputs {
e.Add(1)
go func(in io.Reader) {
go func(in PluginReader) {
defer e.Done()
if err := CopyMulty(e.quit, in, plugins.Outputs...); err != nil {
Debug(2, "Error during copy: ", err)
e.Close()
if err := CopyMulty(in, plugins.Outputs...); err != nil {
Debug(2, fmt.Sprintf("[EMITTER] error during copy: %q", err))
}
}(in)
}
for _, out := range plugins.Outputs {
if r, ok := out.(io.Reader); ok {
e.Add(1)
go func(r io.Reader) {
defer e.Done()
if err := CopyMulty(e.quit, r, plugins.Outputs...); err != nil {
Debug(2, "Error during copy: ", err)
e.Close()
}
}(r)
}
}
}
}
func (e *emitter) close() {
select {
case <-e.quit:
default:
close(e.quit)
}
}
// Close closes all the goroutine and waits for it to finish.
func (e *emitter) Close() {
e.close()
for _, p := range e.plugins.Inputs {
func (e *Emitter) Close() {
for _, p := range e.plugins.All {
if cp, ok := p.(io.Closer); ok {
cp.Close()
}
}
for _, p := range e.plugins.Outputs {
if cp, ok := p.(io.Closer); ok {
cp.Close()
}
}
e.plugins = nil // avoid further accidental usage
e.plugins.All = nil // avoid further accidental close
}
// CopyMulty copies from 1 reader to multiple writers
func CopyMulty(stop chan int, src io.Reader, writers ...io.Writer) error {
buf := make([]byte, Settings.CopyBufferSize)
func CopyMulty(src PluginReader, writers ...PluginWriter) error {
wIndex := 0
modifier := NewHTTPModifier(&Settings.ModifierConfig)
filteredRequests := make(map[string]time.Time)
filteredRequestsLastCleanTime := time.Now()
filteredRequests := make(map[string]int64)
filteredRequestsLastCleanTime := time.Now().UnixNano()
filteredCount := 0
i := 0
for {
var nr int
nr, err := src.Read(buf)
select {
case <-stop:
return nil
default:
}
msg, err := src.PluginRead()
if err != nil {
if err == ErrorStopped || err == io.EOF {
return nil
}
return err
}
_maxN := nr
if nr > 500 {
_maxN = 500
}
if nr > 0 {
payload := buf[:nr]
meta := payloadMeta(payload)
if msg != nil && len(msg.Data) > 0 {
if len(msg.Data) > int(Settings.CopyBufferSize) {
msg.Data = msg.Data[:Settings.CopyBufferSize]
}
meta := payloadMeta(msg.Meta)
if len(meta) < 3 {
Debug(2, "[EMITTER] Found malformed record", string(payload[0:_maxN]), nr, "from:", src)
Debug(2, fmt.Sprintf("[EMITTER] Found malformed record %q from %q", msg.Meta, src))
continue
}
requestID := string(meta[1])
Debug(3, "[EMITTER] input:", string(payload[0:_maxN]), nr, "from:", src)
requestID := byteutils.SliceToString(meta[1])
// start a subroutine only when necessary
if Settings.Verbose >= 3 {
Debug(3, "[EMITTER] input: ", byteutils.SliceToString(msg.Meta[:len(msg.Meta)-1]), " from: ", src)
}
if modifier != nil {
if isRequestPayload(payload) {
headSize := bytes.IndexByte(payload, '\n') + 1
body := payload[headSize:]
originalBodyLen := len(body)
body = modifier.Rewrite(body)
Debug(3, "[EMITTER] modifier:", requestID, "from:", src)
if isRequestPayload(msg.Meta) {
msg.Data = modifier.Rewrite(msg.Data)
// If modifier tells to skip request
if len(body) == 0 {
filteredRequests[requestID] = time.Now()
if len(msg.Data) == 0 {
filteredRequests[requestID] = time.Now().UnixNano()
filteredCount++
continue
}
if originalBodyLen != len(body) {
payload = append(payload[:headSize], body...)
}
Debug(3, "[EMITTER] Rewritten input:", len(payload), "First %d bytes:", _maxN, string(payload[0:_maxN]))
Debug(3, "[EMITTER] Rewritten input:", requestID, "from:", src)
} else {
if _, ok := filteredRequests[requestID]; ok {
delete(filteredRequests, requestID)
filteredCount--
continue
}
}
}
if Settings.PrettifyHTTP {
payload = prettifyHTTP(payload)
if len(payload) == 0 {
msg.Data = prettifyHTTP(msg.Data)
if len(msg.Data) == 0 {
continue
}
}
@@ -189,15 +134,15 @@ func CopyMulty(stop chan int, src io.Reader, writers ...io.Writer) error {
log.Fatal("Detailed TCP sessions work only with PRO license")
}
hasher := fnv.New32a()
// First 20 bytes contain tcp session
id := payloadID(payload)
hasher.Write(id)
hasher.Write(meta[1])
wIndex = int(hasher.Sum32()) % len(writers)
writers[wIndex].Write(payload)
if _, err := writers[wIndex].PluginWrite(msg); err != nil {
return err
}
} else {
// Simple round robin
if _, err := writers[wIndex].Write(payload); err != nil {
if _, err := writers[wIndex].PluginWrite(msg); err != nil {
return err
}
@@ -209,7 +154,7 @@ func CopyMulty(stop chan int, src io.Reader, writers ...io.Writer) error {
}
} else {
for _, dst := range writers {
if _, err := dst.Write(payload); err != nil {
if _, err := dst.PluginWrite(msg); err != nil {
return err
}
}
@@ -217,19 +162,18 @@ func CopyMulty(stop chan int, src io.Reader, writers ...io.Writer) error {
}
// Run GC on each 1000 request
if i%1000 == 0 {
if filteredCount > 0 && filteredCount%1000 == 0 {
// Clean up filtered requests for which we didn't get a response to filter
now := time.Now()
if now.Sub(filteredRequestsLastCleanTime) > 60*time.Second {
now := time.Now().UnixNano()
if now-filteredRequestsLastCleanTime > int64(60*time.Second) {
for k, v := range filteredRequests {
if now.Sub(v) > 60*time.Second {
if now-v > int64(60*time.Second) {
delete(filteredRequests, k)
filteredCount--
}
}
filteredRequestsLastCleanTime = time.Now()
filteredRequestsLastCleanTime = time.Now().UnixNano()
}
}
i++
}
}
emitter_test.go
+35 -43
View File
@@ -2,7 +2,6 @@ package main
import (
"fmt"
"io"
"os"
"sync"
"sync/atomic"
@@ -18,20 +17,19 @@ func TestMain(m *testing.M) {
func TestEmitter(t *testing.T) {
wg := new(sync.WaitGroup)
quit := make(chan int)
input := NewTestInput()
output := NewTestOutput(func(data []byte) {
output := NewTestOutput(func(*Message) {
wg.Done()
})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
plugins.All = append(plugins.All, input, output)
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
for i := 0; i < 1000; i++ {
@@ -45,32 +43,31 @@ func TestEmitter(t *testing.T) {
func TestEmitterFiltered(t *testing.T) {
wg := new(sync.WaitGroup)
quit := make(chan int)
input := NewTestInput()
input.skipHeader = true
output := NewTestOutput(func(data []byte) {
output := NewTestOutput(func(*Message) {
wg.Done()
})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
plugins.All = append(plugins.All, input, output)
methods := HTTPMethods{[]byte("GET")}
Settings.ModifierConfig = HTTPModifierConfig{Methods: methods}
emitter := &emitter{quit: quit}
emitter := &Emitter{}
go emitter.Start(plugins, "")
wg.Add(2)
id := uuid()
reqh := payloadHeader(RequestPayload, id, time.Now().UnixNano(), -1)
reqb := append(reqh, []byte("GET / HTTP/1.1\r\nHost: www.w3.org\r\nUser-Agent: Go 1.1 package http\r\nAccept-Encoding: gzip\r\n\r\n")...)
reqb := append(reqh, []byte("POST / HTTP/1.1\r\nHost: www.w3.org\r\nUser-Agent: Go 1.1 package http\r\nAccept-Encoding: gzip\r\n\r\n")...)
resh := payloadHeader(ResponsePayload, id, time.Now().UnixNano()+1, 1)
respb := append(resh, []byte("HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n")...)
@@ -80,7 +77,7 @@ func TestEmitterFiltered(t *testing.T) {
id = uuid()
reqh = payloadHeader(RequestPayload, id, time.Now().UnixNano(), -1)
reqb = append(reqh, []byte("POST / HTTP/1.1\r\nHost: www.w3.org\r\nUser-Agent: Go 1.1 package http\r\nAccept-Encoding: gzip\r\n\r\n")...)
reqb = append(reqh, []byte("GET / HTTP/1.1\r\nHost: www.w3.org\r\nUser-Agent: Go 1.1 package http\r\nAccept-Encoding: gzip\r\n\r\n")...)
resh = payloadHeader(ResponsePayload, id, time.Now().UnixNano()+1, 1)
respb = append(resh, []byte("HTTP/1.1 200 OK\r\nContent-Length: 0\r\n\r\n")...)
@@ -96,30 +93,29 @@ func TestEmitterFiltered(t *testing.T) {
func TestEmitterSplitRoundRobin(t *testing.T) {
wg := new(sync.WaitGroup)
quit := make(chan int)
input := NewTestInput()
var counter1, counter2 int32
output1 := NewTestOutput(func(data []byte) {
output1 := NewTestOutput(func(*Message) {
atomic.AddInt32(&counter1, 1)
wg.Done()
})
output2 := NewTestOutput(func(data []byte) {
output2 := NewTestOutput(func(*Message) {
atomic.AddInt32(&counter2, 1)
wg.Done()
})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output1, output2},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output1, output2},
}
Settings.SplitOutput = true
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
for i := 0; i < 1000; i++ {
@@ -140,31 +136,30 @@ func TestEmitterSplitRoundRobin(t *testing.T) {
func TestEmitterRoundRobin(t *testing.T) {
wg := new(sync.WaitGroup)
quit := make(chan int)
input := NewTestInput()
var counter1, counter2 int32
output1 := NewTestOutput(func(data []byte) {
atomic.AddInt32(&counter1, 1)
output1 := NewTestOutput(func(*Message) {
counter1++
wg.Done()
})
output2 := NewTestOutput(func(data []byte) {
atomic.AddInt32(&counter2, 1)
output2 := NewTestOutput(func(*Message) {
counter2++
wg.Done()
})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output1, output2},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output1, output2},
}
plugins.All = append(plugins.All, input, output1, output2)
Settings.SplitOutput = true
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
for i := 0; i < 1000; i++ {
@@ -186,36 +181,34 @@ func TestEmitterSplitSession(t *testing.T) {
wg := new(sync.WaitGroup)
wg.Add(200)
quit := make(chan int)
input := NewTestInput()
input.skipHeader = true
var counter1, counter2 int32
output1 := NewTestOutput(func(data []byte) {
if payloadID(data)[0] == 'a' {
atomic.AddInt32(&counter1, 1)
output1 := NewTestOutput(func(msg *Message) {
if payloadID(msg.Meta)[0] == 'a' {
counter1++
}
wg.Done()
})
output2 := NewTestOutput(func(data []byte) {
if payloadID(data)[0] == 'b' {
atomic.AddInt32(&counter2, 1)
output2 := NewTestOutput(func(msg *Message) {
if payloadID(msg.Meta)[0] == 'b' {
counter2++
}
wg.Done()
})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output1, output2},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output1, output2},
}
Settings.SplitOutput = true
Settings.RecognizeTCPSessions = true
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
for i := 0; i < 200; i++ {
@@ -242,21 +235,20 @@ func TestEmitterSplitSession(t *testing.T) {
func BenchmarkEmitter(b *testing.B) {
wg := new(sync.WaitGroup)
quit := make(chan int)
input := NewTestInput()
output := NewTestOutput(func(data []byte) {
output := NewTestOutput(func(*Message) {
wg.Done()
})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
plugins.All = append(plugins.All, input, output)
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
b.ResetTimer()
gor.go
+1 -1
View File
@@ -74,7 +74,7 @@ func main() {
}
closeCh := make(chan int)
emitter := NewEmitter(closeCh)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
if Settings.ExitAfter > 0 {
log.Printf("Running gor for a duration of %s\n", Settings.ExitAfter)
http_modifier_settings.go
+30 -12
View File
@@ -19,10 +19,9 @@ type HTTPModifierConfig struct {
HeaderBasicAuthFilters HTTPHeaderBasicAuthFilters `json:"http-basic-auth-filter"`
HeaderHashFilters HTTPHashFilters `json:"http-header-limiter"`
ParamHashFilters HTTPHashFilters `json:"http-param-limiter"`
Params HTTPParams `json:"http-set-param"`
Headers HTTPHeaders `json:"http-set-header"`
Methods HTTPMethods `json:"http-allow-method"`
Params HTTPParams `json:"http-set-param"`
Headers HTTPHeaders `json:"http-set-header"`
Methods HTTPMethods `json:"http-allow-method"`
}
//
@@ -40,10 +39,11 @@ func (h *HTTPHeaderFilters) String() string {
return fmt.Sprint(*h)
}
// Set method to implement flags.Value
func (h *HTTPHeaderFilters) Set(value string) error {
valArr := strings.SplitN(value, ":", 2)
if len(valArr) < 2 {
return errors.New("need both header and value, colon-delimited (ex. user_id:^169$).")
return errors.New("need both header and value, colon-delimited (ex. user_id:^169$)")
}
val := strings.TrimSpace(valArr[1])
r, err := regexp.Compile(val)
@@ -63,13 +63,14 @@ type basicAuthFilter struct {
regexp *regexp.Regexp
}
// HTTPHeaderFilters holds list of headers and their regexps
// HTTPHeaderBasicAuthFilters holds list of regxp to match basic Auth header values
type HTTPHeaderBasicAuthFilters []basicAuthFilter
func (h *HTTPHeaderBasicAuthFilters) String() string {
return fmt.Sprint(*h)
}
// Set method to implement flags.Value
func (h *HTTPHeaderBasicAuthFilters) Set(value string) error {
r, err := regexp.Compile(value)
if err != nil {
@@ -89,12 +90,14 @@ type hashFilter struct {
percent uint32
}
// HTTPHashFilters represents a slice of header hash filters
type HTTPHashFilters []hashFilter
func (h *HTTPHashFilters) String() string {
return fmt.Sprint(*h)
}
// Set method to implement flags.Value
func (h *HTTPHashFilters) Set(value string) error {
valArr := strings.SplitN(value, ":", 2)
if len(valArr) < 2 {
@@ -129,23 +132,26 @@ func (h *HTTPHashFilters) Set(value string) error {
//
// Handling of --http-set-header option
//
type HTTPHeaders []HTTPHeader
type HTTPHeader struct {
type httpHeader struct {
Name string
Value string
}
// HTTPHeaders is a slice of headers that must appended
type HTTPHeaders []httpHeader
func (h *HTTPHeaders) String() string {
return fmt.Sprint(*h)
}
// Set method to implement flags.Value
func (h *HTTPHeaders) Set(value string) error {
v := strings.SplitN(value, ":", 2)
if len(v) != 2 {
return errors.New("Expected `Key: Value`")
}
header := HTTPHeader{
header := httpHeader{
strings.TrimSpace(v[0]),
strings.TrimSpace(v[1]),
}
@@ -157,23 +163,26 @@ func (h *HTTPHeaders) Set(value string) error {
//
// Handling of --http-set-param option
//
type HTTPParams []HTTPParam
type HTTPParam struct {
type httpParam struct {
Name []byte
Value []byte
}
// HTTPParams filters for --http-set-param
type HTTPParams []httpParam
func (h *HTTPParams) String() string {
return fmt.Sprint(*h)
}
// Set method to implement flags.Value
func (h *HTTPParams) Set(value string) error {
v := strings.SplitN(value, "=", 2)
if len(v) != 2 {
return errors.New("Expected `Key=Value`")
}
param := HTTPParam{
param := httpParam{
[]byte(strings.TrimSpace(v[0])),
[]byte(strings.TrimSpace(v[1])),
}
@@ -185,12 +194,15 @@ func (h *HTTPParams) Set(value string) error {
//
// Handling of --http-allow-method option
//
// HTTPMethods holds values for method allowed
type HTTPMethods [][]byte
func (h *HTTPMethods) String() string {
return fmt.Sprint(*h)
}
// Set method to implement flags.Value
func (h *HTTPMethods) Set(value string) error {
*h = append(*h, []byte(value))
return nil
@@ -204,12 +216,14 @@ type urlRewrite struct {
target []byte
}
// URLRewriteMap holds regexp and data to modify URL
type URLRewriteMap []urlRewrite
func (r *URLRewriteMap) String() string {
return fmt.Sprint(*r)
}
// Set method to implement flags.Value
func (r *URLRewriteMap) Set(value string) error {
valArr := strings.SplitN(value, ":", 2)
if len(valArr) < 2 {
@@ -232,12 +246,14 @@ type headerRewrite struct {
target []byte
}
// HeaderRewriteMap holds regexp and data to rewrite headers
type HeaderRewriteMap []headerRewrite
func (r *HeaderRewriteMap) String() string {
return fmt.Sprint(*r)
}
// Set method to implement flags.Value
func (r *HeaderRewriteMap) Set(value string) error {
headerArr := strings.SplitN(value, ":", 2)
if len(headerArr) < 2 {
@@ -266,12 +282,14 @@ type urlRegexp struct {
regexp *regexp.Regexp
}
// HTTPURLRegexp a slice of regexp to match URLs
type HTTPURLRegexp []urlRegexp
func (r *HTTPURLRegexp) String() string {
return fmt.Sprint(*r)
}
// Set method to implement flags.Value
func (r *HTTPURLRegexp) Set(value string) error {
regexp, err := regexp.Compile(value)
http_prettifier.go
+7 -10
View File
@@ -12,25 +12,22 @@ import (
)
func prettifyHTTP(p []byte) []byte {
headSize := bytes.IndexByte(p, '\n') + 1
head := p[:headSize]
body := p[headSize:]
tEnc := bytes.Equal(proto.Header(body, []byte("Transfer-Encoding")), []byte("chunked"))
cEnc := bytes.Equal(proto.Header(body, []byte("Content-Encoding")), []byte("gzip"))
tEnc := bytes.Equal(proto.Header(p, []byte("Transfer-Encoding")), []byte("chunked"))
cEnc := bytes.Equal(proto.Header(p, []byte("Content-Encoding")), []byte("gzip"))
if !(tEnc || cEnc) {
return p
}
headersPos := proto.MIMEHeadersEndPos(body)
headersPos := proto.MIMEHeadersEndPos(p)
if headersPos < 5 || headersPos > len(body) {
if headersPos < 5 || headersPos > len(p) {
return p
}
headers := body[:headersPos]
content := body[headersPos:]
headers := p[:headersPos]
content := p[headersPos:]
if tEnc {
buf := bytes.NewReader(content)
@@ -64,7 +61,7 @@ func prettifyHTTP(p []byte) []byte {
headers = proto.SetHeader(headers, []byte("Content-Length"), []byte(newLen))
}
newPayload := append(append(head, headers...), content...)
newPayload := append(headers, content...)
return newPayload
}
http_prettifier_test.go
+8 -7
View File
@@ -5,6 +5,8 @@ import (
"compress/gzip"
"strconv"
"testing"
"github.com/buger/goreplay/proto"
)
func TestHTTPPrettifierGzip(t *testing.T) {
@@ -15,22 +17,21 @@ func TestHTTPPrettifierGzip(t *testing.T) {
size := strconv.Itoa(len(b.Bytes()))
payload := []byte("2 1 1\nHTTP/1.1 200 OK\r\nContent-Length: " + size + "\r\nContent-Encoding: gzip\r\n\r\n")
payload := []byte("HTTP/1.1 200 OK\r\nContent-Length: " + size + "\r\nContent-Encoding: gzip\r\n\r\n")
payload = append(payload, b.Bytes()...)
newPayload := prettifyHTTP(payload)
if string(newPayload) != "2 1 1\nHTTP/1.1 200 OK\r\nContent-Length: 4\r\n\r\ntest" {
t.Error("Payload not match:", string(newPayload))
if string(newPayload) != "HTTP/1.1 200 OK\r\nContent-Length: 4\r\n\r\ntest" {
t.Errorf("Payload not match %q", string(newPayload))
}
}
func TestHTTPPrettifierChunked(t *testing.T) {
payload := []byte("POST / HTTP/1.1\r\nHost: www.w3.org\r\nTransfer-Encoding: chunked\r\n\r\n4\r\nWiki\r\n5\r\npedia\r\ne\r\n in\r\n\r\nchunks.\r\n0\r\n\r\n")
newPayload := prettifyHTTP(payload)
if string(newPayload) != "POST / HTTP/1.1\r\nHost: www.w3.org\r\nContent-Length: 23\r\n\r\nWikipedia in\r\n\r\nchunks." {
t.Error("Payload not match:", string(newPayload))
payload = prettifyHTTP(payload)
if string(proto.Header(payload, []byte("Content-Length"))) != "23" {
t.Errorf("payload should have content length of 23")
}
}
input_dummy.go
+18 -6
View File
@@ -7,24 +7,30 @@ import (
// DummyInput used for debugging. It generate 1 "GET /"" request per second.
type DummyInput struct {
data chan []byte
quit chan struct{}
}
// NewDummyInput constructor for DummyInput
func NewDummyInput(options string) (di *DummyInput) {
di = new(DummyInput)
di.data = make(chan []byte)
di.quit = make(chan struct{})
go di.emit()
return
}
func (i *DummyInput) Read(data []byte) (int, error) {
buf := <-i.data
copy(data, buf)
return len(buf), nil
// PluginRead reads message from this plugin
func (i *DummyInput) PluginRead() (*Message, error) {
var msg Message
select {
case <-i.quit:
return nil, ErrorStopped
case buf := <-i.data:
msg.Meta, msg.Data = payloadMetaWithBody(buf)
return &msg, nil
}
}
func (i *DummyInput) emit() {
@@ -46,3 +52,9 @@ func (i *DummyInput) emit() {
func (i *DummyInput) String() string {
return "Dummy Input"
}
// Close closes this plugins
func (i *DummyInput) Close() error {
close(i.quit)
return nil
}
input_file.go
+18 -15
View File
@@ -5,8 +5,8 @@ import (
"bytes"
"compress/gzip"
"errors"
"fmt"
"io"
"log"
"os"
"path/filepath"
"strconv"
@@ -64,6 +64,8 @@ func (f *fileInputReader) ReadPayload() []byte {
return f.data
}
// Close closes this plugin
func (f *fileInputReader) Close() error {
if atomic.LoadInt32(&f.closed) == 0 {
atomic.StoreInt32(&f.closed, 1)
@@ -73,7 +75,7 @@ func (f *fileInputReader) Close() error {
return nil
}
func NewFileInputReader(path string) *fileInputReader {
func newFileInputReader(path string) *fileInputReader {
var file io.ReadCloser
var err error
@@ -84,7 +86,7 @@ func NewFileInputReader(path string) *fileInputReader {
}
if err != nil {
log.Println(err)
Debug(0, fmt.Sprintf("[INPUT-FILE] err: %q", err))
return nil
}
@@ -92,7 +94,7 @@ func NewFileInputReader(path string) *fileInputReader {
if strings.HasSuffix(path, ".gz") {
gzReader, err := gzip.NewReader(file)
if err != nil {
log.Println(err)
Debug(0, fmt.Sprintf("[INPUT-FILE] err: %q", err))
return nil
}
r.reader = bufio.NewReader(gzReader)
@@ -153,7 +155,7 @@ func (i *FileInput) init() (err error) {
resp, err := svc.ListObjects(params)
if err != nil {
log.Println("Error while retreiving list of files from S3", i.path, err)
Debug(0, "[INPUT-FILE] Error while retreiving list of files from S3", i.path, err)
return err
}
@@ -162,34 +164,35 @@ func (i *FileInput) init() (err error) {
}
} else {
if matches, err = filepath.Glob(i.path); err != nil {
log.Println("Wrong file pattern", i.path, err)
Debug(0, "[INPUT-FILE] Wrong file pattern", i.path, err)
return
}
}
if len(matches) == 0 {
log.Println("No files match pattern: ", i.path)
Debug(0, "[INPUT-FILE] No files match pattern: ", i.path)
return errors.New("No matching files")
}
i.readers = make([]*fileInputReader, len(matches))
for idx, p := range matches {
i.readers[idx] = NewFileInputReader(p)
i.readers[idx] = newFileInputReader(p)
}
return nil
}
func (i *FileInput) Read(data []byte) (int, error) {
var buf []byte
// PluginRead reads message from this plugin
func (i *FileInput) PluginRead() (*Message, error) {
var msg Message
select {
case <-i.exit:
return 0, ErrorStopped
case buf = <-i.data:
return nil, ErrorStopped
case buf := <-i.data:
msg.Meta, msg.Data = payloadMetaWithBody(buf)
return &msg, nil
}
n := copy(data, buf)
return n, nil
}
func (i *FileInput) String() string {
@@ -256,7 +259,7 @@ func (i *FileInput) emit() {
}
}
log.Printf("FileInput: end of file '%s'\n", i.path)
Debug(0, fmt.Sprintf("[INPUT-FILE] FileInput: end of file '%s'\n", i.path))
}
input_file_test.go
+48 -65
View File
@@ -4,9 +4,7 @@ import (
"bytes"
"errors"
"fmt"
"io"
"io/ioutil"
"log"
"math/rand"
"os"
"sync"
@@ -14,20 +12,18 @@ import (
"time"
)
var _ = log.Println
func TestInputFileWithGET(t *testing.T) {
input := NewTestInput()
rg := NewRequestGenerator([]io.Reader{input}, func() { input.EmitGET() }, 1)
readPayloads := [][]byte{}
rg := NewRequestGenerator([]PluginReader{input}, func() { input.EmitGET() }, 1)
readPayloads := []*Message{}
// Given a capture file with a GET request
expectedCaptureFile := CreateCaptureFile(rg)
defer expectedCaptureFile.TearDown()
// When the request is read from the capture file
err := ReadFromCaptureFile(expectedCaptureFile.file, 1, func(data []byte) {
readPayloads = append(readPayloads, Duplicate(data))
err := ReadFromCaptureFile(expectedCaptureFile.file, 1, func(msg *Message) {
readPayloads = append(readPayloads, msg)
})
// The read request should match the original request
@@ -41,18 +37,17 @@ func TestInputFileWithGET(t *testing.T) {
}
func TestInputFileWithPayloadLargerThan64Kb(t *testing.T) {
input := NewTestInput()
rg := NewRequestGenerator([]io.Reader{input}, func() { input.EmitSizedPOST(64 * 1024) }, 1)
readPayloads := [][]byte{}
rg := NewRequestGenerator([]PluginReader{input}, func() { input.EmitSizedPOST(64 * 1024) }, 1)
readPayloads := []*Message{}
// Given a capture file with a request over 64Kb
expectedCaptureFile := CreateCaptureFile(rg)
defer expectedCaptureFile.TearDown()
// When the request is read from the capture file
err := ReadFromCaptureFile(expectedCaptureFile.file, 1, func(data []byte) {
readPayloads = append(readPayloads, Duplicate(data))
err := ReadFromCaptureFile(expectedCaptureFile.file, 1, func(msg *Message) {
readPayloads = append(readPayloads, msg)
})
// The read request should match the original request
@@ -69,19 +64,19 @@ func TestInputFileWithPayloadLargerThan64Kb(t *testing.T) {
func TestInputFileWithGETAndPOST(t *testing.T) {
input := NewTestInput()
rg := NewRequestGenerator([]io.Reader{input}, func() {
rg := NewRequestGenerator([]PluginReader{input}, func() {
input.EmitGET()
input.EmitPOST()
}, 2)
readPayloads := [][]byte{}
readPayloads := []*Message{}
// Given a capture file with a GET request
expectedCaptureFile := CreateCaptureFile(rg)
defer expectedCaptureFile.TearDown()
// When the requests are read from the capture file
err := ReadFromCaptureFile(expectedCaptureFile.file, 2, func(data []byte) {
readPayloads = append(readPayloads, Duplicate(data))
err := ReadFromCaptureFile(expectedCaptureFile.file, 2, func(msg *Message) {
readPayloads = append(readPayloads, msg)
})
// The read requests should match the original request
@@ -113,12 +108,11 @@ func TestInputFileMultipleFilesWithRequestsOnly(t *testing.T) {
file2.Close()
input := NewFileInput(fmt.Sprintf("/tmp/%d*", rnd), false)
buf := make([]byte, 1000)
for i := '1'; i <= '4'; i++ {
n, _ := input.Read(buf)
if buf[4] != byte(i) {
t.Error("Should emit requests in right order", string(buf[:n]))
msg, _ := input.PluginRead()
if msg.Meta[4] != byte(i) {
t.Error("Should emit requests in right order", string(msg.Meta))
}
}
@@ -140,11 +134,10 @@ func TestInputFileRequestsWithLatency(t *testing.T) {
file.Write([]byte(payloadSeparator))
input := NewFileInput(fmt.Sprintf("/tmp/%d", rnd), false)
buf := make([]byte, 1000)
start := time.Now().UnixNano()
for i := 0; i < 3; i++ {
input.Read(buf)
input.PluginRead()
}
end := time.Now().UnixNano()
@@ -181,17 +174,16 @@ func TestInputFileMultipleFilesWithRequestsAndResponses(t *testing.T) {
file2.Close()
input := NewFileInput(fmt.Sprintf("/tmp/%d*", rnd), false)
buf := make([]byte, 1000)
for i := '1'; i <= '4'; i++ {
n, _ := input.Read(buf)
if buf[0] != '1' && buf[4] != byte(i) {
t.Error("Shound emit requests in right order", string(buf[:n]))
msg, _ := input.PluginRead()
if msg.Meta[0] != '1' && msg.Meta[4] != byte(i) {
t.Error("Shound emit requests in right order", string(msg.Meta))
}
n, _ = input.Read(buf)
if buf[0] != '2' && buf[4] != byte(i) {
t.Error("Shound emit responses in right order", string(buf[:n]))
msg, _ = input.PluginRead()
if msg.Meta[0] != '2' && msg.Meta[4] != byte(i) {
t.Error("Shound emit responses in right order", string(msg.Meta))
}
}
@@ -210,11 +202,10 @@ func TestInputFileLoop(t *testing.T) {
file.Close()
input := NewFileInput(fmt.Sprintf("/tmp/%d", rnd), true)
buf := make([]byte, 1000)
// Even if we have just 2 requests in file, it should indifinitly loop
for i := 0; i < 1000; i++ {
input.Read(buf)
input.PluginRead()
}
input.Close()
@@ -226,22 +217,21 @@ func TestInputFileCompressed(t *testing.T) {
output := NewFileOutput(fmt.Sprintf("/tmp/%d_0.gz", rnd), &FileOutputConfig{FlushInterval: time.Minute, Append: true})
for i := 0; i < 1000; i++ {
output.Write([]byte("1 1 1\r\ntest"))
output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")})
}
name1 := output.file.Name()
output.Close()
output2 := NewFileOutput(fmt.Sprintf("/tmp/%d_1.gz", rnd), &FileOutputConfig{FlushInterval: time.Minute, Append: true})
for i := 0; i < 1000; i++ {
output2.Write([]byte("1 1 1\r\ntest"))
output2.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")})
}
name2 := output2.file.Name()
output2.Close()
input := NewFileInput(fmt.Sprintf("/tmp/%d*", rnd), false)
buf := make([]byte, 1000)
for i := 0; i < 2000; i++ {
input.Read(buf)
input.PluginRead()
}
os.Remove(name1)
@@ -249,14 +239,14 @@ func TestInputFileCompressed(t *testing.T) {
}
type CaptureFile struct {
data [][]byte
msgs []*Message
file *os.File
}
func NewExpectedCaptureFile(data [][]byte, file *os.File) *CaptureFile {
func NewExpectedCaptureFile(msgs []*Message, file *os.File) *CaptureFile {
ecf := new(CaptureFile)
ecf.file = file
ecf.data = data
ecf.msgs = msgs
return ecf
}
@@ -267,12 +257,12 @@ func (expectedCaptureFile *CaptureFile) TearDown() {
}
type RequestGenerator struct {
inputs []io.Reader
inputs []PluginReader
emit func()
wg *sync.WaitGroup
}
func NewRequestGenerator(inputs []io.Reader, emit func(), count int) (rg *RequestGenerator) {
func NewRequestGenerator(inputs []PluginReader, emit func(), count int) (rg *RequestGenerator) {
rg = new(RequestGenerator)
rg.inputs = inputs
rg.emit = emit
@@ -281,14 +271,17 @@ func NewRequestGenerator(inputs []io.Reader, emit func(), count int) (rg *Reques
return
}
func (expectedCaptureFile *CaptureFile) PayloadsEqual(other [][]byte) bool {
func (expectedCaptureFile *CaptureFile) PayloadsEqual(other []*Message) bool {
if len(expectedCaptureFile.data) != len(other) {
if len(expectedCaptureFile.msgs) != len(other) {
return false
}
for i, payload := range other {
if !bytes.Equal(expectedCaptureFile.data[i], payload) {
if !bytes.Equal(expectedCaptureFile.msgs[i].Meta, payload.Meta) {
return false
}
if !bytes.Equal(expectedCaptureFile.msgs[i].Data, payload.Data) {
return false
}
}
@@ -303,26 +296,24 @@ func CreateCaptureFile(requestGenerator *RequestGenerator) *CaptureFile {
panic(err)
}
quit := make(chan int)
readPayloads := [][]byte{}
output := NewTestOutput(func(data []byte) {
readPayloads = append(readPayloads, Duplicate(data))
readPayloads := []*Message{}
output := NewTestOutput(func(msg *Message) {
readPayloads = append(readPayloads, msg)
requestGenerator.wg.Done()
})
outputFile := NewFileOutput(f.Name(), &FileOutputConfig{FlushInterval: time.Minute, Append: true})
outputFile := NewFileOutput(f.Name(), &FileOutputConfig{FlushInterval: time.Second, Append: true})
plugins := &InOutPlugins{
Inputs: requestGenerator.inputs,
Outputs: []io.Writer{output, outputFile},
Outputs: []PluginWriter{output, outputFile},
}
for _, input := range requestGenerator.inputs {
plugins.All = append(plugins.All, input)
}
plugins.All = append(plugins.All, output, outputFile)
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
requestGenerator.emit()
@@ -336,23 +327,22 @@ func CreateCaptureFile(requestGenerator *RequestGenerator) *CaptureFile {
}
func ReadFromCaptureFile(captureFile *os.File, count int, callback writeCallback) (err error) {
quit := make(chan int)
wg := new(sync.WaitGroup)
input := NewFileInput(captureFile.Name(), false)
output := NewTestOutput(func(data []byte) {
callback(data)
output := NewTestOutput(func(msg *Message) {
callback(msg)
wg.Done()
})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
plugins.All = append(plugins.All, input, output)
wg.Add(count)
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
done := make(chan int, 1)
@@ -371,10 +361,3 @@ func ReadFromCaptureFile(captureFile *os.File, count int, callback writeCallback
return
}
func Duplicate(data []byte) (duplicate []byte) {
duplicate = make([]byte, len(data))
copy(duplicate, data)
return
}
input_http.go
+8 -16
View File
@@ -27,25 +27,17 @@ func NewHTTPInput(address string) (i *HTTPInput) {
return
}
func (i *HTTPInput) Read(data []byte) (int, error) {
var buf []byte
// PluginRead reads message from this plugin
func (i *HTTPInput) PluginRead() (*Message, error) {
var msg Message
select {
case <-i.stop:
return 0, ErrorStopped
case buf = <-i.data:
return nil, ErrorStopped
case buf := <-i.data:
msg.Data = buf
msg.Meta = payloadHeader(RequestPayload, uuid(), time.Now().UnixNano(), -1)
return &msg, nil
}
header := payloadHeader(RequestPayload, uuid(), time.Now().UnixNano(), -1)
n := copy(data, header)
if len(data) > len(header) {
n += copy(data[len(header):], buf)
}
dis := len(header) + len(buf) - n
if dis > 0 {
Debug(2, "[INPUT-HTTP] discarded", dis, "increase copy buffer size")
}
return n, nil
}
// Close closes this plugin
input_http_test.go
+9 -15
View File
@@ -2,33 +2,29 @@ package main
import (
"bytes"
"io"
"net/http"
"strings"
"sync"
"testing"
"time"
"github.com/buger/goreplay/proto"
)
func TestHTTPInput(t *testing.T) {
wg := new(sync.WaitGroup)
quit := make(chan int)
input := NewHTTPInput("127.0.0.1:0")
time.Sleep(time.Millisecond)
output := NewTestOutput(func(data []byte) {
output := NewTestOutput(func(*Message) {
wg.Done()
})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
plugins.All = append(plugins.All, input, output)
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
address := strings.Replace(input.address, "[::]", "127.0.0.1", -1)
@@ -44,27 +40,25 @@ func TestHTTPInput(t *testing.T) {
func TestInputHTTPLargePayload(t *testing.T) {
wg := new(sync.WaitGroup)
quit := make(chan int, 1)
const n = 10 << 20 // 10MB
var large [n]byte
large[n-1] = '0'
input := NewHTTPInput("127.0.0.1:0")
time.Sleep(time.Millisecond)
output := NewTestOutput(func(data []byte) {
_len := len(proto.Body(payloadBody(data)))
output := NewTestOutput(func(msg *Message) {
_len := len(msg.Data)
if _len >= n { // considering http body CRLF
t.Errorf("expected body to be >= %d", n)
}
wg.Done()
})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
plugins.All = append(plugins.All, input, output)
emitter := NewEmitter(quit)
emitter := NewEmitter()
defer emitter.Close()
go emitter.Start(plugins, Settings.Middleware)
input_kafka.go
+32 -14
View File
@@ -15,6 +15,7 @@ type KafkaInput struct {
config *InputKafkaConfig
consumers []sarama.PartitionConsumer
messages chan *sarama.ConsumerMessage
quit chan struct{}
}
// NewKafkaInput creates instance of kafka consumer client with TLS config
@@ -43,6 +44,7 @@ func NewKafkaInput(address string, config *InputKafkaConfig, tlsConfig *KafkaTLS
config: config,
consumers: make([]sarama.PartitionConsumer, len(partitions)),
messages: make(chan *sarama.ConsumerMessage, 256),
quit: make(chan struct{}),
}
for index, partition := range partitions {
@@ -74,29 +76,45 @@ func (i *KafkaInput) ErrorHandler(consumer sarama.PartitionConsumer) {
}
}
func (i *KafkaInput) Read(data []byte) (int, error) {
message := <-i.messages
if !i.config.UseJSON {
copy(data, message.Value)
return len(message.Value), nil
// PluginRead a reads message from this plugin
func (i *KafkaInput) PluginRead() (*Message, error) {
var message *sarama.ConsumerMessage
var msg Message
select {
case <-i.quit:
return nil, ErrorStopped
case message = <-i.messages:
}
var kafkaMessage KafkaMessage
json.Unmarshal(message.Value, &kafkaMessage)
msg.Data = message.Value
if i.config.UseJSON {
buf, err := kafkaMessage.Dump()
if err != nil {
Debug(1, "Failed to decode access log entry:", err)
return 0, err
var kafkaMessage KafkaMessage
json.Unmarshal(message.Value, &kafkaMessage)
var err error
msg.Data, err = kafkaMessage.Dump()
if err != nil {
Debug(1, "[INPUT-KAFKA] failed to decode access log entry:", err)
return nil, err
}
}
n := copy(data, buf)
// does it have meta
if isOriginPayload(msg.Data) {
msg.Meta, msg.Data = payloadMetaWithBody(msg.Data)
}
return n, nil
return &msg, nil
}
func (i *KafkaInput) String() string {
return "Kafka Input: " + i.config.Host + "/" + i.config.Topic
}
// Close closes this plugin
func (i *KafkaInput) Close() error {
close(i.quit)
return nil
}
input_kafka_test.go
+8 -10
View File
@@ -20,17 +20,16 @@ func TestInputKafkaRAW(t *testing.T) {
consumer: consumer,
Topic: "test",
UseJSON: false,
},nil)
}, nil)
buf := make([]byte, 1024)
n, err := input.Read(buf)
msg, err := input.PluginRead()
if err != nil {
t.Fatal(err)
}
if string(buf[:n]) != "1 2 3\nGET / HTTP1.1\r\nHeader: 1\r\n\r\n" {
t.Error("Message not properly decoded: ", string(buf[:n]), n)
if string(append(msg.Meta, msg.Data...)) != "1 2 3\nGET / HTTP1.1\r\nHeader: 1\r\n\r\n" {
t.Error("Message not properly decoded")
}
}
@@ -47,16 +46,15 @@ func TestInputKafkaJSON(t *testing.T) {
consumer: consumer,
Topic: "test",
UseJSON: true,
},nil)
}, nil)
buf := make([]byte, 1024)
n, err := input.Read(buf)
msg, err := input.PluginRead()
if err != nil {
t.Fatal(err)
}
if string(buf[:n]) != "1 2 3\nGET / HTTP/1.1\r\nHeader: 1\r\n\r\n" {
t.Error("Message not properly decoded: ", string(buf[:n]), n)
if string(append(msg.Meta, msg.Data...)) != "1 2 3\nGET / HTTP/1.1\r\nHeader: 1\r\n\r\n" {
t.Error("Message not properly decoded")
}
}
input_raw.go
+17 -22
View File
@@ -102,44 +102,39 @@ func NewRAWInput(address string, config RAWInputConfig) (i *RAWInput) {
return
}
func (i *RAWInput) Read(data []byte) (n int, err error) {
var msg *tcp.Message
var buf []byte
// PluginRead reads meassage from this plugin
func (i *RAWInput) PluginRead() (*Message, error) {
var msgTCP *tcp.Message
var msg Message
select {
case <-i.quit:
return 0, ErrorStopped
case msg = <-i.message:
buf = msg.Data()
return nil, ErrorStopped
case msgTCP = <-i.message:
msg.Data = msgTCP.Data()
}
var header []byte
var msgType byte = ResponsePayload
if msg.IsIncoming {
if msgTCP.IsIncoming {
msgType = RequestPayload
if i.RealIPHeader != "" {
buf = proto.SetHeader(buf, []byte(i.RealIPHeader), []byte(msg.SrcAddr))
msg.Data = proto.SetHeader(msg.Data, []byte(i.RealIPHeader), []byte(msgTCP.SrcAddr))
}
}
header = payloadHeader(msgType, msg.UUID(), msg.Start.UnixNano(), msg.End.UnixNano()-msg.Start.UnixNano())
msg.Meta = payloadHeader(msgType, msgTCP.UUID(), msgTCP.Start.UnixNano(), msgTCP.End.UnixNano()-msgTCP.Start.UnixNano())
n = copy(data, header)
if len(data) > len(header) {
n += copy(data[len(header):], buf)
}
// to be removed....
if msg.Truncated || len(header)+len(buf)-n > 0 {
go Debug(2, "[INPUT-RAW] message truncated, increase copy-buffer-size")
if msgTCP.Truncated {
Debug(2, "[INPUT-RAW] message truncated, increase copy-buffer-size")
}
// to be removed...
if msg.TimedOut {
go Debug(2, "[INPUT-RAW] message timeout reached, increase input-raw-expire")
if msgTCP.TimedOut {
Debug(2, "[INPUT-RAW] message timeout reached, increase input-raw-expire")
}
if i.Stats {
stat := msg.Stats
stat := msgTCP.Stats
go i.addStats(stat)
}
msg = nil
return n, nil
msgTCP = nil
return &msg, nil
}
func (i *RAWInput) listen(address string) {
input_raw_test.go
+45 -48
View File
@@ -2,7 +2,6 @@ package main
import (
"bytes"
"io"
"io/ioutil"
"net"
"net/http"
@@ -23,9 +22,8 @@ const testRawExpire = time.Millisecond * 200
func TestRAWInputIPv4(t *testing.T) {
wg := new(sync.WaitGroup)
quit := make(chan int)
listener, err := net.Listen("tcp", ":0")
listener, err := net.Listen("tcp4", "127.0.0.1:0")
if err != nil {
t.Error(err)
return
@@ -49,32 +47,31 @@ func TestRAWInputIPv4(t *testing.T) {
TrackResponse: true,
RealIPHeader: "X-Real-IP",
}
input := NewRAWInput(":"+port, conf)
input := NewRAWInput(listener.Addr().String(), conf)
output := NewTestOutput(func(data []byte) {
if data[0] == '1' {
body := payloadBody(data)
if len(proto.Header(body, []byte("X-Real-IP"))) == 0 {
t.Error("Should have X-Real-IP header", string(body))
output := NewTestOutput(func(msg *Message) {
if msg.Meta[0] == '1' {
if len(proto.Header(msg.Data, []byte("X-Real-IP"))) == 0 {
t.Error("Should have X-Real-IP header")
}
atomic.AddInt64(&reqCounter, 1)
reqCounter++
} else {
atomic.AddInt64(&respCounter, 1)
respCounter++
}
wg.Done()
})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
plugins.All = append(plugins.All, input, output)
addr := "http://127.0.0.1:" + port
emitter := NewEmitter(quit)
emitter := NewEmitter()
defer emitter.Close()
go emitter.Start(plugins, Settings.Middleware)
for i := 0; i < 10; i++ {
for i := 0; i < 1; i++ {
wg.Add(2)
_, err = http.Get(addr)
if err != nil {
@@ -91,7 +88,6 @@ func TestRAWInputIPv4(t *testing.T) {
func TestRAWInputNoKeepAlive(t *testing.T) {
wg := new(sync.WaitGroup)
quit := make(chan int)
listener, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
@@ -117,8 +113,8 @@ func TestRAWInputNoKeepAlive(t *testing.T) {
}
input := NewRAWInput(":"+port, conf)
var respCounter, reqCounter int64
output := NewTestOutput(func(data []byte) {
if data[0] == '1' {
output := NewTestOutput(func(msg *Message) {
if msg.Meta[0] == '1' {
atomic.AddInt64(&reqCounter, 1)
} else {
atomic.AddInt64(&respCounter, 1)
@@ -127,14 +123,14 @@ func TestRAWInputNoKeepAlive(t *testing.T) {
})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
plugins.All = append(plugins.All, input, output)
addr := "http://127.0.0.1:" + port
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
for i := 0; i < 10; i++ {
@@ -157,7 +153,6 @@ func TestRAWInputNoKeepAlive(t *testing.T) {
func TestRAWInputIPv6(t *testing.T) {
wg := new(sync.WaitGroup)
quit := make(chan int)
listener, err := net.Listen("tcp", "[::1]:0")
if err != nil {
@@ -183,8 +178,8 @@ func TestRAWInputIPv6(t *testing.T) {
}
input := NewRAWInput(originAddr, conf)
output := NewTestOutput(func(data []byte) {
if data[0] == '1' {
output := NewTestOutput(func(msg *Message) {
if msg.Meta[0] == '1' {
atomic.AddInt64(&reqCounter, 1)
} else {
atomic.AddInt64(&respCounter, 1)
@@ -193,11 +188,11 @@ func TestRAWInputIPv6(t *testing.T) {
})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
emitter := NewEmitter(quit)
emitter := NewEmitter()
addr := "http://" + originAddr
go emitter.Start(plugins, Settings.Middleware)
for i := 0; i < 10; i++ {
@@ -220,7 +215,6 @@ func TestRAWInputIPv6(t *testing.T) {
func TestInputRAWChunkedEncoding(t *testing.T) {
wg := new(sync.WaitGroup)
quit := make(chan int)
fileContent, _ := ioutil.ReadFile("README.md")
@@ -257,12 +251,12 @@ func TestInputRAWChunkedEncoding(t *testing.T) {
httpOutput := NewHTTPOutput(replay.URL, &HTTPOutputConfig{})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{httpOutput},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{httpOutput},
}
plugins.All = append(plugins.All, input, httpOutput)
emitter := NewEmitter(quit)
emitter := NewEmitter()
defer emitter.Close()
go emitter.Start(plugins, Settings.Middleware)
wg.Add(2)
@@ -278,17 +272,15 @@ func TestInputRAWChunkedEncoding(t *testing.T) {
}
func BenchmarkRAWInputWithReplay(b *testing.B) {
var respCounter, reqCounter, replayCounter uint64
var respCounter, reqCounter, replayCounter uint32
wg := &sync.WaitGroup{}
wg.Add(b.N * 3) // reqCounter + replayCounter + respCounter
quit := make(chan int)
listener, err := net.Listen("tcp", "127.0.0.1:0")
listener, err := net.Listen("tcp4", "127.0.0.1:0")
if err != nil {
b.Error(err)
return
}
listener0, err := net.Listen("tcp", "127.0.0.1:0")
listener0, err := net.Listen("tcp4", "127.0.0.1:0")
if err != nil {
b.Error(err)
return
@@ -305,6 +297,8 @@ func BenchmarkRAWInputWithReplay(b *testing.B) {
replay := http.Server{
Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
atomic.AddUint32(&replayCounter, 1)
w.Write(nil)
wg.Done()
}),
}
@@ -320,34 +314,37 @@ func BenchmarkRAWInputWithReplay(b *testing.B) {
}
input := NewRAWInput(originAddr, conf)
testOutput := NewTestOutput(func(data []byte) {
if data[0] == '1' {
atomic.AddUint64(&reqCounter, 1)
testOutput := NewTestOutput(func(msg *Message) {
if msg.Meta[0] == '1' {
reqCounter++
} else {
atomic.AddUint64(&respCounter, 1)
respCounter++
}
wg.Done()
})
httpOutput := NewHTTPOutput(replayAddr, &HTTPOutputConfig{})
httpOutput := NewHTTPOutput("http://"+replayAddr, &HTTPOutputConfig{})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{testOutput, httpOutput},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{testOutput, httpOutput},
}
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
now := time.Now()
addr := "http://" + originAddr
b.ResetTimer()
for i := 0; i < b.N; i++ {
_, err = http.Get(addr)
wg.Add(3) // reqCounter + replayCounter + respCounter
resp, err := http.Get(addr)
if err != nil {
b.Log(err)
wg.Add(-3)
}
resp.Body.Close()
}
wg.Wait()
b.Logf("%d/%d Requests, %d/%d Responses, %d/%d Replayed in %s\n", reqCounter, b.N, respCounter, b.N, replayCounter, b.N, time.Since(now))
b.ReportMetric(float64(reqCounter), "requests")
b.ReportMetric(float64(respCounter), "responses")
b.ReportMetric(float64(replayCounter), "replayed")
emitter.Close()
}
input_tcp.go
+43 -29
View File
@@ -8,18 +8,18 @@ import (
"io"
"log"
"net"
"os"
)
// TCPInput used for internal communication
type TCPInput struct {
data chan []byte
data chan *Message
listener net.Listener
address string
config *TCPInputConfig
stop chan bool // Channel used only to indicate goroutine should shutdown
}
// TCPInputConfig represents configuration of a TCP input plugin
type TCPInputConfig struct {
Secure bool `json:"input-tcp-secure"`
CertificatePath string `json:"input-tcp-certificate"`
@@ -29,7 +29,7 @@ type TCPInputConfig struct {
// NewTCPInput constructor for TCPInput, accepts address with port
func NewTCPInput(address string, config *TCPInputConfig) (i *TCPInput) {
i = new(TCPInput)
i.data = make(chan []byte, 1000)
i.data = make(chan *Message, 1000)
i.address = address
i.config = config
i.stop = make(chan bool)
@@ -39,20 +39,21 @@ func NewTCPInput(address string, config *TCPInputConfig) (i *TCPInput) {
return
}
func (i *TCPInput) Read(data []byte) (int, error) {
var buf []byte
// PluginRead returns data and details read from plugin
func (i *TCPInput) PluginRead() (msg *Message, err error) {
select {
case <-i.stop:
return 0, ErrorStopped
case buf = <-i.data:
return nil, ErrorStopped
case msg = <-i.data:
return msg, nil
}
copy(data, buf)
return len(buf), nil
}
// Close closes the plugin
func (i *TCPInput) Close() error {
close(i.stop)
i.listener.Close()
return nil
}
@@ -60,64 +61,67 @@ func (i *TCPInput) listen(address string) {
if i.config.Secure {
cer, err := tls.LoadX509KeyPair(i.config.CertificatePath, i.config.KeyPath)
if err != nil {
log.Fatal("Error while loading --input-file certificate:", err)
log.Fatalln("error while loading --input-tcp TLS certificate:", err)
}
config := &tls.Config{Certificates: []tls.Certificate{cer}}
listener, err := tls.Listen("tcp", address, config)
if err != nil {
log.Fatal("Can't start --input-tcp with secure connection:", err)
log.Fatalln("[INPUT-TCP] failed to start INPUT-TCP listener:", err)
}
i.listener = listener
} else {
listener, err := net.Listen("tcp", address)
if err != nil {
log.Fatal("Can't start:", err)
log.Fatalln("failed to start INPUT-TCP listener:", err)
}
i.listener = listener
}
go func() {
for {
conn, err := i.listener.Accept()
if err != nil {
log.Println("Error while Accept()", err)
if err == nil {
go i.handleConnection(conn)
continue
}
go i.handleConnection(conn)
if isTemporaryNetworkError(err) {
continue
}
if operr, ok := err.(*net.OpError); ok && operr.Err.Error() != "use of closed network connection" {
Debug(0, fmt.Sprintf("[INPUT-TCP] listener closed, err: %q", err))
}
break
}
}()
}
var payloadSeparatorAsBytes = []byte(payloadSeparator)
func (i *TCPInput) handleConnection(conn net.Conn) {
defer conn.Close()
payloadSeparatorAsBytes := []byte(payloadSeparator)
reader := bufio.NewReader(conn)
var buffer bytes.Buffer
for {
line, err := reader.ReadBytes('\n')
if err != nil {
if isTemporaryNetworkError(err) {
continue
}
if err != io.EOF {
fmt.Fprintln(os.Stderr, "Unexpected error in input tcp connection:", err)
Debug(0, fmt.Sprintf("[INPUT-TCP] connection error: %q", err))
}
break
}
if bytes.Equal(payloadSeparatorAsBytes[1:], line) {
asBytes := buffer.Bytes()
// unread the '\n' before monkeys
buffer.UnreadByte()
var msg Message
msg.Meta, msg.Data = payloadMetaWithBody(buffer.Bytes())
i.data <- &msg
buffer.Reset()
newBuf := make([]byte, len(asBytes)-1)
copy(newBuf, asBytes)
i.data <- newBuf
} else {
buffer.Write(line)
}
@@ -127,3 +131,13 @@ func (i *TCPInput) handleConnection(conn net.Conn) {
func (i *TCPInput) String() string {
return "TCP input: " + i.address
}
func isTemporaryNetworkError(err error) bool {
if nerr, ok := err.(net.Error); ok && nerr.Temporary() {
return true
}
if operr, ok := err.(*net.OpError); ok && operr.Temporary() {
return true
}
return false
}
input_tcp_test.go
+15 -15
View File
@@ -7,7 +7,6 @@ import (
"crypto/tls"
"crypto/x509"
"encoding/pem"
"io"
"io/ioutil"
"log"
"math/big"
@@ -20,20 +19,19 @@ import (
func TestTCPInput(t *testing.T) {
wg := new(sync.WaitGroup)
quit := make(chan int)
input := NewTCPInput("127.0.0.1:0", &TCPInputConfig{})
output := NewTestOutput(func(data []byte) {
output := NewTestOutput(func(*Message) {
wg.Done()
})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
plugins.All = append(plugins.All, input, output)
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
tcpAddr, err := net.ResolveTCPAddr("tcp", input.listener.Addr().String())
@@ -43,7 +41,6 @@ func TestTCPInput(t *testing.T) {
}
conn, err := net.DialTCP("tcp", nil, tcpAddr)
if err != nil {
log.Fatal(err)
}
@@ -52,10 +49,14 @@ func TestTCPInput(t *testing.T) {
for i := 0; i < 100; i++ {
wg.Add(1)
conn.Write(msg)
conn.Write([]byte(payloadSeparator))
if _, err = conn.Write(msg); err == nil {
_, err = conn.Write(payloadSeparatorAsBytes)
}
if err != nil {
t.Error(err)
return
}
}
wg.Wait()
emitter.Close()
}
@@ -99,24 +100,23 @@ func TestTCPInputSecure(t *testing.T) {
}()
wg := new(sync.WaitGroup)
quit := make(chan int)
input := NewTCPInput("127.0.0.1:0", &TCPInputConfig{
Secure: true,
CertificatePath: serverCertPemFile.Name(),
KeyPath: serverPrivPemFile.Name(),
})
output := NewTestOutput(func(data []byte) {
output := NewTestOutput(func(*Message) {
wg.Done()
})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
plugins.All = append(plugins.All, input, output)
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
conf := &tls.Config{
limiter.go
+18 -13
View File
@@ -20,8 +20,8 @@ type Limiter struct {
}
func parseLimitOptions(options string) (limit int, isPercent bool) {
if strings.Contains(options, "%") {
limit, _ = strconv.Atoi(strings.Split(options, "%")[0])
if n := strings.Index(options, "%"); n > 0 {
limit, _ = strconv.Atoi(options[:n])
isPercent = true
} else {
limit, _ = strconv.Atoi(options)
@@ -33,7 +33,7 @@ func parseLimitOptions(options string) (limit int, isPercent bool) {
// NewLimiter constructor for Limiter, accepts plugin and options
// `options` allow to sprcify relatve or absolute limiting
func NewLimiter(plugin interface{}, options string) io.ReadWriter {
func NewLimiter(plugin interface{}, options string) PluginReadWriter {
l := new(Limiter)
l.limit, l.isPercent = parseLimitOptions(options)
l.plugin = plugin
@@ -71,24 +71,29 @@ func (l *Limiter) isLimited() bool {
return false
}
func (l *Limiter) Write(data []byte) (n int, err error) {
// PluginWrite writes message to this plugin
func (l *Limiter) PluginWrite(msg *Message) (n int, err error) {
if l.isLimited() {
return 0, nil
}
n, err = l.plugin.(io.Writer).Write(data)
return
if w, ok := l.plugin.(PluginWriter); ok {
return w.PluginWrite(msg)
}
// avoid further writing
return 0, io.ErrClosedPipe
}
func (l *Limiter) Read(data []byte) (n int, err error) {
if r, ok := l.plugin.(io.Reader); ok {
n, err = r.Read(data)
// PluginRead reads message from this plugin
func (l *Limiter) PluginRead() (msg *Message, err error) {
if r, ok := l.plugin.(PluginReader); ok {
msg, err = r.PluginRead()
} else {
return 0, nil
// avoid further reading
return nil, io.ErrClosedPipe
}
if l.isLimited() {
return 0, nil
return nil, nil
}
return
@@ -100,7 +105,7 @@ func (l *Limiter) String() string {
// Close closes the resources.
func (l *Limiter) Close() error {
if fi, ok := l.plugin.(io.ReadCloser); ok {
if fi, ok := l.plugin.(io.Closer); ok {
fi.Close()
}
return nil
limiter_test.go
+16 -21
View File
@@ -3,28 +3,26 @@
package main
import (
"io"
"sync"
"testing"
)
func TestOutputLimiter(t *testing.T) {
wg := new(sync.WaitGroup)
quit := make(chan int)
input := NewTestInput()
output := NewLimiter(NewTestOutput(func(data []byte) {
output := NewLimiter(NewTestOutput(func(*Message) {
wg.Done()
}), "10")
wg.Add(10)
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
plugins.All = append(plugins.All, input, output)
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
for i := 0; i < 100; i++ {
@@ -37,21 +35,20 @@ func TestOutputLimiter(t *testing.T) {
func TestInputLimiter(t *testing.T) {
wg := new(sync.WaitGroup)
quit := make(chan int)
input := NewLimiter(NewTestInput(), "10")
output := NewTestOutput(func(data []byte) {
output := NewTestOutput(func(*Message) {
wg.Done()
})
wg.Add(10)
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
plugins.All = append(plugins.All, input, output)
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
for i := 0; i < 100; i++ {
@@ -65,20 +62,19 @@ func TestInputLimiter(t *testing.T) {
// Should limit all requests
func TestPercentLimiter1(t *testing.T) {
wg := new(sync.WaitGroup)
quit := make(chan int)
input := NewTestInput()
output := NewLimiter(NewTestOutput(func(data []byte) {
output := NewLimiter(NewTestOutput(func(*Message) {
wg.Done()
}), "0%")
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
plugins.All = append(plugins.All, input, output)
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
for i := 0; i < 100; i++ {
@@ -91,21 +87,20 @@ func TestPercentLimiter1(t *testing.T) {
// Should not limit at all
func TestPercentLimiter2(t *testing.T) {
wg := new(sync.WaitGroup)
quit := make(chan int)
input := NewTestInput()
output := NewLimiter(NewTestOutput(func(data []byte) {
output := NewLimiter(NewTestOutput(func(*Message) {
wg.Done()
}), "100%")
wg.Add(100)
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
plugins.All = append(plugins.All, input, output)
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
for i := 0; i < 100; i++ {
middleware.go
+37 -52
View File
@@ -2,6 +2,7 @@ package main
import (
"bufio"
"context"
"encoding/hex"
"fmt"
"io"
@@ -9,32 +10,29 @@ import (
"os"
"os/exec"
"strings"
"sync"
)
// Middleware represents a middleware object
type Middleware struct {
command string
data chan []byte
mu sync.Mutex
Stdin io.Writer
Stdout io.Reader
stop chan bool // Channel used only to indicate goroutine should shutdown
command string
data chan *Message
Stdin io.Writer
Stdout io.Reader
commandCancel context.CancelFunc
stop chan bool // Channel used only to indicate goroutine should shutdown
}
// NewMiddleware returns new middleware
func NewMiddleware(command string) *Middleware {
m := new(Middleware)
m.command = command
m.data = make(chan []byte, 1000)
m.data = make(chan *Message, 1000)
m.stop = make(chan bool)
commands := strings.Split(command, " ")
cmd := exec.Command(commands[0], commands[1:]...)
ctx, cancl := context.WithCancel(context.Background())
m.commandCancel = cancl
cmd := exec.CommandContext(ctx, commands[0], commands[1:]...)
m.Stdout, _ = cmd.StdoutPipe()
m.Stdin, _ = cmd.StdinPipe()
@@ -61,45 +59,32 @@ func NewMiddleware(command string) *Middleware {
}
// ReadFrom start a worker to read from this plugin
func (m *Middleware) ReadFrom(plugin io.Reader) {
func (m *Middleware) ReadFrom(plugin PluginReader) {
Debug(2, "[MIDDLEWARE-MASTER] Starting reading from", plugin)
go m.copy(m.Stdin, plugin)
}
func (m *Middleware) copy(to io.Writer, from io.Reader) {
buf := make([]byte, 5*1024*1024)
dst := make([]byte, len(buf)*4)
func (m *Middleware) copy(to io.Writer, from PluginReader) {
var buf, dst []byte
for {
nr, _ := from.Read(buf)
if nr == 0 || nr > len(buf) {
msg, err := from.PluginRead()
if err != nil {
return
}
if msg == nil || len(msg.Data) == 0 {
continue
}
payload := buf[0:nr]
if Settings.PrettifyHTTP {
payload = prettifyHTTP(payload)
nr = len(payload)
if nr*2 > len(dst) {
continue
}
buf = prettifyHTTP(msg.Data)
}
dst = make([]byte, len(buf)*2+1)
hex.Encode(dst, buf)
dst[len(buf)*2] = '\n'
if Settings.PrettifyHTTP {
payload = prettifyHTTP(payload)
nr = len(payload)
}
to.Write(dst)
hex.Encode(dst, payload)
dst[nr*2] = '\n'
m.mu.Lock()
to.Write(dst[0 : nr*2+1])
m.mu.Unlock()
Debug(3, "[MIDDLEWARE-MASTER] Sending:", string(buf[0:nr]), "From:", from)
}
}
@@ -117,41 +102,41 @@ func (m *Middleware) read(from io.Reader) {
}
}
buf := make([]byte, len(line)/2)
buf := make([]byte, len(line)/2-1)
if _, err := hex.Decode(buf, line[:len(line)-1]); err != nil {
fmt.Fprintln(os.Stderr, "Failed to decode input payload", err, len(line), string(line[:len(line)-1]))
Debug(0, fmt.Sprintf("[MIDDLEWARE] failed to decode err: %q", err))
continue
}
Debug(3, "[MIDDLEWARE-MASTER] Received:", string(buf))
var msg Message
msg.Meta, msg.Data = payloadMetaWithBody(buf)
select {
case <-m.stop:
return
case m.data <- buf:
case m.data <- &msg:
}
}
return
}
func (m *Middleware) Read(data []byte) (int, error) {
var buf []byte
// PluginRead reads message from this plugin
func (m *Middleware) PluginRead() (msg *Message, err error) {
select {
case <-m.stop:
return 0, ErrorStopped
case buf = <-m.data:
return nil, ErrorStopped
case msg = <-m.data:
}
n := copy(data, buf)
return n, nil
return
}
func (m *Middleware) String() string {
return fmt.Sprintf("Modifying traffic using '%s' command", m.command)
return fmt.Sprintf("Modifying traffic using %q command", m.command)
}
// Close closes this plugin
func (m *Middleware) Close() error {
m.commandCancel()
close(m.stop)
return nil
}
output_binary.go
+40 -44
View File
@@ -1,7 +1,6 @@
package main
import (
"io"
"sync/atomic"
"time"
@@ -25,30 +24,27 @@ type BinaryOutput struct {
// alignment. atomic.* functions crash on 32bit machines if operand is not
// aligned at 64bit. See https://github.com/golang/go/issues/599
activeWorkers int64
address string
queue chan []byte
responses chan response
needWorker chan int
config *BinaryOutputConfig
queueStats *GorStat
address string
queue chan *Message
responses chan response
needWorker chan int
quit chan struct{}
config *BinaryOutputConfig
queueStats *GorStat
}
// NewBinaryOutput constructor for BinaryOutput
// Initialize workers
func NewBinaryOutput(address string, config *BinaryOutputConfig) io.Writer {
func NewBinaryOutput(address string, config *BinaryOutputConfig) PluginReadWriter {
o := new(BinaryOutput)
o.address = address
o.config = config
o.queue = make(chan []byte, 1000)
o.queue = make(chan *Message, 1000)
o.responses = make(chan response, 1000)
o.needWorker = make(chan int, 1)
o.quit = make(chan struct{})
// Initial workers count
if o.config.Workers == 0 {
@@ -89,8 +85,8 @@ func (o *BinaryOutput) startWorker() {
for {
select {
case data := <-o.queue:
o.sendRequest(client, data)
case msg := <-o.queue:
o.sendRequest(client, msg)
deathCount = 0
case <-time.After(time.Millisecond * 100):
// When dynamic scaling enabled workers die after 2s of inactivity
@@ -113,15 +109,13 @@ func (o *BinaryOutput) startWorker() {
}
}
func (o *BinaryOutput) Write(data []byte) (n int, err error) {
if !isRequestPayload(data) {
return len(data), nil
// PluginWrite writes a message tothis plugin
func (o *BinaryOutput) PluginWrite(msg *Message) (n int, err error) {
if !isRequestPayload(msg.Meta) {
return len(msg.Data), nil
}
buf := make([]byte, len(data))
copy(buf, data)
o.queue <- buf
o.queue <- msg
if o.config.Workers == 0 {
workersCount := atomic.LoadInt64(&o.activeWorkers)
@@ -131,37 +125,33 @@ func (o *BinaryOutput) Write(data []byte) (n int, err error) {
}
}
return len(data), nil
return len(msg.Data) + len(msg.Meta), nil
}
func (o *BinaryOutput) Read(data []byte) (int, error) {
resp := <-o.responses
// PluginRead reads a message from this plugin
func (o *BinaryOutput) PluginRead() (*Message, error) {
var resp response
var msg Message
select {
case <-o.quit:
return nil, ErrorStopped
case resp = <-o.responses:
}
msg.Data = resp.payload
msg.Meta = payloadHeader(ReplayedResponsePayload, resp.uuid, resp.startedAt, resp.roundTripTime)
Debug(2, "[OUTPUT-TCP] Received response:", string(resp.payload))
header := payloadHeader(ReplayedResponsePayload, resp.uuid, resp.startedAt, resp.roundTripTime)
copy(data[0:len(header)], header)
copy(data[len(header):], resp.payload)
return len(resp.payload) + len(header), nil
return &msg, nil
}
func (o *BinaryOutput) sendRequest(client *TCPClient, request []byte) {
meta := payloadMeta(request)
if len(meta) < 2 {
func (o *BinaryOutput) sendRequest(client *TCPClient, msg *Message) {
if !isRequestPayload(msg.Meta) {
return
}
if !isRequestPayload(request) {
return
}
uuid := meta[1]
body := payloadBody(request)
uuid := payloadID(msg.Meta)
start := time.Now()
resp, err := client.Send(body)
resp, err := client.Send(msg.Data)
stop := time.Now()
if err != nil {
@@ -176,3 +166,9 @@ func (o *BinaryOutput) sendRequest(client *TCPClient, request []byte) {
func (o *BinaryOutput) String() string {
return "Binary output: " + o.address
}
// Close closes this plugin for reading
func (o *BinaryOutput) Close() error {
close(o.quit)
return nil
}
output_dummy.go
+9 -3
View File
@@ -15,9 +15,15 @@ func NewDummyOutput() (di *DummyOutput) {
return
}
func (i *DummyOutput) Write(data []byte) (int, error) {
n, err := os.Stdout.Write(data)
os.Stdout.Write([]byte{'\n'})
// PluginWrite writes message to this plugin
func (i *DummyOutput) PluginWrite(msg *Message) (int, error) {
var n, nn int
var err error
n, err = os.Stdout.Write(msg.Meta)
nn, err = os.Stdout.Write(msg.Data)
n += nn
nn, err = os.Stdout.Write(payloadSeparatorAsBytes)
n += nn
return n, err
}
output_file.go
+12 -9
View File
@@ -195,10 +195,11 @@ func (o *FileOutput) updateName() {
o.Unlock()
}
func (o *FileOutput) Write(data []byte) (n int, err error) {
// PluginWrite writes message to this plugin
func (o *FileOutput) PluginWrite(msg *Message) (n int, err error) {
if o.requestPerFile {
o.Lock()
meta := payloadMeta(data)
meta := payloadMeta(msg.Meta)
o.currentID = meta[1]
o.payloadType = meta[0]
o.Unlock()
@@ -227,10 +228,12 @@ func (o *FileOutput) Write(data []byte) (n int, err error) {
o.QueueLength = 0
}
n, _ = o.writer.Write(data)
nSeparator, _ := o.writer.Write([]byte(payloadSeparator))
n += nSeparator
var nn int
n, err = o.writer.Write(msg.Meta)
nn, err = o.writer.Write(msg.Data)
n += nn
nn, err = o.writer.Write(payloadSeparatorAsBytes)
n += nn
o.totalFileSize += size.Size(n)
o.QueueLength++
@@ -239,14 +242,14 @@ func (o *FileOutput) Write(data []byte) (n int, err error) {
return n, errors.New("File output reached size limit")
}
return n, nil
return n, err
}
func (o *FileOutput) flush() {
// Don't exit on panic
defer func() {
if r := recover(); r != nil {
log.Println("PANIC while file flush: ", r, o, string(debug.Stack()))
Debug(0, "[OUTPUT-FILE] PANIC while file flush: ", r, o, string(debug.Stack()))
}
}()
@@ -263,7 +266,7 @@ func (o *FileOutput) flush() {
if stat, err := o.file.Stat(); err == nil {
o.chunkSize = int(stat.Size())
} else {
log.Println("Error accessing file sats", err)
Debug(0, "[OUTPUT-HTTP] error accessing file size", err)
}
}
}
output_file_test.go
+26 -29
View File
@@ -2,7 +2,6 @@ package main
import (
"fmt"
"io"
"math/rand"
"os"
"reflect"
@@ -17,18 +16,17 @@ import (
func TestFileOutput(t *testing.T) {
wg := new(sync.WaitGroup)
quit := make(chan int)
input := NewTestInput()
output := NewFileOutput("/tmp/test_requests.gor", &FileOutputConfig{FlushInterval: time.Minute, Append: true})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
plugins.All = append(plugins.All, input, output)
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
for i := 0; i < 100; i++ {
@@ -42,19 +40,18 @@ func TestFileOutput(t *testing.T) {
var counter int64
input2 := NewFileInput("/tmp/test_requests.gor", false)
output2 := NewTestOutput(func(data []byte) {
output2 := NewTestOutput(func(*Message) {
atomic.AddInt64(&counter, 1)
wg.Done()
})
plugins2 := &InOutPlugins{
Inputs: []io.Reader{input2},
Outputs: []io.Writer{output2},
Inputs: []PluginReader{input2},
Outputs: []PluginWriter{output2},
}
plugins2.All = append(plugins2.All, input2, output2)
quit2 := make(chan int)
emitter2 := NewEmitter(quit2)
emitter2 := NewEmitter()
go emitter2.Start(plugins2, Settings.Middleware)
wg.Wait()
@@ -91,16 +88,16 @@ func TestFileOutputMultipleFiles(t *testing.T) {
t.Error("Should not initialize file if no writes")
}
output.Write([]byte("1 1 1\r\ntest"))
output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")})
name1 := output.file.Name()
output.Write([]byte("1 1 1\r\ntest"))
output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")})
name2 := output.file.Name()
time.Sleep(time.Second)
output.updateName()
output.Write([]byte("1 1 1\r\ntest"))
output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")})
name3 := output.file.Name()
if name2 != name1 {
@@ -122,16 +119,16 @@ func TestFileOutputFilePerRequest(t *testing.T) {
t.Error("Should not initialize file if no writes")
}
output.Write([]byte("1 1 1\ntest"))
output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")})
name1 := output.file.Name()
output.Write([]byte("1 2 1\ntest"))
output.PluginWrite(&Message{Meta: []byte("1 2 1\r\n"), Data: []byte("test")})
name2 := output.file.Name()
time.Sleep(time.Second)
output.updateName()
output.Write([]byte("1 3 1\ntest"))
output.PluginWrite(&Message{Meta: []byte("1 3 1\r\n"), Data: []byte("test")})
name3 := output.file.Name()
if name3 == name2 || name2 == name1 || name3 == name1 {
@@ -151,7 +148,7 @@ func TestFileOutputCompression(t *testing.T) {
}
for i := 0; i < 1000; i++ {
output.Write([]byte("1 1 1\r\ntest"))
output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")})
}
name := output.file.Name()
@@ -210,15 +207,15 @@ func TestFileOutputAppendQueueLimitOverflow(t *testing.T) {
output := NewFileOutput(name, &FileOutputConfig{Append: false, FlushInterval: time.Minute, QueueLimit: 2})
output.Write([]byte("1 1 1\r\ntest"))
output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")})
name1 := output.file.Name()
output.Write([]byte("1 1 1\r\ntest"))
output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")})
name2 := output.file.Name()
output.updateName()
output.Write([]byte("1 1 1\r\ntest"))
output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")})
name3 := output.file.Name()
if name2 != name1 || name1 != fmt.Sprintf("/tmp/%d_0", rnd) {
@@ -239,15 +236,15 @@ func TestFileOutputAppendQueueLimitNoOverflow(t *testing.T) {
output := NewFileOutput(name, &FileOutputConfig{Append: false, FlushInterval: time.Minute, QueueLimit: 3})
output.Write([]byte("1 1 1\r\ntest"))
output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")})
name1 := output.file.Name()
output.Write([]byte("1 1 1\r\ntest"))
output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")})
name2 := output.file.Name()
output.updateName()
output.Write([]byte("1 1 1\r\ntest"))
output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")})
name3 := output.file.Name()
if name2 != name1 || name1 != fmt.Sprintf("/tmp/%d_0", rnd) {
@@ -268,15 +265,15 @@ func TestFileOutputAppendQueueLimitGzips(t *testing.T) {
output := NewFileOutput(name, &FileOutputConfig{Append: false, FlushInterval: time.Minute, QueueLimit: 2})
output.Write([]byte("1 1 1\r\ntest"))
output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")})
name1 := output.file.Name()
output.Write([]byte("1 1 1\r\ntest"))
output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")})
name2 := output.file.Name()
output.updateName()
output.Write([]byte("1 1 1\r\ntest"))
output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")})
name3 := output.file.Name()
if name2 != name1 || name1 != fmt.Sprintf("/tmp/%d_0.gz", rnd) {
@@ -311,16 +308,16 @@ func TestFileOutputAppendSizeLimitOverflow(t *testing.T) {
output := NewFileOutput(name, &FileOutputConfig{Append: false, FlushInterval: time.Minute, SizeLimit: size.Size(2 * messageSize)})
output.Write([]byte("1 1 1\r\ntest"))
output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")})
name1 := output.file.Name()
output.Write([]byte("1 1 1\r\ntest"))
output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")})
name2 := output.file.Name()
output.flush()
output.updateName()
output.Write([]byte("1 1 1\r\ntest"))
output.PluginWrite(&Message{Meta: []byte("1 1 1\r\n"), Data: []byte("test")})
name3 := output.file.Name()
if name2 != name1 || name1 != fmt.Sprintf("/tmp/%d_0", rnd) {
output_http.go
+26 -31
View File
@@ -5,7 +5,6 @@ import (
"bytes"
"crypto/tls"
"fmt"
"io"
"log"
"math"
"net/http"
@@ -59,14 +58,14 @@ type HTTPOutput struct {
elasticSearch *ESPlugin
client *HTTPClient
stopWorker chan struct{}
queue chan []byte
queue chan *Message
responses chan response
stop chan bool // Channel used only to indicate goroutine should shutdown
}
// NewHTTPOutput constructor for HTTPOutput
// Initialize workers
func NewHTTPOutput(address string, config *HTTPOutputConfig) io.Writer {
func NewHTTPOutput(address string, config *HTTPOutputConfig) PluginReadWriter {
o := new(HTTPOutput)
var err error
config.url, err = url.Parse(address)
@@ -110,7 +109,7 @@ func NewHTTPOutput(address string, config *HTTPOutputConfig) io.Writer {
o.queueStats = NewGorStat("output_http", o.config.StatsMs)
}
o.queue = make(chan []byte, o.config.QueueLen)
o.queue = make(chan *Message, o.config.QueueLen)
o.responses = make(chan response, o.config.QueueLen)
// it should not be buffered to avoid races
o.stopWorker = make(chan struct{})
@@ -160,23 +159,22 @@ func (o *HTTPOutput) startWorker() {
select {
case <-o.stopWorker:
return
case data := <-o.queue:
o.sendRequest(o.client, data)
case msg := <-o.queue:
o.sendRequest(o.client, msg)
}
}
}
func (o *HTTPOutput) Write(data []byte) (n int, err error) {
if !isRequestPayload(data) {
return len(data), nil
// PluginWrite writes message to this plugin
func (o *HTTPOutput) PluginWrite(msg *Message) (n int, err error) {
if !isRequestPayload(msg.Meta) {
return len(msg.Data), nil
}
buf := make([]byte, len(data))
copy(buf, data)
select {
case <-o.stop:
return 0, ErrorStopped
case o.queue <- buf:
case o.queue <- msg:
}
if o.config.Stats {
@@ -189,51 +187,48 @@ func (o *HTTPOutput) Write(data []byte) (n int, err error) {
atomic.AddInt32(&o.activeWorkers, 1)
}
}
return len(data), nil
return len(msg.Data) + len(msg.Meta), nil
}
func (o *HTTPOutput) Read(data []byte) (int, error) {
// PluginRead reads message from this plugin
func (o *HTTPOutput) PluginRead() (*Message, error) {
var resp response
var msg Message
select {
case <-o.stop:
return 0, ErrorStopped
return nil, ErrorStopped
case resp = <-o.responses:
msg.Data = resp.payload
}
header := payloadHeader(ReplayedResponsePayload, resp.uuid, resp.roundTripTime, resp.startedAt)
n := copy(data, header)
if len(data) > len(header) {
n += copy(data[len(header):], resp.payload)
}
dis := len(header) + len(data) - n
if dis > 0 {
Debug(2, fmt.Sprintf("[OUTPUT-HTTP] %dB discarded increase copy buffer size", dis))
}
msg.Meta = payloadHeader(ReplayedResponsePayload, resp.uuid, resp.roundTripTime, resp.startedAt)
return n, nil
return &msg, nil
}
func (o *HTTPOutput) sendRequest(client *HTTPClient, request []byte) {
if !isRequestPayload(request) {
func (o *HTTPOutput) sendRequest(client *HTTPClient, msg *Message) {
if !isRequestPayload(msg.Meta) {
return
}
uuid := payloadID(request)
body := payloadBody(request)
uuid := payloadID(msg.Meta)
start := time.Now()
resp, err := client.Send(body)
resp, err := client.Send(msg.Data)
stop := time.Now()
if err != nil {
Debug(1, fmt.Sprintf("[HTTP-OUTPUT] error when sending: %q", err))
return
}
if resp == nil {
return
}
if o.config.TrackResponses {
o.responses <- response{resp, uuid, start.UnixNano(), stop.UnixNano() - start.UnixNano()}
}
if o.elasticSearch != nil {
o.elasticSearch.ResponseAnalyze(request, resp, start, stop)
o.elasticSearch.ResponseAnalyze(msg.Data, resp, start, stop)
}
}
output_http_test.go
+29 -33
View File
@@ -1,7 +1,6 @@
package main
import (
"io"
"io/ioutil"
"net/http"
"net/http/httptest"
@@ -12,7 +11,6 @@ import (
func TestHTTPOutput(t *testing.T) {
wg := new(sync.WaitGroup)
quit := make(chan int)
input := NewTestInput()
@@ -38,27 +36,27 @@ func TestHTTPOutput(t *testing.T) {
}))
defer server.Close()
headers := HTTPHeaders{HTTPHeader{"User-Agent", "Gor"}}
headers := HTTPHeaders{httpHeader{"User-Agent", "Gor"}}
methods := HTTPMethods{[]byte("GET"), []byte("PUT"), []byte("POST")}
Settings.ModifierConfig = HTTPModifierConfig{Headers: headers, Methods: methods}
httpOutput := NewHTTPOutput(server.URL, &HTTPOutputConfig{TrackResponses: true})
output := NewTestOutput(func(data []byte) {
httpOutput := NewHTTPOutput(server.URL, &HTTPOutputConfig{TrackResponses: false})
output := NewTestOutput(func(*Message) {
wg.Done()
})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{httpOutput, output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{httpOutput, output},
}
plugins.All = append(plugins.All, input, output, httpOutput)
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
for i := 0; i < 10; i++ {
// 2 http-output, 2 - test output request, 2 - test output http response
wg.Add(6) // OPTIONS should be ignored
// 2 http-output, 2 - test output request
wg.Add(4) // OPTIONS should be ignored
input.EmitPOST()
input.EmitOPTIONS()
input.EmitGET()
@@ -72,7 +70,6 @@ func TestHTTPOutput(t *testing.T) {
func TestHTTPOutputKeepOriginalHost(t *testing.T) {
wg := new(sync.WaitGroup)
quit := make(chan int)
input := NewTestInput()
@@ -85,18 +82,18 @@ func TestHTTPOutputKeepOriginalHost(t *testing.T) {
}))
defer server.Close()
headers := HTTPHeaders{HTTPHeader{"Host", "custom-host.com"}}
headers := HTTPHeaders{httpHeader{"Host", "custom-host.com"}}
Settings.ModifierConfig = HTTPModifierConfig{Headers: headers}
output := NewHTTPOutput(server.URL, &HTTPOutputConfig{OriginalHost: true, SkipVerify: true})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
plugins.All = append(plugins.All, input, output)
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
wg.Add(1)
@@ -109,7 +106,6 @@ func TestHTTPOutputKeepOriginalHost(t *testing.T) {
func TestHTTPOutputSSL(t *testing.T) {
wg := new(sync.WaitGroup)
quit := make(chan int)
// Origing and Replay server initialization
server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -120,12 +116,12 @@ func TestHTTPOutputSSL(t *testing.T) {
output := NewHTTPOutput(server.URL, &HTTPOutputConfig{SkipVerify: true})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
plugins.All = append(plugins.All, input, output)
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
wg.Add(2)
@@ -139,7 +135,6 @@ func TestHTTPOutputSSL(t *testing.T) {
func TestHTTPOutputSessions(t *testing.T) {
wg := new(sync.WaitGroup)
quit := make(chan int)
input := NewTestInput()
input.skipHeader = true
@@ -150,25 +145,27 @@ func TestHTTPOutputSessions(t *testing.T) {
defer server.Close()
Settings.RecognizeTCPSessions = true
Settings.SplitOutput = true
output := NewHTTPOutput(server.URL, &HTTPOutputConfig{})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
emitter := NewEmitter(quit)
plugins.All = append(plugins.All, input, output)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
uuid1 := []byte("1234567890123456789a0000")
uuid2 := []byte("1234567890123456789d0000")
for i := 0; i < 100; i++ {
for i := 0; i < 10; i++ {
wg.Add(1) // OPTIONS should be ignored
copy(uuid1[20:], randByte(4))
input.EmitBytes([]byte("1 " + string(uuid1) + " 1\n" + "GET / HTTP/1.1\r\n\r\n"))
}
for i := 0; i < 100; i++ {
for i := 0; i < 10; i++ {
wg.Add(1) // OPTIONS should be ignored
copy(uuid2[20:], randByte(4))
input.EmitBytes([]byte("1 " + string(uuid2) + " 1\n" + "GET / HTTP/1.1\r\n\r\n"))
@@ -179,11 +176,11 @@ func TestHTTPOutputSessions(t *testing.T) {
emitter.Close()
Settings.RecognizeTCPSessions = false
Settings.SplitOutput = false
}
func BenchmarkHTTPOutput(b *testing.B) {
wg := new(sync.WaitGroup)
quit := make(chan int)
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
wg.Done()
@@ -194,12 +191,12 @@ func BenchmarkHTTPOutput(b *testing.B) {
output := NewHTTPOutput(server.URL, &HTTPOutputConfig{WorkersMax: 1})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
plugins.All = append(plugins.All, input, output)
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
for i := 0; i < b.N; i++ {
@@ -213,7 +210,6 @@ func BenchmarkHTTPOutput(b *testing.B) {
func BenchmarkHTTPOutputTLS(b *testing.B) {
wg := new(sync.WaitGroup)
quit := make(chan int)
server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
wg.Done()
@@ -224,12 +220,12 @@ func BenchmarkHTTPOutputTLS(b *testing.B) {
output := NewHTTPOutput(server.URL, &HTTPOutputConfig{SkipVerify: true, WorkersMax: 1})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
plugins.All = append(plugins.All, input, output)
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
for i := 0; i < b.N; i++ {
output_kafka.go
+7 -7
View File
@@ -2,7 +2,6 @@ package main
import (
"encoding/json"
"io"
"log"
"strings"
"time"
@@ -24,7 +23,7 @@ type KafkaOutput struct {
const KafkaOutputFrequency = 500
// NewKafkaOutput creates instance of kafka producer client with TLS config
func NewKafkaOutput (address string, config *OutputKafkaConfig, tlsConfig *KafkaTLSConfig) io.Writer {
func NewKafkaOutput(address string, config *OutputKafkaConfig, tlsConfig *KafkaTLSConfig) PluginWriter {
c := NewKafkaConfig(tlsConfig)
var producer sarama.AsyncProducer
@@ -63,20 +62,21 @@ func (o *KafkaOutput) ErrorHandler() {
}
}
func (o *KafkaOutput) Write(data []byte) (n int, err error) {
// PluginWrite writes a message to this plugin
func (o *KafkaOutput) PluginWrite(msg *Message) (n int, err error) {
var message sarama.StringEncoder
if !o.config.UseJSON {
message = sarama.StringEncoder(byteutils.SliceToString(data))
message = sarama.StringEncoder(byteutils.SliceToString(msg.Meta) + byteutils.SliceToString(msg.Data))
} else {
mimeHeader := proto.ParseHeaders(data)
mimeHeader := proto.ParseHeaders(msg.Data)
var header map[string]string
for k, v := range mimeHeader {
header[k] = strings.Join(v, ", ")
}
meta := payloadMeta(data)
req := payloadBody(data)
meta := payloadMeta(msg.Meta)
req := msg.Data
kafkaMessage := KafkaMessage{
ReqURL: byteutils.SliceToString(proto.Path(req)),
output_kafka_test.go
+4 -4
View File
@@ -17,16 +17,16 @@ func TestOutputKafkaRAW(t *testing.T) {
producer: producer,
Topic: "test",
UseJSON: false,
},nil)
}, nil)
output.Write([]byte("1 2 3\nGET / HTTP1.1\r\nHeader: 1\r\n\r\n"))
output.PluginWrite(&Message{Meta: []byte("1 2 3\n"), Data: []byte("GET / HTTP1.1\r\nHeader: 1\r\n\r\n")})
resp := <-producer.Successes()
data, _ := resp.Value.Encode()
if string(data) != "1 2 3\nGET / HTTP1.1\r\nHeader: 1\r\n\r\n" {
t.Error("Message not properly encoded: ", string(data))
t.Errorf("Message not properly encoded: %q", data)
}
}
@@ -42,7 +42,7 @@ func TestOutputKafkaJSON(t *testing.T) {
UseJSON: true,
}, nil)
output.Write([]byte("1 2 3\nGET / HTTP1.1\r\nHeader: 1\r\n\r\n"))
output.PluginWrite(&Message{Meta: []byte("1 2 3\n"), Data: []byte("GET / HTTP1.1\r\nHeader: 1\r\n\r\n")})
resp := <-producer.Successes()
output_null.go
+3 -2
View File
@@ -9,8 +9,9 @@ func NewNullOutput() (o *NullOutput) {
return new(NullOutput)
}
func (o *NullOutput) Write(data []byte) (int, error) {
return len(data), nil
// PluginWrite writes message to this plugin
func (o *NullOutput) PluginWrite(msg *Message) (int, error) {
return len(msg.Data) + len(msg.Meta), nil
}
func (o *NullOutput) String() string {
output_s3.go
+3 -2
View File
@@ -67,8 +67,9 @@ func (o *S3Output) connect() {
}
}
func (o *S3Output) Write(data []byte) (n int, err error) {
return o.buffer.Write(data)
// PluginWrite writes message to this plugin
func (o *S3Output) PluginWrite(msg *Message) (n int, err error) {
return o.buffer.PluginWrite(msg)
}
func (o *S3Output) String() string {
output_tcp.go
+17 -19
View File
@@ -4,7 +4,6 @@ import (
"crypto/tls"
"fmt"
"hash/fnv"
"io"
"net"
"time"
)
@@ -15,7 +14,7 @@ import (
type TCPOutput struct {
address string
limit int
buf []chan []byte
buf []chan *Message
bufStats *GorStat
config *TCPOutputConfig
workerIndex uint32
@@ -31,7 +30,7 @@ type TCPOutputConfig struct {
// NewTCPOutput constructor for TCPOutput
// Initialize X workers which hold keep-alive connection
func NewTCPOutput(address string, config *TCPOutputConfig) io.Writer {
func NewTCPOutput(address string, config *TCPOutputConfig) PluginWriter {
o := new(TCPOutput)
o.address = address
@@ -42,9 +41,9 @@ func NewTCPOutput(address string, config *TCPOutputConfig) io.Writer {
}
// create X buffers and send the buffer index to the worker
o.buf = make([]chan []byte, o.config.Workers)
o.buf = make([]chan *Message, o.config.Workers)
for i := 0; i < o.config.Workers; i++ {
o.buf[i] = make(chan []byte, 100)
o.buf[i] = make(chan *Message, 100)
go o.worker(i)
}
@@ -73,14 +72,16 @@ func (o *TCPOutput) worker(bufferIndex int) {
defer conn.Close()
for {
data := <-o.buf[bufferIndex]
if _, err = conn.Write(data); err == nil {
_, err = conn.Write([]byte(payloadSeparator))
msg := <-o.buf[bufferIndex]
if _, err = conn.Write(msg.Meta); err == nil {
if _, err = conn.Write(msg.Data); err == nil {
_, err = conn.Write(payloadSeparatorAsBytes)
}
}
if err != nil {
Debug(2, "INFO: TCP output connection closed, reconnecting")
o.buf[bufferIndex] <- data
o.buf[bufferIndex] <- msg
go o.worker(bufferIndex)
break
}
@@ -99,23 +100,20 @@ func (o *TCPOutput) getBufferIndex(data []byte) int {
}
func (o *TCPOutput) Write(data []byte) (n int, err error) {
if !isOriginPayload(data) {
return len(data), nil
// PluginWrite writes message to this plugin
func (o *TCPOutput) PluginWrite(msg *Message) (n int, err error) {
if !isOriginPayload(msg.Meta) {
return len(msg.Data), nil
}
// We have to copy, because sending data in multiple threads
newBuf := make([]byte, len(data))
copy(newBuf, data)
bufferIndex := o.getBufferIndex(data)
o.buf[bufferIndex] <- newBuf
bufferIndex := o.getBufferIndex(msg.Data)
o.buf[bufferIndex] <- msg
if Settings.OutputTCPStats {
o.bufStats.Write(len(o.buf[bufferIndex]))
}
return len(data), nil
return len(msg.Data) + len(msg.Meta), nil
}
func (o *TCPOutput) connect(address string) (conn net.Conn, err error) {
output_tcp_test.go
+7 -13
View File
@@ -2,7 +2,6 @@ package main
import (
"bufio"
"io"
"log"
"net"
"sync"
@@ -11,9 +10,7 @@ import (
)
func TestTCPOutput(t *testing.T) {
Settings.Verbose = 2
wg := new(sync.WaitGroup)
quit := make(chan int)
listener := startTCP(func(data []byte) {
wg.Done()
@@ -22,12 +19,11 @@ func TestTCPOutput(t *testing.T) {
output := NewTCPOutput(listener.Addr().String(), &TCPOutputConfig{Workers: 10})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
plugins.All = append(plugins.All, input, output)
emitter := NewEmitter(quit)
emitter := NewEmitter()
go emitter.Start(plugins, Settings.Middleware)
for i := 0; i < 10; i++ {
@@ -68,7 +64,6 @@ func startTCP(cb func([]byte)) net.Listener {
func BenchmarkTCPOutput(b *testing.B) {
wg := new(sync.WaitGroup)
quit := make(chan int)
listener := startTCP(func(data []byte) {
wg.Done()
@@ -82,12 +77,11 @@ func BenchmarkTCPOutput(b *testing.B) {
output := NewTCPOutput(listener.Addr().String(), &TCPOutputConfig{Workers: 10})
plugins := &InOutPlugins{
Inputs: []io.Reader{input},
Outputs: []io.Writer{output},
Inputs: []PluginReader{input},
Outputs: []PluginWriter{output},
}
plugins.All = append(plugins.All, input, output)
emitter := NewEmitter(quit)
emitter := NewEmitter()
// avoid counting above initialization
b.ResetTimer()
go emitter.Start(plugins, Settings.Middleware)
@@ -109,7 +103,7 @@ func TestStickyDisable(t *testing.T) {
func TestBufferDistribution(t *testing.T) {
numberOfWorkers := 10
numberOfMessages := 1000000
numberOfMessages := 10000
percentDistributionErrorRange := 20
buffer := make([]int, numberOfWorkers)
plugins.go
+30 -15
View File
@@ -1,15 +1,36 @@
package main
import (
"io"
"reflect"
"strings"
)
// Message represents data accross plugins
type Message struct {
Meta []byte // metadata
Data []byte // actual data
}
// PluginReader is an interface for input plugins
type PluginReader interface {
PluginRead() (msg *Message, err error)
}
// PluginWriter is an interface for output plugins
type PluginWriter interface {
PluginWrite(msg *Message) (n int, err error)
}
// PluginReadWriter is an interface for plugins that support reading and writing
type PluginReadWriter interface {
PluginReader
PluginWriter
}
// InOutPlugins struct for holding references to plugins
type InOutPlugins struct {
Inputs []io.Reader
Outputs []io.Writer
Inputs []PluginReader
Outputs []PluginWriter
All []interface{}
}
@@ -48,27 +69,21 @@ func (plugins *InOutPlugins) registerPlugin(constructor interface{}, options ...
// Calling our constructor with list of given options
plugin := vc.Call(vo)[0].Interface()
pluginWrapper := plugin
if limit != "" {
pluginWrapper = NewLimiter(plugin, limit)
} else {
pluginWrapper = plugin
plugin = NewLimiter(plugin, limit)
}
_, isR := plugin.(io.Reader)
_, isW := plugin.(io.Writer)
// Some of the output can be Readers as well because return responses
if isR && !isW {
plugins.Inputs = append(plugins.Inputs, pluginWrapper.(io.Reader))
if r, ok := plugin.(PluginReader); ok {
plugins.Inputs = append(plugins.Inputs, r)
}
if isW {
plugins.Outputs = append(plugins.Outputs, pluginWrapper.(io.Writer))
if w, ok := plugin.(PluginWriter); ok {
plugins.Outputs = append(plugins.Outputs, w)
}
plugins.All = append(plugins.All, plugin)
}
// NewPlugins specify and initialize all available plugins
plugins_test.go
+2 -2
View File
@@ -12,8 +12,8 @@ func TestPluginsRegistration(t *testing.T) {
plugins := NewPlugins()
if len(plugins.Inputs) != 2 {
t.Errorf("Should be 2 inputs %d", len(plugins.Inputs))
if len(plugins.Inputs) != 3 {
t.Errorf("Should be 3 inputs got %d", len(plugins.Inputs))
}
if _, ok := plugins.Inputs[0].(*DummyInput); !ok {
protocol.go
+10
View File
@@ -66,6 +66,16 @@ func payloadMeta(payload []byte) [][]byte {
return bytes.Split(payload[:headerSize], []byte{' '})
}
func payloadMetaWithBody(payload []byte) (meta, body []byte) {
if i := bytes.IndexByte(payload, '\n'); i > 0 && len(payload) > i+1 {
meta = payload[:i+1]
body = payload[i+1:]
return
}
// we assume the message did not have meta data
return nil, payload
}
func payloadID(payload []byte) (id []byte) {
meta := payloadMeta(payload)
settings.go
-20
View File
@@ -93,14 +93,10 @@ func init() {
}
flag.BoolVar(&Settings.SplitOutput, "split-output", false, "By default each output gets same traffic. If set to `true` it splits traffic equally among all outputs.")
flag.BoolVar(&Settings.RecognizeTCPSessions, "recognize-tcp-sessions", false, "[PRO] If turned on http output will create separate worker for each TCP session. Splitting output will session based as well.")
flag.Var(&Settings.InputDummy, "input-dummy", "Used for testing outputs. Emits 'Get /' request every 1s")
flag.Var(&Settings.OutputDummy, "output-dummy", "DEPRECATED: use --output-stdout instead")
flag.BoolVar(&Settings.OutputStdout, "output-stdout", false, "Used for testing inputs. Just prints to console data coming from inputs.")
flag.BoolVar(&Settings.OutputNull, "output-null", false, "Used for testing inputs. Drops all requests.")
flag.Var(&Settings.InputTCP, "input-tcp", "Used for internal communication between Gor instances. Example: \n\t# Receive requests from other Gor instances on 28020 port, and redirect output to staging\n\tgor --input-tcp :28020 --output-http staging.com")
@@ -192,32 +188,16 @@ func init() {
flag.StringVar(&Settings.KafkaTLSConfig.ClientKey, "kafka-tls-client-key", "", "Client Key for Kafka TLS Config (mandatory with to kafka-tls-client-cert and kafka-tls-client-key)")
flag.Var(&Settings.ModifierConfig.Headers, "http-set-header", "Inject additional headers to http reqest:\n\tgor --input-raw :8080 --output-http staging.com --http-set-header 'User-Agent: Gor'")
flag.Var(&Settings.ModifierConfig.Headers, "output-http-header", "WARNING: `--output-http-header` DEPRECATED, use `--http-set-header` instead")
flag.Var(&Settings.ModifierConfig.HeaderRewrite, "http-rewrite-header", "Rewrite the request header based on a mapping:\n\tgor --input-raw :8080 --output-http staging.com --http-rewrite-header Host: (.*).example.com,$1.beta.example.com")
flag.Var(&Settings.ModifierConfig.Params, "http-set-param", "Set request url param, if param already exists it will be overwritten:\n\tgor --input-raw :8080 --output-http staging.com --http-set-param api_key=1")
flag.Var(&Settings.ModifierConfig.Methods, "http-allow-method", "Whitelist of HTTP methods to replay. Anything else will be dropped:\n\tgor --input-raw :8080 --output-http staging.com --http-allow-method GET --http-allow-method OPTIONS")
flag.Var(&Settings.ModifierConfig.URLRegexp, "http-allow-url", "A regexp to match requests against. Filter get matched against full url with domain. Anything else will be dropped:\n\t gor --input-raw :8080 --output-http staging.com --http-allow-url ^www.")
flag.Var(&Settings.ModifierConfig.URLNegativeRegexp, "http-disallow-url", "A regexp to match requests against. Filter get matched against full url with domain. Anything else will be forwarded:\n\t gor --input-raw :8080 --output-http staging.com --http-disallow-url ^www.")
flag.Var(&Settings.ModifierConfig.URLRewrite, "http-rewrite-url", "Rewrite the request url based on a mapping:\n\tgor --input-raw :8080 --output-http staging.com --http-rewrite-url /v1/user/([^\\/]+)/ping:/v2/user/$1/ping")
flag.Var(&Settings.ModifierConfig.URLRewrite, "output-http-rewrite-url", "WARNING: `--output-http-rewrite-url` DEPRECATED, use `--http-rewrite-url` instead")
flag.Var(&Settings.ModifierConfig.HeaderFilters, "http-allow-header", "A regexp to match a specific header against. Requests with non-matching headers will be dropped:\n\t gor --input-raw :8080 --output-http staging.com --http-allow-header api-version:^v1")
flag.Var(&Settings.ModifierConfig.HeaderFilters, "output-http-header-filter", "WARNING: `--output-http-header-filter` DEPRECATED, use `--http-allow-header` instead")
flag.Var(&Settings.ModifierConfig.HeaderNegativeFilters, "http-disallow-header", "A regexp to match a specific header against. Requests with matching headers will be dropped:\n\t gor --input-raw :8080 --output-http staging.com --http-disallow-header \"User-Agent: Replayed by Gor\"")
flag.Var(&Settings.ModifierConfig.HeaderBasicAuthFilters, "http-basic-auth-filter", "A regexp to match the decoded basic auth string against. Requests with non-matching headers will be dropped:\n\t gor --input-raw :8080 --output-http staging.com --http-basic-auth-filter \"^customer[0-9].*\"")
flag.Var(&Settings.ModifierConfig.HeaderHashFilters, "http-header-limiter", "Takes a fraction of requests, consistently taking or rejecting a request based on the FNV32-1A hash of a specific header:\n\t gor --input-raw :8080 --output-http staging.com --http-header-limiter user-id:25%")
flag.Var(&Settings.ModifierConfig.HeaderHashFilters, "output-http-header-hash-filter", "WARNING: `output-http-header-hash-filter` DEPRECATED, use `--http-header-hash-limiter` instead")
flag.Var(&Settings.ModifierConfig.ParamHashFilters, "http-param-limiter", "Takes a fraction of requests, consistently taking or rejecting a request based on the FNV32-1A hash of a specific GET param:\n\t gor --input-raw :8080 --output-http staging.com --http-param-limiter user_id:25%")
// default values, using for tests
test_input.go
+8 -9
View File
@@ -25,22 +25,21 @@ func NewTestInput() (i *TestInput) {
return
}
func (i *TestInput) Read(data []byte) (int, error) {
// PluginRead reads message from this plugin
func (i *TestInput) PluginRead() (*Message, error) {
var msg Message
select {
case buf := <-i.data:
var header []byte
msg.Data = buf
if !i.skipHeader {
header = payloadHeader(RequestPayload, uuid(), time.Now().UnixNano(), -1)
copy(data[0:len(header)], header)
copy(data[len(header):], buf)
msg.Meta = payloadHeader(RequestPayload, uuid(), time.Now().UnixNano(), -1)
} else {
copy(data, buf)
msg.Meta, msg.Data = payloadMetaWithBody(msg.Data)
}
return len(buf) + len(header), nil
return &msg, nil
case <-i.stop:
return 0, ErrorStopped
return nil, ErrorStopped
}
}
test_output.go
+9 -8
View File
@@ -1,6 +1,6 @@
package main
type writeCallback func(data []byte)
type writeCallback func(*Message)
// TestOutput used in testing to intercept any output into callback
type TestOutput struct {
@@ -8,19 +8,20 @@ type TestOutput struct {
}
// NewTestOutput constructor for TestOutput, accepts callback which get called on each incoming Write
func NewTestOutput(cb writeCallback) (i *TestOutput) {
i = new(TestOutput)
func NewTestOutput(cb writeCallback) PluginWriter {
i := new(TestOutput)
i.cb = cb
return
return i
}
func (i *TestOutput) Write(data []byte) (int, error) {
i.cb(data)
// PluginWrite write message to this plugin
func (i *TestOutput) PluginWrite(msg *Message) (int, error) {
i.cb(msg)
return len(data), nil
return len(msg.Data) + len(msg.Meta), nil
}
func (i *TestOutput) String() string {
return "Test Input"
return "Test Output"
}