mirror of
https://github.com/buger/goreplay.git
synced 2026-10-19 16:15:42 +00:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5734ddba20 | ||
|
|
1558c1d9bc | ||
|
|
a71245be23 | ||
|
|
40a3c51c9c | ||
|
|
5e2bbf8228 | ||
|
|
a3b6be8b82 | ||
|
|
3aec926ad9 | ||
|
|
ec4dbf2061 | ||
|
|
6d812ceb7f | ||
|
|
20435af11b | ||
|
|
f719071358 | ||
|
|
716fb72e27 | ||
|
|
3fcdb6f1bc | ||
|
|
715ce88a63 | ||
|
|
bcf81c6410 | ||
|
|
358cb78a53 | ||
|
|
48aa326ba6 | ||
|
|
5d8ca525a4 | ||
|
|
3635d66a76 | ||
|
|
a4e78cf22c | ||
|
|
8aedc15027 | ||
|
|
30d66d31de | ||
|
|
b1d7c4134b | ||
|
|
66e680a56f | ||
|
|
a4ed67acd6 | ||
|
|
ee9ff8163b | ||
|
|
9519b9a9f5 | ||
|
|
3401dfbf2d | ||
|
|
cfa14b0a32 | ||
|
|
85065cb41a | ||
|
|
7100ac3593 | ||
|
|
d5184c6cbd | ||
|
|
4d7d4671c4 | ||
|
|
8c668d252c | ||
|
|
887293230b |
+1
-1
@@ -3,7 +3,7 @@ dist: focal
|
||||
language: go
|
||||
go: 1.14
|
||||
|
||||
script: sudo -E bash -c "source /etc/profile && eval '$(gimme 1.14)' && go test ./... -v -timeout 120s"
|
||||
script: sudo -E bash -c "source /etc/profile && eval '$(gimme 1.14)' && go test ./... -v -timeout 120s -race"
|
||||
|
||||
before_install:
|
||||
- sudo apt-get install libpcap-dev -y
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
FROM alpine:latest as builder
|
||||
RUN apk add --no-cache ca-certificates openssl
|
||||
RUN wget https://github.com/buger/goreplay/releases/download/v1.0.0/gor_1.0.0_x64.tar.gz -O gor.tar.gz
|
||||
RUN wget https://github.com/buger/goreplay/releases/download/v1.2.0/gor_v1.2.0_x64.tar.gz -O gor.tar.gz
|
||||
RUN tar xzf gor.tar.gz
|
||||
|
||||
FROM scratch
|
||||
|
||||
@@ -130,7 +130,7 @@ bash:
|
||||
|
||||
|
||||
FPMCOMMON= \
|
||||
--name goreplay \
|
||||
--name gor \
|
||||
--description "GoReplay is an open-source network monitoring tool which can record your live traffic, and use it for shadowing, load testing, monitoring and detailed analysis." \
|
||||
-v $(VERSION) \
|
||||
--vendor "Leonid Bugaev" \
|
||||
|
||||
@@ -2,7 +2,6 @@
|
||||
package byteutils
|
||||
|
||||
import (
|
||||
"reflect"
|
||||
"unsafe"
|
||||
)
|
||||
|
||||
@@ -48,9 +47,6 @@ func Replace(a []byte, from, to int, new []byte) []byte {
|
||||
}
|
||||
|
||||
// SliceToString preferred for large body payload (zero allocation and faster)
|
||||
func SliceToString(buf *[]byte, s *string) {
|
||||
bHeader := (*reflect.SliceHeader)(unsafe.Pointer(buf))
|
||||
sHeader := (*reflect.StringHeader)(unsafe.Pointer(s))
|
||||
sHeader.Data = bHeader.Data
|
||||
sHeader.Len = bHeader.Len
|
||||
func SliceToString(buf []byte) string {
|
||||
return *(*string)(unsafe.Pointer(&buf))
|
||||
}
|
||||
|
||||
@@ -32,9 +32,10 @@ func TestReplace(t *testing.T) {
|
||||
}
|
||||
|
||||
func BenchmarkStringtoSlice(b *testing.B) {
|
||||
b.StopTimer()
|
||||
buf := make([]byte, b.N)
|
||||
b.StartTimer()
|
||||
s := new(string)
|
||||
SliceToString(&buf, s)
|
||||
var s string
|
||||
var buf [1 << 20]byte
|
||||
for i := 0; i < b.N; i++ {
|
||||
s = SliceToString(buf[:])
|
||||
}
|
||||
_ = s // avoid gc to optimize away the loop body
|
||||
}
|
||||
|
||||
+83
-20
@@ -6,6 +6,7 @@ import (
|
||||
"io/ioutil"
|
||||
"net"
|
||||
"os"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -257,41 +258,103 @@ func BenchmarkPcapFile(b *testing.B) {
|
||||
b.Logf("%d/%d packets in %s", pckts, b.N, time.Since(now))
|
||||
}
|
||||
|
||||
func BenchmarkPcap(b *testing.B) {
|
||||
now := time.Now()
|
||||
var err error
|
||||
// used to benchmark sock engine
|
||||
var buf [1024]byte
|
||||
|
||||
l, err := NewListener(LoopBack.Name, 8000, "", EnginePcap, true)
|
||||
func init() {
|
||||
for i := 0; i < len(buf); i++ {
|
||||
buf[i] = 0xff
|
||||
}
|
||||
}
|
||||
|
||||
func handler(n, counter *int32) Handler {
|
||||
return func(p gopacket.Packet) {
|
||||
nn := int32(len(p.Data()))
|
||||
atomic.AddInt32(n, nn)
|
||||
atomic.AddInt32(counter, 1)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkPcap(b *testing.B) {
|
||||
var err error
|
||||
n := new(int32)
|
||||
counter := new(int32)
|
||||
l, err := NewListener(LoopBack.Name, 8000, "", EnginePcap, false)
|
||||
if err != nil {
|
||||
b.Errorf("expected error to be nil, got %v", err)
|
||||
b.Error(err)
|
||||
return
|
||||
}
|
||||
l.PcapOptions.BPFFilter = "udp dst port 8000 and host 127.0.0.1"
|
||||
err = l.Activate()
|
||||
if err != nil {
|
||||
b.Errorf("expected error to be nil, got %v", err)
|
||||
b.Error(err)
|
||||
return
|
||||
}
|
||||
defer l.Handles[LoopBack.Name].(*pcap.Handle).Close()
|
||||
for i := 0; i < b.N; i++ {
|
||||
_, _ = net.Dial("tcp", "127.0.0.1:8000")
|
||||
errCh := l.ListenBackground(context.Background(), handler(n, counter))
|
||||
select {
|
||||
case <-l.Reading:
|
||||
case err = <-errCh:
|
||||
b.Error(err)
|
||||
return
|
||||
}
|
||||
sts, _ := l.Handles[LoopBack.Name].(*pcap.Handle).Stats()
|
||||
b.Logf("%d packets in %s", sts.PacketsReceived, time.Since(now))
|
||||
var conn net.Conn
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
b.StopTimer()
|
||||
conn, err = net.Dial("udp", "127.0.0.1:8000")
|
||||
if err != nil {
|
||||
b.Error(err)
|
||||
return
|
||||
}
|
||||
b.StartTimer()
|
||||
_, err = conn.Write(buf[:])
|
||||
if err != nil {
|
||||
b.Error(err)
|
||||
return
|
||||
}
|
||||
}
|
||||
b.ReportMetric(float64(atomic.LoadInt32(n)), "buf")
|
||||
b.ReportMetric(float64(atomic.LoadInt32(counter)), "packets")
|
||||
}
|
||||
|
||||
func BenchmarkRawSocket(b *testing.B) {
|
||||
now := time.Now()
|
||||
var err error
|
||||
|
||||
l, err := NewListener(LoopBack.Name, 0, "", EngineRawSocket, true)
|
||||
err = l.Activate()
|
||||
n := new(int32)
|
||||
counter := new(int32)
|
||||
l, err := NewListener(LoopBack.Name, 8000, "", EngineRawSocket, false)
|
||||
if err != nil {
|
||||
b.Error(err)
|
||||
return
|
||||
}
|
||||
defer l.Handles[LoopBack.Name].(*SockRaw).Close()
|
||||
for i := 0; i < b.N; i++ {
|
||||
_, _ = net.Dial("tcp", "127.0.0.1:8000")
|
||||
l.PcapOptions.BPFFilter = "udp dst port 8000 and host 127.0.0.1"
|
||||
err = l.Activate()
|
||||
if err != nil {
|
||||
b.Error(err)
|
||||
return
|
||||
}
|
||||
sts, _ := l.Handles[LoopBack.Name].(*SockRaw).Stats()
|
||||
b.Logf("%d packets in %s", sts.Packets, time.Since(now))
|
||||
errCh := l.ListenBackground(context.Background(), handler(n, counter))
|
||||
select {
|
||||
case <-l.Reading:
|
||||
case err = <-errCh:
|
||||
b.Error(err)
|
||||
return
|
||||
}
|
||||
var conn net.Conn
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
b.StopTimer()
|
||||
conn, err = net.Dial("udp", "127.0.0.1:8000")
|
||||
if err != nil {
|
||||
b.Error(err)
|
||||
return
|
||||
}
|
||||
b.StartTimer()
|
||||
_, err = conn.Write(buf[:])
|
||||
if err != nil {
|
||||
b.Error(err)
|
||||
return
|
||||
}
|
||||
}
|
||||
b.ReportMetric(float64(atomic.LoadInt32(n)), "buf")
|
||||
b.ReportMetric(float64(atomic.LoadInt32(counter)), "packets")
|
||||
}
|
||||
|
||||
+1
-1
@@ -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)),
|
||||
|
||||
+93
-124
@@ -1,35 +1,46 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"hash/fnv"
|
||||
"io"
|
||||
"log"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/buger/goreplay/byteutils"
|
||||
"github.com/buger/goreplay/stat"
|
||||
)
|
||||
|
||||
type emitter struct {
|
||||
sync.Mutex
|
||||
sync.WaitGroup
|
||||
quit chan int
|
||||
plugins *InOutPlugins
|
||||
// EmitterConfig configurations associated with emitters
|
||||
type EmitterConfig struct {
|
||||
Stats bool
|
||||
SplitOutput bool
|
||||
RecognizeTCPSessions bool
|
||||
}
|
||||
|
||||
// NewEmitter creates and initializes new `emitter` object.
|
||||
func NewEmitter(quit chan int) *emitter {
|
||||
return &emitter{
|
||||
quit: quit,
|
||||
}
|
||||
// Emitter represents an abject to manage plugins communication
|
||||
type Emitter struct {
|
||||
sync.WaitGroup
|
||||
plugins *InOutPlugins
|
||||
st *stat.Writer
|
||||
}
|
||||
|
||||
// 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) {
|
||||
defer e.Wait()
|
||||
func (e *Emitter) Start(plugins *InOutPlugins, middlewareCmd string) {
|
||||
if Settings.CopyBufferSize < 1 {
|
||||
Settings.CopyBufferSize = 5 << 20
|
||||
}
|
||||
e.plugins = plugins
|
||||
if Settings.EmitterConfig.Stats {
|
||||
e.st = initializeStatWriter("meta.csv", "type: GOR_STAT\ntitle: Messages Meta")
|
||||
statsWriteRecord(e.st, []string{"type", "id", "timestamp", "latency"}, "[EMITTER]")
|
||||
}
|
||||
|
||||
if middlewareCmd != "" {
|
||||
middleware := NewMiddleware(middlewareCmd)
|
||||
@@ -38,147 +49,102 @@ 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 := e.CopyMulty(middleware, plugins.Outputs...); err != nil {
|
||||
Debug(1, 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 := e.CopyMulty(in, plugins.Outputs...); err != nil {
|
||||
Debug(1, 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()
|
||||
}
|
||||
if len(e.plugins.All) > 0 {
|
||||
// wait for everything to stop
|
||||
e.Wait()
|
||||
}
|
||||
e.plugins.All = nil // avoid Close to make changes again
|
||||
if e.st != nil {
|
||||
e.st.Close()
|
||||
}
|
||||
e.plugins = nil // avoid further accidental usage
|
||||
}
|
||||
|
||||
// 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 (e *Emitter) 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)
|
||||
|
||||
if e.st != nil {
|
||||
statsWriteRecord(e.st, splitUnsafe(meta), "[EMITTER]")
|
||||
}
|
||||
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,27 +155,23 @@ 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
|
||||
}
|
||||
|
||||
wIndex++
|
||||
|
||||
if wIndex >= len(writers) {
|
||||
wIndex = 0
|
||||
}
|
||||
wIndex = (wIndex + 1) % len(writers)
|
||||
}
|
||||
} else {
|
||||
for _, dst := range writers {
|
||||
if _, err := dst.Write(payload); err != nil {
|
||||
if _, err := dst.PluginWrite(msg); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
@@ -217,19 +179,26 @@ 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++
|
||||
}
|
||||
}
|
||||
|
||||
func splitUnsafe(data [][]byte) []string {
|
||||
s := make([]string, len(data), len(data))
|
||||
for i := range data {
|
||||
s[i] = byteutils.SliceToString(data[i])
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
+35
-43
@@ -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()
|
||||
|
||||
@@ -4,7 +4,6 @@ package main
|
||||
|
||||
import (
|
||||
"flag"
|
||||
"fmt"
|
||||
"log"
|
||||
"net/http"
|
||||
"net/http/httputil"
|
||||
@@ -74,13 +73,13 @@ 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)
|
||||
|
||||
time.AfterFunc(Settings.ExitAfter, func() {
|
||||
fmt.Printf("gor run timeout %s\n", Settings.ExitAfter)
|
||||
log.Printf("gor run timeout %s\n", Settings.ExitAfter)
|
||||
close(closeCh)
|
||||
})
|
||||
}
|
||||
|
||||
-64
@@ -1,64 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"runtime"
|
||||
"strconv"
|
||||
"time"
|
||||
)
|
||||
|
||||
type GorStat struct {
|
||||
statName string
|
||||
rateMs int
|
||||
latest int
|
||||
mean int
|
||||
max int
|
||||
count int
|
||||
}
|
||||
|
||||
func NewGorStat(statName string, rateMs int) (s *GorStat) {
|
||||
s = new(GorStat)
|
||||
s.statName = statName
|
||||
s.rateMs = rateMs
|
||||
s.latest = 0
|
||||
s.mean = 0
|
||||
s.max = 0
|
||||
s.count = 0
|
||||
|
||||
if Settings.Stats {
|
||||
go s.reportStats()
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (s *GorStat) Write(latest int) {
|
||||
if Settings.Stats {
|
||||
if latest > s.max {
|
||||
s.max = latest
|
||||
}
|
||||
if latest != 0 {
|
||||
s.mean = ((s.mean * s.count) + latest) / (s.count + 1)
|
||||
}
|
||||
s.latest = latest
|
||||
s.count = s.count + 1
|
||||
}
|
||||
}
|
||||
|
||||
func (s *GorStat) Reset() {
|
||||
s.latest = 0
|
||||
s.max = 0
|
||||
s.mean = 0
|
||||
s.count = 0
|
||||
}
|
||||
|
||||
func (s *GorStat) String() string {
|
||||
return s.statName + ":" + strconv.Itoa(s.latest) + "," + strconv.Itoa(s.mean) + "," + strconv.Itoa(s.max) + "," + strconv.Itoa(s.count) + "," + strconv.Itoa(s.count/(s.rateMs/1000.0)) + "," + strconv.Itoa(runtime.NumGoroutine())
|
||||
}
|
||||
|
||||
func (s *GorStat) reportStats() {
|
||||
Debug(0, "\n", s.statName+":latest,mean,max,count,count/second,gcount")
|
||||
for {
|
||||
Debug(0, "\n", s)
|
||||
s.Reset()
|
||||
time.Sleep(time.Duration(s.rateMs) * time.Millisecond)
|
||||
}
|
||||
}
|
||||
-521
@@ -1,521 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"crypto/tls"
|
||||
"encoding/base64"
|
||||
"io"
|
||||
"log"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httputil"
|
||||
"net/url"
|
||||
"runtime/debug"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
"github.com/buger/goreplay/proto"
|
||||
)
|
||||
|
||||
var httpMu sync.Mutex
|
||||
|
||||
const (
|
||||
readChunkSize = 64 * 1024
|
||||
maxResponseSize = 1073741824
|
||||
)
|
||||
|
||||
var chunkedSuffix = []byte("0\r\n\r\n")
|
||||
|
||||
var defaultPorts = map[string]string{
|
||||
"http": "80",
|
||||
"https": "443",
|
||||
}
|
||||
|
||||
type HTTPClientConfig struct {
|
||||
FollowRedirects int
|
||||
Debug bool
|
||||
OriginalHost bool
|
||||
ConnectionTimeout time.Duration
|
||||
Timeout time.Duration
|
||||
ResponseBufferSize int
|
||||
CompatibilityMode bool
|
||||
}
|
||||
|
||||
type HTTPClient struct {
|
||||
baseURL string
|
||||
scheme string
|
||||
host string
|
||||
auth string
|
||||
conn net.Conn
|
||||
proxy *url.URL
|
||||
proxyAuth string
|
||||
respBuf []byte
|
||||
config *HTTPClientConfig
|
||||
goClient *http.Client
|
||||
redirectsCount int
|
||||
}
|
||||
|
||||
func NewHTTPClient(baseURL string, config *HTTPClientConfig) *HTTPClient {
|
||||
if !strings.HasPrefix(baseURL, "http") {
|
||||
baseURL = "http://" + baseURL
|
||||
}
|
||||
|
||||
u, _ := url.Parse(baseURL)
|
||||
|
||||
if config.Timeout == 0 {
|
||||
config.Timeout = time.Second
|
||||
}
|
||||
|
||||
config.ConnectionTimeout = config.Timeout
|
||||
|
||||
if config.ResponseBufferSize == 0 {
|
||||
config.ResponseBufferSize = 100 * 1024 // 100kb
|
||||
}
|
||||
|
||||
client := new(HTTPClient)
|
||||
client.baseURL = u.String()
|
||||
client.host = u.Host
|
||||
client.scheme = u.Scheme
|
||||
client.respBuf = make([]byte, config.ResponseBufferSize)
|
||||
client.config = config
|
||||
|
||||
if config.CompatibilityMode {
|
||||
client.goClient = &http.Client{
|
||||
// #TODO
|
||||
// CheckRedirect: redirectPolicyFunc,
|
||||
}
|
||||
}
|
||||
|
||||
if u.User != nil {
|
||||
client.auth = "Basic " + base64.StdEncoding.EncodeToString([]byte(u.User.String()))
|
||||
}
|
||||
|
||||
client.proxy, _ = http.ProxyFromEnvironment(&http.Request{URL: u})
|
||||
|
||||
if client.isProxy() && client.proxy.User != nil {
|
||||
client.proxyAuth = "Basic " + base64.StdEncoding.EncodeToString([]byte(client.proxy.User.String()))
|
||||
}
|
||||
|
||||
return client
|
||||
}
|
||||
|
||||
func (c *HTTPClient) Connect() (err error) {
|
||||
c.Disconnect()
|
||||
|
||||
var toDial string
|
||||
if !strings.Contains(c.host, ":") {
|
||||
toDial = c.host + ":" + defaultPorts[c.scheme]
|
||||
} else {
|
||||
toDial = c.host
|
||||
}
|
||||
|
||||
if c.isProxy() {
|
||||
if c.proxy.Scheme != "http" {
|
||||
panic("Unsupported HTTP Proxy method")
|
||||
}
|
||||
Debug(3, "[HTTPClient] Connecting to proxy", c.proxy.String(), "<>", toDial)
|
||||
c.conn, err = net.DialTimeout("tcp", c.proxy.Host, c.config.ConnectionTimeout)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if c.scheme == "https" {
|
||||
c.conn.Write([]byte("CONNECT " + toDial + " HTTP/1.1\r\n"))
|
||||
if c.proxyAuth != "" {
|
||||
c.conn.Write([]byte("Proxy-Authorization: " + c.proxyAuth + "\r\n"))
|
||||
}
|
||||
c.conn.Write([]byte("\r\n"))
|
||||
br := bufio.NewReader(c.conn)
|
||||
l, _, err := br.ReadLine()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(l) < 12 {
|
||||
panic("HTTP proxy did not respond correctly")
|
||||
}
|
||||
status := l[9:12]
|
||||
if !bytes.Equal(status, []byte("200")) {
|
||||
panic("HTTP proxy did not respond correctly")
|
||||
}
|
||||
for {
|
||||
// Read until we find the empty line
|
||||
l, _, err := br.ReadLine()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if len(l) == 0 {
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
Debug(3, "[HTTPClient] Proxy successfully connected")
|
||||
} else {
|
||||
c.conn, err = net.DialTimeout("tcp", toDial, c.config.ConnectionTimeout)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if c.scheme == "https" {
|
||||
// Wrap our socket in TLS
|
||||
Debug(3, "[HTTPClient] Wrapping socket in TLS", c.host)
|
||||
tlsConn := tls.Client(c.conn, &tls.Config{InsecureSkipVerify: true, ServerName: c.host})
|
||||
|
||||
if err = tlsConn.Handshake(); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
c.conn = tlsConn
|
||||
Debug(3, "[HTTPClient] Successfully wrapped in TLS")
|
||||
}
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
func (c *HTTPClient) Disconnect() {
|
||||
if c.conn != nil {
|
||||
c.conn.Close()
|
||||
c.conn = nil
|
||||
Debug(3, "[HTTP] Disconnected: ", c.baseURL)
|
||||
}
|
||||
}
|
||||
|
||||
func (c *HTTPClient) isAlive(readBytes *int) bool {
|
||||
// Ready 1 byte from socket without timeout to check if it not closed
|
||||
c.conn.SetReadDeadline(time.Now().Add(time.Millisecond))
|
||||
n, err := c.conn.Read(c.respBuf[:1])
|
||||
|
||||
if err == io.EOF {
|
||||
Debug(3, "[HTTPClient] connection closed, reconnecting")
|
||||
return false
|
||||
}
|
||||
|
||||
if err == syscall.EPIPE {
|
||||
Debug(3, "Detected broken pipe.", err)
|
||||
return false
|
||||
}
|
||||
if n != 0 {
|
||||
*readBytes += n
|
||||
Debug(3, "[HTTPClient] isAlive readBytes ", *readBytes)
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func (c *HTTPClient) SendGoClient(data []byte) ([]byte, error) {
|
||||
var req *http.Request
|
||||
var resp *http.Response
|
||||
var err error
|
||||
|
||||
req, err = http.ReadRequest(bufio.NewReader(bytes.NewBuffer(data)))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if !c.config.OriginalHost {
|
||||
req.Host = c.host
|
||||
}
|
||||
|
||||
if c.auth != "" {
|
||||
req.Header.Add("Authorization", c.auth)
|
||||
}
|
||||
|
||||
req.URL, _ = url.ParseRequestURI(c.scheme + "://" + c.host + req.RequestURI)
|
||||
req.RequestURI = ""
|
||||
|
||||
resp, err = c.goClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return httputil.DumpResponse(resp, true)
|
||||
}
|
||||
|
||||
func (c *HTTPClient) Send(data []byte) (response []byte, err error) {
|
||||
// Don't exit on panic
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
Debug(3, "[HTTPClient]", r, string(data))
|
||||
|
||||
if _, ok := r.(error); ok {
|
||||
log.Println("[HTTPClient] Failed to send request: ", string(data))
|
||||
log.Println("[HTTPClient] Response: ", string(response))
|
||||
log.Println("PANIC: pkg:", r, string(debug.Stack()))
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
if c.config.CompatibilityMode {
|
||||
return c.SendGoClient(data)
|
||||
}
|
||||
|
||||
var readBytes int
|
||||
if c.conn == nil || !c.isAlive(&readBytes) {
|
||||
Debug(3, "[HTTPClient] Connecting:", c.baseURL)
|
||||
if err = c.Connect(); err != nil {
|
||||
Debug(1, "[HTTPClient] Connection error:", err)
|
||||
response = errorPayload(HTTP_CONNECTION_ERROR)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
timeout := time.Now().Add(c.config.Timeout)
|
||||
|
||||
c.conn.SetWriteDeadline(timeout)
|
||||
|
||||
if !c.config.OriginalHost {
|
||||
data = proto.SetHost(data, []byte(c.baseURL), []byte(c.host))
|
||||
}
|
||||
|
||||
if c.isProxy() && c.scheme == "http" {
|
||||
path := proto.Path(data)
|
||||
if len(path) > 0 && path[0] == '/' {
|
||||
data = proto.SetPath(data, c.proxyPath(path))
|
||||
if c.proxyAuth != "" {
|
||||
data = proto.SetHeader(data, []byte("Proxy-Authorization"), []byte(c.proxyAuth))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if c.auth != "" {
|
||||
data = proto.SetHeader(data, []byte("Authorization"), []byte(c.auth))
|
||||
}
|
||||
|
||||
if c.config.Debug {
|
||||
Debug(3, "[HTTPClient] Sending:", string(data))
|
||||
}
|
||||
|
||||
return c.send(data, readBytes, timeout)
|
||||
}
|
||||
|
||||
func (c *HTTPClient) send(data []byte, readBytes int, timeout time.Time) (response []byte, err error) {
|
||||
var payload []byte
|
||||
var n int
|
||||
if _, err = c.conn.Write(data); err != nil {
|
||||
Debug(1, "[HTTPClient] Write error:", err, c.baseURL)
|
||||
response = errorPayload(HTTP_TIMEOUT)
|
||||
c.Disconnect()
|
||||
return
|
||||
}
|
||||
|
||||
var currentChunk []byte
|
||||
timeout = time.Now().Add(c.config.Timeout)
|
||||
chunked := false
|
||||
contentLength := -1
|
||||
currentContentLength := 0
|
||||
chunks := 0
|
||||
|
||||
for {
|
||||
c.conn.SetReadDeadline(timeout)
|
||||
|
||||
if readBytes < len(c.respBuf) {
|
||||
n, err = c.conn.Read(c.respBuf[readBytes:])
|
||||
readBytes += n
|
||||
chunks++
|
||||
|
||||
// First chunk
|
||||
if chunked || contentLength != -1 {
|
||||
currentContentLength += n
|
||||
} else {
|
||||
// If headers are finished
|
||||
var firstEmptyLine = bytes.Index(c.respBuf[:readBytes], proto.EmptyLine)
|
||||
if firstEmptyLine != -1 {
|
||||
if bytes.Equal(proto.Header(c.respBuf[:readBytes], []byte("Transfer-Encoding")), []byte("chunked")) {
|
||||
chunked = true
|
||||
} else {
|
||||
status, _ := strconv.Atoi(string(proto.Status(c.respBuf[:readBytes])))
|
||||
// We want to soak up all 100 Continues received to get the real result code
|
||||
if status >= 100 && status < 200 {
|
||||
timeout = time.Now().Add(c.config.Timeout)
|
||||
var deleteLen = firstEmptyLine + len(proto.EmptyLine)
|
||||
copy(c.respBuf, c.respBuf[deleteLen:readBytes])
|
||||
readBytes -= deleteLen
|
||||
chunks--
|
||||
continue
|
||||
} else if status == 204 || status == 304 {
|
||||
contentLength = 0
|
||||
break
|
||||
} else {
|
||||
l := proto.Header(c.respBuf[:readBytes], []byte("Content-Length"))
|
||||
if len(l) > 0 {
|
||||
contentLength, _ = strconv.Atoi(string(l))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
currentContentLength += len(proto.Body(c.respBuf[:readBytes]))
|
||||
}
|
||||
}
|
||||
|
||||
if chunked {
|
||||
// Check if chunked message finished
|
||||
if bytes.HasSuffix(c.respBuf[:readBytes], chunkedSuffix) {
|
||||
break
|
||||
}
|
||||
} else if contentLength != -1 {
|
||||
if currentContentLength > contentLength {
|
||||
Debug(3, "[HTTPClient] disconnected, wrong length", currentContentLength, contentLength)
|
||||
c.Disconnect()
|
||||
break
|
||||
} else if currentContentLength == contentLength {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
if err == io.EOF {
|
||||
err = nil
|
||||
}
|
||||
break
|
||||
}
|
||||
} else {
|
||||
if currentChunk == nil {
|
||||
currentChunk = make([]byte, readChunkSize)
|
||||
}
|
||||
|
||||
n, err = c.conn.Read(currentChunk)
|
||||
|
||||
readBytes += int(n)
|
||||
chunks++
|
||||
currentContentLength += n
|
||||
|
||||
if chunked {
|
||||
// Check if chunked message finished
|
||||
if bytes.HasSuffix(currentChunk[:n], chunkedSuffix) {
|
||||
break
|
||||
}
|
||||
} else if contentLength != -1 {
|
||||
if currentContentLength > contentLength {
|
||||
Debug(3, "[HTTPClient] disconnected, wrong length", currentContentLength, contentLength)
|
||||
c.Disconnect()
|
||||
break
|
||||
} else if currentContentLength == contentLength {
|
||||
break
|
||||
}
|
||||
} else {
|
||||
Debug(3, "[HTTPClient] disconnected, can't find Content-Length or Chunked")
|
||||
c.Disconnect()
|
||||
break
|
||||
}
|
||||
|
||||
if err == io.EOF {
|
||||
break
|
||||
} else if err != nil {
|
||||
Debug(3, "[HTTPClient] Read the whole body error:", err, c.baseURL)
|
||||
break
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
if readBytes >= maxResponseSize {
|
||||
Debug(3, "[HTTPClient] Body is more than the max size", maxResponseSize,
|
||||
c.baseURL)
|
||||
break
|
||||
}
|
||||
|
||||
// For following chunks expect less timeout
|
||||
timeout = time.Now().Add(c.config.Timeout / 5)
|
||||
}
|
||||
|
||||
if err != nil && readBytes == 0 {
|
||||
Debug(3, "[HTTPClient] Response read timeout error", err, c.conn, readBytes, string(c.respBuf[:readBytes]))
|
||||
response = errorPayload(HTTP_TIMEOUT)
|
||||
c.Disconnect()
|
||||
return
|
||||
}
|
||||
|
||||
if readBytes < 4 || string(c.respBuf[:4]) != "HTTP" {
|
||||
maxRead := 100
|
||||
if readBytes < maxRead {
|
||||
maxRead = readBytes
|
||||
}
|
||||
Debug(3, "[HTTPClient] Response read unknown error", err, c.conn, readBytes, string(c.respBuf[:maxRead]))
|
||||
response = errorPayload(HTTP_UNKNOWN_ERROR)
|
||||
c.Disconnect()
|
||||
return
|
||||
}
|
||||
|
||||
if readBytes > len(c.respBuf) {
|
||||
readBytes = len(c.respBuf)
|
||||
}
|
||||
payload = make([]byte, readBytes)
|
||||
copy(payload, c.respBuf[:readBytes])
|
||||
|
||||
if c.config.Debug {
|
||||
Debug(3, "[HTTPClient] Received:", string(payload))
|
||||
}
|
||||
|
||||
if c.config.FollowRedirects > 0 && c.redirectsCount < c.config.FollowRedirects {
|
||||
status := payload[9:12]
|
||||
|
||||
// 3xx requests
|
||||
if status[0] == '3' {
|
||||
c.redirectsCount++
|
||||
|
||||
location := proto.Header(payload, []byte("Location"))
|
||||
redirectPayload := proto.SetPath(data, location)
|
||||
|
||||
Debug(3, "[HTTPClient] Redirecting to: "+string(location))
|
||||
|
||||
return c.Send(redirectPayload)
|
||||
}
|
||||
}
|
||||
|
||||
if bytes.Equal(proto.Status(payload), []byte("400")) {
|
||||
Debug(3, "[HTTPClient] Closed connection on 400 response")
|
||||
c.Disconnect()
|
||||
}
|
||||
|
||||
c.redirectsCount = 0
|
||||
|
||||
return payload, err
|
||||
}
|
||||
|
||||
func (c *HTTPClient) Get(path string) (response []byte, err error) {
|
||||
payload := "GET " + path + " HTTP/1.1\r\n\r\n"
|
||||
|
||||
return c.Send([]byte(payload))
|
||||
}
|
||||
|
||||
func (c *HTTPClient) Post(path string, body []byte) (response []byte, err error) {
|
||||
payload := "POST " + path + " HTTP/1.1\r\n"
|
||||
payload += "Content-Length: " + strconv.Itoa(len(body)) + "\r\n\r\n"
|
||||
payload += string(body)
|
||||
|
||||
return c.Send([]byte(payload))
|
||||
}
|
||||
|
||||
func (c *HTTPClient) proxyPath(path []byte) []byte {
|
||||
return append([]byte(c.scheme+"://"+c.host), path...)
|
||||
}
|
||||
|
||||
func (c *HTTPClient) isProxy() bool {
|
||||
return c.proxy != nil
|
||||
}
|
||||
|
||||
const (
|
||||
// https://support.cloudflare.com/hc/en-us/articles/200171936-Error-520-Web-server-is-returning-an-unknown-error
|
||||
HTTP_UNKNOWN_ERROR = "520"
|
||||
// https://support.cloudflare.com/hc/en-us/articles/200171916-Error-521-Web-server-is-down
|
||||
HTTP_CONNECTION_ERROR = "521"
|
||||
// https://support.cloudflare.com/hc/en-us/articles/200171906-Error-522-Connection-timed-out
|
||||
HTTP_CONNECTION_TIMEOUT = "522"
|
||||
// https://support.cloudflare.com/hc/en-us/articles/200171946-Error-523-Origin-is-unreachable
|
||||
HTTP_UNREACHABLE = "523"
|
||||
// https://support.cloudflare.com/hc/en-us/articles/200171926-Error-524-A-timeout-occurred
|
||||
HTTP_TIMEOUT = "524"
|
||||
)
|
||||
|
||||
var errorPayloadTemplate = "HTTP/1.1 202 Accepted\r\nDate: Mon, 17 Aug 2015 14:10:11 GMT\r\nContent-Length: 0\r\nContent-Type: text/plain; charset=utf-8\r\n\r\n"
|
||||
|
||||
func errorPayload(errorCode string) []byte {
|
||||
payload := make([]byte, len(errorPayloadTemplate))
|
||||
copy(payload, errorPayloadTemplate)
|
||||
|
||||
copy(payload[29:58], []byte(time.Now().Format(time.RFC1123)))
|
||||
copy(payload[9:12], errorCode)
|
||||
|
||||
return payload
|
||||
}
|
||||
@@ -1,532 +0,0 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"crypto/rand"
|
||||
"io/ioutil"
|
||||
_ "log"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/http/httputil"
|
||||
_ "reflect"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/buger/goreplay/proto"
|
||||
)
|
||||
|
||||
func TestHTTPClientURLPort(t *testing.T) {
|
||||
c1 := NewHTTPClient("http://example.com", &HTTPClientConfig{})
|
||||
if c1.baseURL != "http://example.com" {
|
||||
t.Error("Sould not add 80 port for http:", c1.baseURL)
|
||||
}
|
||||
|
||||
c2 := NewHTTPClient("https://example.com", &HTTPClientConfig{})
|
||||
if c2.baseURL != "https://example.com" {
|
||||
t.Error("Sould not add 443 port for https:", c2.baseURL)
|
||||
}
|
||||
|
||||
c3 := NewHTTPClient("https://example.com:1", &HTTPClientConfig{})
|
||||
if c3.baseURL != "https://example.com:1" {
|
||||
t.Error("Sould use specified port:", c3.baseURL)
|
||||
}
|
||||
|
||||
c4 := NewHTTPClient("example.com", &HTTPClientConfig{})
|
||||
if c4.baseURL != "http://example.com" {
|
||||
t.Error("Sould not add default protocol:", c4.baseURL)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPClientSend(t *testing.T) {
|
||||
wg := new(sync.WaitGroup)
|
||||
|
||||
payload := func(reqType string) []byte {
|
||||
switch reqType {
|
||||
case "GET":
|
||||
return []byte("GET / HTTP/1.1\r\n\r\n")
|
||||
case "POST":
|
||||
return []byte("POST /post HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2")
|
||||
case "POST_CHUNKED":
|
||||
return []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")
|
||||
}
|
||||
|
||||
return []byte("")
|
||||
}
|
||||
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
if r.Method == "POST" {
|
||||
defer r.Body.Close()
|
||||
body, _ := ioutil.ReadAll(r.Body)
|
||||
|
||||
if len(r.TransferEncoding) > 0 && r.TransferEncoding[0] == "chunked" {
|
||||
if string(body) != "Wikipedia in\r\n\r\nchunks." {
|
||||
t.Error("Wrong POST body:", body, string(body))
|
||||
}
|
||||
} else {
|
||||
if string(body) != "a=1&b=2" {
|
||||
buf, _ := httputil.DumpRequest(r, true)
|
||||
t.Error("Wrong POST body:", string(body), string(buf))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
wg.Done()
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewHTTPClient(server.URL, &HTTPClientConfig{Debug: true})
|
||||
|
||||
wg.Add(4)
|
||||
client.Send(payload("POST"))
|
||||
client.Send(payload("GET"))
|
||||
client.Send(payload("POST_CHUNKED"))
|
||||
client.Send(payload("POST"))
|
||||
|
||||
wg.Wait()
|
||||
|
||||
client = NewHTTPClient(server.URL, &HTTPClientConfig{Debug: true, CompatibilityMode: true})
|
||||
|
||||
wg.Add(4)
|
||||
|
||||
if _, err := client.Send(payload("POST")); err != nil {
|
||||
t.Fatal(err.Error())
|
||||
}
|
||||
if _, err := client.Send(payload("GET")); err != nil {
|
||||
t.Fatal(err.Error())
|
||||
}
|
||||
if _, err := client.Send(payload("POST_CHUNKED")); err != nil {
|
||||
t.Fatal(err.Error())
|
||||
}
|
||||
if _, err := client.Send(payload("POST")); err != nil {
|
||||
t.Fatal(err.Error())
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func TestHTTPClientResonseByClose(t *testing.T) {
|
||||
wg := new(sync.WaitGroup)
|
||||
|
||||
payload := []byte("GET / HTTP/1.1\r\n\r\n")
|
||||
ln, _ := net.Listen("tcp", ":0")
|
||||
go func() {
|
||||
for {
|
||||
conn, _ := ln.Accept()
|
||||
buf := make([]byte, 4096)
|
||||
conn.Read(buf)
|
||||
|
||||
conn.Write([]byte("HTTP/1.1 200 OK\r\n\r\n"))
|
||||
conn.Write([]byte("ab"))
|
||||
conn.Close()
|
||||
|
||||
wg.Done()
|
||||
}
|
||||
}()
|
||||
|
||||
client := NewHTTPClient(ln.Addr().String(), &HTTPClientConfig{Debug: true})
|
||||
|
||||
wg.Add(1)
|
||||
resp, _ := client.Send(payload)
|
||||
|
||||
if !bytes.Equal(resp, []byte("HTTP/1.1 200 OK\r\n\r\nab")) {
|
||||
t.Error("Should return valid response", string(resp))
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
// https://github.com/buger/gor/issues/184
|
||||
func TestHTTPClientResponseBuffer(t *testing.T) {
|
||||
testCases := []struct {
|
||||
name string
|
||||
responseSize int
|
||||
buffserSize int
|
||||
expectedSize int
|
||||
timeout time.Duration
|
||||
}{
|
||||
{"Chunked, buffer overflow", 10 * 1024, 1024, 1024, 50 * time.Millisecond},
|
||||
|
||||
{"Chunked, fits buffer", 10 * 1024, 64 * 1024, 10*1024 + 145 /* headers length + chunked meta */, 50 * time.Millisecond},
|
||||
|
||||
{"Content-Length, buffer overflow", 1024, 1000, 1000, 50 * time.Millisecond},
|
||||
|
||||
{"Content-Length, fits buffer", 1024, 64 * 1024, 1024 + 118, 50 * time.Millisecond},
|
||||
}
|
||||
|
||||
for _, tc := range testCases {
|
||||
wg := new(sync.WaitGroup)
|
||||
|
||||
payload := []byte("GET / HTTP/1.1\r\n\r\n")
|
||||
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
size := tc.responseSize // 1kb
|
||||
rb := make([]byte, size)
|
||||
rand.Read(rb)
|
||||
|
||||
w.Write(rb[:size/2])
|
||||
w.Write(rb[size/2:])
|
||||
|
||||
wg.Done()
|
||||
}))
|
||||
|
||||
client := NewHTTPClient(server.URL, &HTTPClientConfig{Debug: true, ResponseBufferSize: tc.buffserSize, Timeout: 100 * time.Millisecond})
|
||||
|
||||
wg.Add(2)
|
||||
|
||||
start := time.Now()
|
||||
client.Send(payload)
|
||||
resp, err := client.Send(payload)
|
||||
stop := time.Now()
|
||||
|
||||
if err != nil {
|
||||
t.Error("Request error", err)
|
||||
}
|
||||
|
||||
if stop.Sub(start) > tc.timeout {
|
||||
t.Error("Request took too long", stop.Sub(start), tc.timeout)
|
||||
}
|
||||
|
||||
if len(resp) != tc.expectedSize {
|
||||
t.Error(tc.name, " - Wrong response size:", tc.expectedSize, len(resp))
|
||||
} else {
|
||||
if !bytes.Equal(resp[0:8], []byte("HTTP/1.1")) {
|
||||
t.Error(tc.name, " - Response buffer contains data from previous request", string(resp), len(resp))
|
||||
}
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
server.Close()
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPClientHTTPSSend(t *testing.T) {
|
||||
wg := new(sync.WaitGroup)
|
||||
|
||||
payload := func(reqType string) []byte {
|
||||
switch reqType {
|
||||
case "GET":
|
||||
return []byte("GET / HTTP/1.1\r\n\r\n")
|
||||
case "POST":
|
||||
return []byte("POST /post HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.org\r\n\r\na=1&b=2")
|
||||
case "POST_CHUNKED":
|
||||
return []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")
|
||||
}
|
||||
|
||||
return []byte("")
|
||||
}
|
||||
|
||||
server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
if r.Method == "POST" {
|
||||
defer r.Body.Close()
|
||||
body, _ := ioutil.ReadAll(r.Body)
|
||||
|
||||
if len(r.TransferEncoding) > 0 && r.TransferEncoding[0] == "chunked" {
|
||||
if string(body) != "Wikipedia in\r\n\r\nchunks." {
|
||||
t.Error("Wrong POST body:", body, string(body))
|
||||
}
|
||||
} else {
|
||||
if string(body) != "a=1&b=2" {
|
||||
buf, _ := httputil.DumpRequest(r, true)
|
||||
t.Error("Wrong POST body:", string(body), string(buf))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
wg.Done()
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewHTTPClient(server.URL, &HTTPClientConfig{})
|
||||
|
||||
wg.Add(4)
|
||||
client.Send(payload("POST"))
|
||||
client.Send(payload("GET"))
|
||||
client.Send(payload("POST_CHUNKED"))
|
||||
client.Send(payload("POST"))
|
||||
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func TestHTTPClientServerInstantDisconnect(t *testing.T) {
|
||||
wg := new(sync.WaitGroup)
|
||||
|
||||
GETPayload := []byte("GET / HTTP/1.1\r\n\r\n")
|
||||
|
||||
ln, _ := net.Listen("tcp", "127.0.0.1:0")
|
||||
defer ln.Close()
|
||||
|
||||
go func() {
|
||||
for {
|
||||
conn, err := ln.Accept()
|
||||
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
conn.Close()
|
||||
wg.Done()
|
||||
}
|
||||
}()
|
||||
|
||||
client := NewHTTPClient(ln.Addr().String(), &HTTPClientConfig{})
|
||||
|
||||
wg.Add(2)
|
||||
client.Send(GETPayload)
|
||||
client.Send(GETPayload)
|
||||
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func TestHTTPClientServerNoKeepAlive(t *testing.T) {
|
||||
wg := new(sync.WaitGroup)
|
||||
|
||||
GETPayload := []byte("GET / HTTP/1.1\r\n\r\n")
|
||||
|
||||
ln, _ := net.Listen("tcp", "127.0.0.1:0")
|
||||
defer ln.Close()
|
||||
|
||||
go func() {
|
||||
for {
|
||||
conn, err := ln.Accept()
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
|
||||
buf := make([]byte, 4096)
|
||||
_, err = conn.Read(buf)
|
||||
if err != nil {
|
||||
t.Error("Error reading:", err.Error())
|
||||
}
|
||||
conn.Write([]byte("OK"))
|
||||
|
||||
// No keep-alive connections
|
||||
conn.Close()
|
||||
|
||||
wg.Done()
|
||||
}
|
||||
}()
|
||||
|
||||
client := NewHTTPClient(ln.Addr().String(), &HTTPClientConfig{})
|
||||
|
||||
wg.Add(2)
|
||||
client.Send(GETPayload)
|
||||
client.Send(GETPayload)
|
||||
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func TestHTTPClientRedirect(t *testing.T) {
|
||||
wg := new(sync.WaitGroup)
|
||||
|
||||
GETPayload := []byte("GET / HTTP/1.1\r\n\r\n")
|
||||
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
if r.URL.Path == "/" {
|
||||
http.Redirect(w, r, "/new", 301)
|
||||
}
|
||||
|
||||
wg.Done()
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewHTTPClient(server.URL, &HTTPClientConfig{FollowRedirects: 1, Debug: false})
|
||||
|
||||
// Should do 2 queries
|
||||
wg.Add(2)
|
||||
client.Send(GETPayload)
|
||||
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func TestHTTPClientRedirectLimit(t *testing.T) {
|
||||
wg := new(sync.WaitGroup)
|
||||
|
||||
GETPayload := []byte("GET / HTTP/1.1\r\n\r\n")
|
||||
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
if r.URL.Path == "/" {
|
||||
http.Redirect(w, r, "/r1", 301)
|
||||
}
|
||||
|
||||
if r.URL.Path == "/r1" {
|
||||
http.Redirect(w, r, "/r2", 301)
|
||||
}
|
||||
|
||||
if r.URL.Path == "/r2" {
|
||||
http.Redirect(w, r, "/new", 301)
|
||||
}
|
||||
|
||||
wg.Done()
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewHTTPClient(server.URL, &HTTPClientConfig{FollowRedirects: 2, Debug: false})
|
||||
|
||||
// Have 3 redirects + 1 GET, but should do only 2 redirects + GET
|
||||
wg.Add(3)
|
||||
client.Send(GETPayload)
|
||||
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func TestHTTPClientKeepHeadersRedirect(t *testing.T) {
|
||||
wg := new(sync.WaitGroup)
|
||||
|
||||
GETPayload := []byte("GET / HTTP/1.1\r\n\r\n")
|
||||
|
||||
GETPayload = proto.AddHeader(GETPayload, []byte("keep-header"), []byte("true"))
|
||||
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
if r.URL.Path == "/" {
|
||||
http.Redirect(w, r, "/new", 301)
|
||||
}
|
||||
|
||||
if r.Header.Get("keep-header") != "true" {
|
||||
t.Errorf("Header keep-header was incorrect, got: %s, want: %s.", r.Header.Get("keep-header"), "true")
|
||||
}
|
||||
|
||||
wg.Done()
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewHTTPClient(server.URL, &HTTPClientConfig{FollowRedirects: 1, Debug: false})
|
||||
|
||||
wg.Add(2)
|
||||
client.Send(GETPayload)
|
||||
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func TestHTTPClientBasicAuth(t *testing.T) {
|
||||
wg := new(sync.WaitGroup)
|
||||
wg.Add(2)
|
||||
|
||||
GETPayload := []byte("GET / HTTP/1.1\r\n\r\n")
|
||||
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
user, pass, _ := r.BasicAuth()
|
||||
|
||||
if user != "user" || pass != "pass" {
|
||||
http.Error(w, "Unauthorized.", 401)
|
||||
wg.Done()
|
||||
return
|
||||
}
|
||||
|
||||
wg.Done()
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewHTTPClient(server.URL, &HTTPClientConfig{Debug: false})
|
||||
resp, _ := client.Send(GETPayload)
|
||||
client.Disconnect()
|
||||
|
||||
if !bytes.Equal(proto.Status(resp), []byte("401")) {
|
||||
t.Error("Should return unauthorized error", string(resp))
|
||||
}
|
||||
|
||||
authUrl := strings.Replace(server.URL, "http://", "http://user:pass@", -1)
|
||||
client = NewHTTPClient(authUrl, &HTTPClientConfig{Debug: false})
|
||||
resp, _ = client.Send(GETPayload)
|
||||
client.Disconnect()
|
||||
|
||||
if !bytes.Equal(proto.Status(resp), []byte("200")) {
|
||||
t.Error("Should return proper response", string(resp))
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
func TestHTTPClientHandleHTTP10(t *testing.T) {
|
||||
wg := new(sync.WaitGroup)
|
||||
|
||||
GETPayload := []byte("GET http://foobar.com/path HTTP/1.0\r\n\r\n")
|
||||
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
|
||||
if r.URL.Path != "/path" {
|
||||
t.Error("Path not match:", r.URL.Path)
|
||||
}
|
||||
|
||||
wg.Done()
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
client := NewHTTPClient(server.URL, &HTTPClientConfig{Debug: true})
|
||||
|
||||
wg.Add(1)
|
||||
client.Send(GETPayload)
|
||||
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
// func TestHTTPClientErrors(t *testing.T) {
|
||||
// req := []byte("GET http://foobar.com/path HTTP/1.0\r\n\r\n")
|
||||
|
||||
// // Port not exists
|
||||
// client := NewHTTPClient("http://127.0.0.1:1", &HTTPClientConfig{Debug: true})
|
||||
// if resp, err := client.Send(req); err != nil {
|
||||
// if s := proto.Status(resp); !bytes.Equal(s, []byte("521")) {
|
||||
// t.Error("Should return status 521 for connection refused, instead:", string(s), err)
|
||||
// }
|
||||
// } else {
|
||||
// t.Error("Should throw error")
|
||||
// }
|
||||
|
||||
// client = NewHTTPClient("http://not.existing", &HTTPClientConfig{Debug: true})
|
||||
// if resp, err := client.Send(req); err != nil {
|
||||
// if s := proto.Status(resp); !bytes.Equal(s, []byte("521")) {
|
||||
// t.Error("Should return status 521 for no such host, instead:", string(s))
|
||||
// }
|
||||
// } else {
|
||||
// t.Error("Should throw error")
|
||||
// }
|
||||
|
||||
// // Non routable IP address to simulate connection timeout
|
||||
// client = NewHTTPClient("http://10.255.255.1", &HTTPClientConfig{Debug: true, ConnectionTimeout: 100 * time.Millisecond})
|
||||
|
||||
// if resp, err := client.Send(req); err != nil {
|
||||
// if s := proto.Status(resp); !bytes.Equal(s, []byte("521")) {
|
||||
// t.Error("Should return status 521 for io/timeout:", string(s))
|
||||
// }
|
||||
// } else {
|
||||
// t.Error("Should throw error")
|
||||
// }
|
||||
|
||||
// // Connecting but io timeout on read
|
||||
// ln, _ := net.Listen("tcp", "127.0.0.1:0")
|
||||
// client = NewHTTPClient("http://"+ln.Addr().String(), &HTTPClientConfig{Debug: true, Timeout: 10 * time.Millisecond})
|
||||
// defer ln.Close()
|
||||
|
||||
// if resp, err := client.Send(req); err != nil {
|
||||
// if s := proto.Status(resp); !bytes.Equal(s, []byte("524")) {
|
||||
// t.Error("Should return status 524 for io read, instead:", string(s))
|
||||
// }
|
||||
// } else {
|
||||
// t.Error("Should throw error")
|
||||
// }
|
||||
|
||||
// // Response read error read tcp [::1]:51128: connection reset by peer &{{0xc20802a000}}
|
||||
// ln1, _ := net.Listen("tcp", "127.0.0.1:0")
|
||||
// go func() {
|
||||
// ln1.Accept()
|
||||
// }()
|
||||
// defer ln1.Close()
|
||||
|
||||
// client = NewHTTPClient("http://"+ln1.Addr().String(), &HTTPClientConfig{Debug: true, Timeout: 10 * time.Millisecond})
|
||||
|
||||
// if resp, err := client.Send(req); err != nil {
|
||||
// if s := proto.Status(resp); !bytes.Equal(s, []byte("524")) {
|
||||
// t.Error("Should return status 524 for connection reset by peer, instead:", string(s))
|
||||
// }
|
||||
// } else {
|
||||
// t.Error("Should throw error")
|
||||
// }
|
||||
// }
|
||||
+30
-12
@@ -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)
|
||||
|
||||
|
||||
+19
-25
@@ -3,6 +3,7 @@ package main
|
||||
import (
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"fmt"
|
||||
"io/ioutil"
|
||||
"net/http/httputil"
|
||||
"strconv"
|
||||
@@ -11,36 +12,25 @@ import (
|
||||
)
|
||||
|
||||
func prettifyHTTP(p []byte) []byte {
|
||||
headSize := bytes.IndexByte(p, '\n') + 1
|
||||
head := p[:headSize]
|
||||
body := p[headSize:]
|
||||
|
||||
headersPos := proto.MIMEHeadersEndPos(body)
|
||||
tEnc := bytes.Equal(proto.Header(p, []byte("Transfer-Encoding")), []byte("chunked"))
|
||||
cEnc := bytes.Equal(proto.Header(p, []byte("Content-Encoding")), []byte("gzip"))
|
||||
|
||||
if headersPos < 5 || headersPos > len(body) {
|
||||
if !(tEnc || cEnc) {
|
||||
return p
|
||||
}
|
||||
|
||||
headers := body[:headersPos]
|
||||
content := body[headersPos:]
|
||||
headersPos := proto.MIMEHeadersEndPos(p)
|
||||
|
||||
var tEnc, cEnc []byte
|
||||
proto.ParseHeaders([][]byte{headers}, func(header, value []byte) {
|
||||
if bytes.EqualFold(header, []byte("Transfer-Encoding")) {
|
||||
tEnc = value
|
||||
}
|
||||
|
||||
if bytes.EqualFold(header, []byte("Content-Encoding")) {
|
||||
cEnc = value
|
||||
}
|
||||
})
|
||||
|
||||
if len(tEnc) == 0 && len(cEnc) == 0 {
|
||||
if headersPos < 5 || headersPos > len(p) {
|
||||
return p
|
||||
}
|
||||
|
||||
if bytes.Equal(tEnc, []byte("chunked")) {
|
||||
buf := bytes.NewBuffer(content)
|
||||
headers := p[:headersPos]
|
||||
content := p[headersPos:]
|
||||
|
||||
if tEnc {
|
||||
buf := bytes.NewReader(content)
|
||||
r := httputil.NewChunkedReader(buf)
|
||||
content, _ = ioutil.ReadAll(r)
|
||||
|
||||
@@ -50,8 +40,8 @@ func prettifyHTTP(p []byte) []byte {
|
||||
headers = proto.SetHeader(headers, []byte("Content-Length"), []byte(newLen))
|
||||
}
|
||||
|
||||
if bytes.Equal(cEnc, []byte("gzip")) {
|
||||
buf := bytes.NewBuffer(content)
|
||||
if cEnc {
|
||||
buf := bytes.NewReader(content)
|
||||
g, err := gzip.NewReader(buf)
|
||||
|
||||
if err != nil {
|
||||
@@ -59,7 +49,11 @@ func prettifyHTTP(p []byte) []byte {
|
||||
return []byte{}
|
||||
}
|
||||
|
||||
content, _ = ioutil.ReadAll(g)
|
||||
content, err = ioutil.ReadAll(g)
|
||||
if err != nil {
|
||||
Debug(1, fmt.Sprintf("[HTTP-PRETTIFIER] %q", err))
|
||||
return p
|
||||
}
|
||||
|
||||
headers = proto.DeleteHeader(headers, []byte("Content-Encoding"))
|
||||
|
||||
@@ -67,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
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
|
||||
+18
-6
@@ -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
|
||||
}
|
||||
|
||||
+18
-15
@@ -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))
|
||||
|
||||
}
|
||||
|
||||
|
||||
+50
-72
@@ -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,24 +134,18 @@ 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()
|
||||
|
||||
var expectedLatency int64 = 250000000 - 100000000
|
||||
var expectedLatency int64 = 300000000 - 100000000
|
||||
realLatency := end - start
|
||||
if realLatency < expectedLatency {
|
||||
if realLatency > expectedLatency {
|
||||
t.Errorf("Should emit requests respecting latency. Expected: %v, real: %v", expectedLatency, realLatency)
|
||||
}
|
||||
|
||||
if realLatency > expectedLatency+10000000 {
|
||||
t.Errorf("Should emit requests respecting latency. Expected: %v, real: %v", expectedLatency, realLatency)
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
func TestInputFileMultipleFilesWithRequestsAndResponses(t *testing.T) {
|
||||
@@ -186,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))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -215,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()
|
||||
@@ -231,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)
|
||||
@@ -254,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
|
||||
}
|
||||
|
||||
@@ -272,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
|
||||
@@ -286,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
|
||||
}
|
||||
}
|
||||
@@ -308,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()
|
||||
@@ -341,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)
|
||||
@@ -376,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
|
||||
}
|
||||
|
||||
+8
-16
@@ -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
|
||||
|
||||
+9
-15
@@ -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)
|
||||
|
||||
|
||||
+34
-22
@@ -15,15 +15,11 @@ type KafkaInput struct {
|
||||
config *InputKafkaConfig
|
||||
consumers []sarama.PartitionConsumer
|
||||
messages chan *sarama.ConsumerMessage
|
||||
quit chan struct{}
|
||||
}
|
||||
|
||||
// NewKafkaInput creates instance of kafka consumer client.
|
||||
func NewKafkaInput(address string, config *InputKafkaConfig) *KafkaInput {
|
||||
return NewKafkaInputWithTLS(address, config, nil)
|
||||
}
|
||||
|
||||
// NewKafkaInputWithTLS creates instance of kafka consumer client with TLS
|
||||
func NewKafkaInputWithTLS(address string, config *InputKafkaConfig, tlsConfig *KafkaTLSConfig) *KafkaInput {
|
||||
// NewKafkaInput creates instance of kafka consumer client with TLS config
|
||||
func NewKafkaInput(address string, config *InputKafkaConfig, tlsConfig *KafkaTLSConfig) *KafkaInput {
|
||||
c := NewKafkaConfig(tlsConfig)
|
||||
|
||||
var con sarama.Consumer
|
||||
@@ -32,7 +28,6 @@ func NewKafkaInputWithTLS(address string, config *InputKafkaConfig, tlsConfig *K
|
||||
con = config.consumer
|
||||
} else {
|
||||
var err error
|
||||
//con, err = sarama.NewConsumer([]string{config.Host}, c)
|
||||
con, err = sarama.NewConsumer(strings.Split(config.Host, ","), c)
|
||||
|
||||
if err != nil {
|
||||
@@ -49,6 +44,7 @@ func NewKafkaInputWithTLS(address string, config *InputKafkaConfig, tlsConfig *K
|
||||
config: config,
|
||||
consumers: make([]sarama.PartitionConsumer, len(partitions)),
|
||||
messages: make(chan *sarama.ConsumerMessage, 256),
|
||||
quit: make(chan struct{}),
|
||||
}
|
||||
|
||||
for index, partition := range partitions {
|
||||
@@ -80,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
|
||||
}
|
||||
}
|
||||
|
||||
copy(data, buf)
|
||||
// does it have meta
|
||||
if isOriginPayload(msg.Data) {
|
||||
msg.Meta, msg.Data = payloadMetaWithBody(msg.Data)
|
||||
}
|
||||
|
||||
return len(buf), 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
|
||||
}
|
||||
|
||||
+8
-10
@@ -20,17 +20,16 @@ func TestInputKafkaRAW(t *testing.T) {
|
||||
consumer: consumer,
|
||||
Topic: "test",
|
||||
UseJSON: false,
|
||||
})
|
||||
}, 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)
|
||||
|
||||
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")
|
||||
}
|
||||
}
|
||||
|
||||
+59
-50
@@ -12,6 +12,7 @@ import (
|
||||
"github.com/buger/goreplay/capture"
|
||||
"github.com/buger/goreplay/proto"
|
||||
"github.com/buger/goreplay/size"
|
||||
"github.com/buger/goreplay/stat"
|
||||
"github.com/buger/goreplay/tcp"
|
||||
)
|
||||
|
||||
@@ -68,7 +69,7 @@ type RAWInputConfig struct {
|
||||
type RAWInput struct {
|
||||
sync.Mutex
|
||||
RAWInputConfig
|
||||
messageStats []tcp.Stats
|
||||
st *stat.Writer
|
||||
listener *capture.Listener
|
||||
message chan *tcp.Message
|
||||
cancelListener context.CancelFunc
|
||||
@@ -96,45 +97,50 @@ func NewRAWInput(address string, config RAWInputConfig) (i *RAWInput) {
|
||||
}
|
||||
i.host = host
|
||||
i.port = uint16(port)
|
||||
if i.Stats {
|
||||
i.st = initializeStatWriter("input_raw_message.csv", "type: GOR_STAT\ntitle: raw messages info")
|
||||
statsWriteRecord(i.st, []string{
|
||||
"length", "lostdata", "source", "destination", "ipversion",
|
||||
"timedout", "truncated", "id",
|
||||
}, "[INPUT-RAW]")
|
||||
}
|
||||
|
||||
i.listen(address)
|
||||
|
||||
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()
|
||||
if i.st != nil {
|
||||
i.st.Close()
|
||||
i.st = nil
|
||||
}
|
||||
return nil, ErrorStopped
|
||||
case msgTCP = <-i.message:
|
||||
msg.Data = msgTCP.Data()
|
||||
}
|
||||
var header []byte
|
||||
|
||||
var msgType byte = ResponsePayload
|
||||
if msg.IsIncoming {
|
||||
timestamp := msgTCP.Start.UnixNano()
|
||||
latency := msgTCP.End.UnixNano() - timestamp
|
||||
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())
|
||||
id := msgTCP.UUID()
|
||||
msg.Meta = payloadHeader(msgType, id, timestamp, latency)
|
||||
|
||||
n = copy(data, header)
|
||||
if len(data) > len(header) {
|
||||
n += copy(data[len(header):], buf)
|
||||
if i.Stats && i.st != nil {
|
||||
statsWriteRecord(i.st, getStats(msgTCP.Stats, id), "[INPUT-RAW]")
|
||||
}
|
||||
dis := len(header) + len(buf) - n
|
||||
if dis > 0 {
|
||||
go Debug(2, "[INPUT-RAW] discarded", dis, "bytes increase copy buffer size")
|
||||
}
|
||||
if msg.Truncated {
|
||||
go Debug(2, "[INPUT-RAW] message truncated, copy-buffer-size")
|
||||
}
|
||||
go i.addStats(msg.Stats)
|
||||
return n, nil
|
||||
msgTCP = nil
|
||||
return &msg, nil
|
||||
}
|
||||
|
||||
func (i *RAWInput) listen(address string) {
|
||||
@@ -149,8 +155,10 @@ func (i *RAWInput) listen(address string) {
|
||||
log.Fatal(err)
|
||||
}
|
||||
pool := tcp.NewMessagePool(i.CopyBufferSize, i.Expire, Debug, i.handler)
|
||||
pool.End = endHint
|
||||
pool.Start = startHint
|
||||
pool.MatchUUID(i.TrackResponse)
|
||||
if i.Protocol == ProtocolHTTP {
|
||||
pool.Start = http1StartHint
|
||||
}
|
||||
var ctx context.Context
|
||||
ctx, i.cancelListener = context.WithCancel(context.Background())
|
||||
errCh := i.listener.ListenBackground(ctx, pool.Handler)
|
||||
@@ -171,13 +179,26 @@ func (i *RAWInput) String() string {
|
||||
}
|
||||
|
||||
// GetStats returns the stats so far and reset the stats
|
||||
func (i *RAWInput) GetStats() []tcp.Stats {
|
||||
i.Lock()
|
||||
defer func() {
|
||||
i.messageStats = []tcp.Stats{}
|
||||
i.Unlock()
|
||||
}()
|
||||
return i.messageStats
|
||||
func getStats(stat tcp.Stats, id []byte) []string {
|
||||
var s [8]string
|
||||
|
||||
s[0] = strconv.Itoa(stat.Length) // length
|
||||
s[1] = strconv.Itoa(stat.LostData) // lostdata
|
||||
s[2] = stat.SrcAddr // source
|
||||
s[3] = stat.DstAddr // destination
|
||||
s[4] = string(stat.IPversion + '0') // ipversion
|
||||
if stat.TimedOut { // timedout
|
||||
s[5] = "true"
|
||||
} else {
|
||||
s[5] = "false"
|
||||
}
|
||||
if stat.Truncated { // truncated
|
||||
s[6] = "true"
|
||||
} else {
|
||||
s[6] = "false"
|
||||
}
|
||||
s[7] = string(id)
|
||||
return s[:]
|
||||
}
|
||||
|
||||
// Close closes the input raw listener
|
||||
@@ -187,22 +208,10 @@ func (i *RAWInput) Close() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (i *RAWInput) addStats(mStats tcp.Stats) {
|
||||
if i.Stats {
|
||||
i.Lock()
|
||||
if len(i.messageStats) >= 10000 {
|
||||
i.messageStats = []tcp.Stats{}
|
||||
}
|
||||
i.messageStats = append(i.messageStats, mStats)
|
||||
|
||||
i.Unlock()
|
||||
func http1StartHint(pckt *tcp.Packet) (isIncoming, isOutgoing bool) {
|
||||
isIncoming = proto.HasRequestTitle(pckt.Payload)
|
||||
if isIncoming {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
func startHint(pckt *tcp.Packet) (isIncoming, isOutgoing bool) {
|
||||
return proto.HasRequestTitle(pckt.Payload), proto.HasResponseTitle(pckt.Payload)
|
||||
}
|
||||
|
||||
func endHint(m *tcp.Message) bool {
|
||||
return proto.HasFullPayload(m.Data())
|
||||
return false, proto.HasResponseTitle(pckt.Payload)
|
||||
}
|
||||
|
||||
+54
-74
@@ -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,35 +47,33 @@ 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)
|
||||
|
||||
client := NewHTTPClient("127.0.0.1:"+port, &HTTPClientConfig{})
|
||||
|
||||
emitter := NewEmitter(quit)
|
||||
addr := "http://127.0.0.1:" + port
|
||||
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 = client.Get("/")
|
||||
_, err = http.Get(addr)
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
return
|
||||
@@ -92,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 {
|
||||
@@ -118,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)
|
||||
@@ -128,20 +123,20 @@ 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)
|
||||
|
||||
client := NewHTTPClient("127.0.0.1:"+port, &HTTPClientConfig{})
|
||||
addr := "http://127.0.0.1:" + port
|
||||
|
||||
emitter := NewEmitter(quit)
|
||||
emitter := NewEmitter()
|
||||
go emitter.Start(plugins, Settings.Middleware)
|
||||
|
||||
for i := 0; i < 10; i++ {
|
||||
// request + response
|
||||
wg.Add(2)
|
||||
_, err = client.Get("/")
|
||||
_, err = http.Get(addr)
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
return
|
||||
@@ -158,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 {
|
||||
@@ -184,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)
|
||||
@@ -194,18 +188,17 @@ func TestRAWInputIPv6(t *testing.T) {
|
||||
})
|
||||
|
||||
plugins := &InOutPlugins{
|
||||
Inputs: []io.Reader{input},
|
||||
Outputs: []io.Writer{output},
|
||||
Inputs: []PluginReader{input},
|
||||
Outputs: []PluginWriter{output},
|
||||
}
|
||||
|
||||
client := NewHTTPClient(originAddr, &HTTPClientConfig{})
|
||||
|
||||
emitter := NewEmitter(quit)
|
||||
emitter := NewEmitter()
|
||||
addr := "http://" + originAddr
|
||||
go emitter.Start(plugins, Settings.Middleware)
|
||||
for i := 0; i < 10; i++ {
|
||||
// request + response
|
||||
wg.Add(2)
|
||||
_, err = client.Get("/")
|
||||
_, err = http.Get(addr)
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
return
|
||||
@@ -222,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")
|
||||
|
||||
@@ -256,15 +248,15 @@ func TestInputRAWChunkedEncoding(t *testing.T) {
|
||||
}))
|
||||
defer replay.Close()
|
||||
|
||||
httpOutput := NewHTTPOutput(replay.URL, &HTTPOutputConfig{Debug: true})
|
||||
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)
|
||||
@@ -280,17 +272,15 @@ func TestInputRAWChunkedEncoding(t *testing.T) {
|
||||
}
|
||||
|
||||
func BenchmarkRAWInputWithReplay(b *testing.B) {
|
||||
var respCounter, reqCounter, replayCounter, capturedBody 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
|
||||
@@ -307,15 +297,9 @@ func BenchmarkRAWInputWithReplay(b *testing.B) {
|
||||
|
||||
replay := http.Server{
|
||||
Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
defer wg.Done()
|
||||
defer r.Body.Close()
|
||||
w.Write([]byte("ab"))
|
||||
atomic.AddUint64(&replayCounter, 1)
|
||||
data, err := ioutil.ReadAll(r.Body)
|
||||
if err != nil && err != io.EOF && err != io.ErrUnexpectedEOF {
|
||||
b.Log(err)
|
||||
}
|
||||
atomic.AddUint64(&capturedBody, uint64(len(data)))
|
||||
atomic.AddUint32(&replayCounter, 1)
|
||||
w.Write(nil)
|
||||
wg.Done()
|
||||
}),
|
||||
}
|
||||
go replay.Serve(listener0)
|
||||
@@ -330,41 +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++
|
||||
}
|
||||
atomic.AddUint64(&capturedBody, uint64(len(data)))
|
||||
wg.Done()
|
||||
})
|
||||
httpOutput := NewHTTPOutput(replayAddr, &HTTPOutputConfig{Debug: false})
|
||||
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()
|
||||
var buf [1 << 20]byte
|
||||
buf[1<<20-1] = 'a'
|
||||
client := NewHTTPClient(originAddr, &HTTPClientConfig{ResponseBufferSize: 2 << 20, CompatibilityMode: true})
|
||||
addr := "http://" + originAddr
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
if i&1 == 0 {
|
||||
_, err = client.Get("/")
|
||||
} else {
|
||||
_, err = client.Post("/", buf[:])
|
||||
}
|
||||
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, %d Bytes in %s\n", reqCounter, b.N, respCounter, b.N, replayCounter, b.N, capturedBody, time.Since(now))
|
||||
b.ReportMetric(float64(reqCounter), "requests")
|
||||
b.ReportMetric(float64(respCounter), "responses")
|
||||
b.ReportMetric(float64(replayCounter), "replayed")
|
||||
emitter.Close()
|
||||
}
|
||||
|
||||
+43
-29
@@ -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
|
||||
}
|
||||
|
||||
+15
-15
@@ -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{
|
||||
|
||||
@@ -4,6 +4,7 @@ import (
|
||||
"bytes"
|
||||
"crypto/tls"
|
||||
"crypto/x509"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io/ioutil"
|
||||
"log"
|
||||
@@ -12,9 +13,8 @@ import (
|
||||
"github.com/buger/goreplay/proto"
|
||||
)
|
||||
|
||||
// KafkaConfig should contains required information to
|
||||
// InputKafkaConfig should contains required information to
|
||||
// build producers.
|
||||
|
||||
type InputKafkaConfig struct {
|
||||
producer sarama.AsyncProducer
|
||||
consumer sarama.Consumer
|
||||
@@ -23,6 +23,7 @@ type InputKafkaConfig struct {
|
||||
UseJSON bool `json:"input-kafka-json-format"`
|
||||
}
|
||||
|
||||
// OutputKafkaConfig is the representation of kfka output configuration
|
||||
type OutputKafkaConfig struct {
|
||||
producer sarama.AsyncProducer
|
||||
consumer sarama.Consumer
|
||||
@@ -34,8 +35,8 @@ type OutputKafkaConfig struct {
|
||||
// KafkaTLSConfig should contains TLS certificates for connecting to secured Kafka clusters
|
||||
type KafkaTLSConfig struct {
|
||||
CACert string `json:"kafka-tls-ca-cert"`
|
||||
clientCert string `json:"kafka-tls-client-cert"`
|
||||
clientKey string `json:"kafka-tls-client-key"`
|
||||
ClientCert string `json:"kafka-tls-client-cert"`
|
||||
ClientKey string `json:"kafka-tls-client-key"`
|
||||
}
|
||||
|
||||
// KafkaMessage should contains catched request information that should be
|
||||
@@ -50,36 +51,45 @@ type KafkaMessage struct {
|
||||
ReqHeaders map[string]string `json:"Req_Headers,omitempty"`
|
||||
}
|
||||
|
||||
|
||||
// NewTLSConfig loads TLS certificates
|
||||
func NewTLSConfig(clientCertFile, clientKeyFile, caCertFile string) (*tls.Config, error) {
|
||||
tlsConfig := tls.Config{}
|
||||
|
||||
if clientCertFile != "" && clientKeyFile == "" {
|
||||
return &tlsConfig, errors.New("Missing key of client certificate in kafka")
|
||||
}
|
||||
if clientCertFile == "" && clientKeyFile != "" {
|
||||
return &tlsConfig, errors.New("missing TLS client certificate in kafka")
|
||||
}
|
||||
// Load client cert
|
||||
cert, err := tls.LoadX509KeyPair(clientCertFile, clientKeyFile)
|
||||
if err != nil {
|
||||
return &tlsConfig, err
|
||||
if (clientCertFile != "") && (clientKeyFile != "") {
|
||||
cert, err := tls.LoadX509KeyPair(clientCertFile, clientKeyFile)
|
||||
if err != nil {
|
||||
return &tlsConfig, err
|
||||
}
|
||||
tlsConfig.Certificates = []tls.Certificate{cert}
|
||||
}
|
||||
tlsConfig.Certificates = []tls.Certificate{cert}
|
||||
|
||||
// Load CA cert
|
||||
caCert, err := ioutil.ReadFile(caCertFile)
|
||||
if err != nil {
|
||||
return &tlsConfig, err
|
||||
if caCertFile != "" {
|
||||
caCert, err := ioutil.ReadFile(caCertFile)
|
||||
if err != nil {
|
||||
return &tlsConfig, err
|
||||
}
|
||||
caCertPool := x509.NewCertPool()
|
||||
caCertPool.AppendCertsFromPEM(caCert)
|
||||
tlsConfig.RootCAs = caCertPool
|
||||
}
|
||||
caCertPool := x509.NewCertPool()
|
||||
caCertPool.AppendCertsFromPEM(caCert)
|
||||
tlsConfig.RootCAs = caCertPool
|
||||
|
||||
return &tlsConfig, err
|
||||
return &tlsConfig, nil
|
||||
}
|
||||
|
||||
// NewKafkaConfig returns Kafka config with or without TLS
|
||||
func NewKafkaConfig(tlsConfig *KafkaTLSConfig) *sarama.Config {
|
||||
config := sarama.NewConfig()
|
||||
// Configuration options go here
|
||||
if (tlsConfig != nil) && (tlsConfig.CACert != "") && (tlsConfig.clientCert != "") && (tlsConfig.clientKey != "") {
|
||||
if tlsConfig != nil && (tlsConfig.ClientCert != "" || tlsConfig.CACert != "") {
|
||||
config.Net.TLS.Enable = true
|
||||
tlsConfig, err := NewTLSConfig(tlsConfig.clientCert, tlsConfig.clientKey, tlsConfig.CACert)
|
||||
tlsConfig, err := NewTLSConfig(tlsConfig.ClientCert, tlsConfig.ClientKey, tlsConfig.CACert)
|
||||
if err != nil {
|
||||
log.Fatal(err)
|
||||
}
|
||||
|
||||
+18
-13
@@ -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
|
||||
|
||||
+16
-21
@@ -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++ {
|
||||
|
||||
+37
-52
@@ -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
|
||||
}
|
||||
|
||||
+179
-178
@@ -1,225 +1,226 @@
|
||||
package main
|
||||
|
||||
// import (
|
||||
// "bytes"
|
||||
// "crypto/rand"
|
||||
// "encoding/hex"
|
||||
// "io"
|
||||
// "net/http"
|
||||
// "net/http/httptest"
|
||||
// "strings"
|
||||
// "sync"
|
||||
// "testing"
|
||||
// "time"
|
||||
import (
|
||||
"bytes"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/http/httputil"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
// "github.com/buger/goreplay/capture"
|
||||
// "github.com/buger/goreplay/proto"
|
||||
// )
|
||||
"github.com/buger/goreplay/capture"
|
||||
"github.com/buger/goreplay/proto"
|
||||
)
|
||||
|
||||
// type fakeServiceCb func(string, int, []byte)
|
||||
type fakeServiceCb func(string, int, []byte)
|
||||
|
||||
// // Simple service that generate token on request, and require this token for accesing to secure area
|
||||
// func NewFakeSecureService(wg *sync.WaitGroup, cb fakeServiceCb) *httptest.Server {
|
||||
// active_tokens := make([]string, 0)
|
||||
// var mu sync.Mutex
|
||||
// Simple service that generate token on request, and require this token for accesing to secure area
|
||||
func NewFakeSecureService(wg *sync.WaitGroup, cb fakeServiceCb) *httptest.Server {
|
||||
activeTokens := make([]string, 0)
|
||||
var mu sync.Mutex
|
||||
|
||||
// server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
|
||||
// mu.Lock()
|
||||
// defer mu.Unlock()
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
|
||||
mu.Lock()
|
||||
defer mu.Unlock()
|
||||
|
||||
// switch req.URL.Path {
|
||||
// case "/token":
|
||||
// // Generate random token
|
||||
// token_length := 10
|
||||
// buf := make([]byte, token_length)
|
||||
// rand.Read(buf)
|
||||
// token := hex.EncodeToString(buf)
|
||||
// active_tokens = append(active_tokens, token)
|
||||
switch req.URL.Path {
|
||||
case "/token":
|
||||
// Generate random token
|
||||
tokenLength := 10
|
||||
buf := make([]byte, tokenLength)
|
||||
rand.Read(buf)
|
||||
token := hex.EncodeToString(buf)
|
||||
activeTokens = append(activeTokens, token)
|
||||
|
||||
// w.Write([]byte(token))
|
||||
w.Write([]byte(token))
|
||||
|
||||
// cb(req.URL.Path, 200, []byte(token))
|
||||
// case "/secure":
|
||||
// token := req.URL.Query().Get("token")
|
||||
// token_found := false
|
||||
cb(req.URL.Path, 200, []byte(token))
|
||||
case "/secure":
|
||||
token := req.URL.Query().Get("token")
|
||||
tokenFound := false
|
||||
|
||||
// for _, t := range active_tokens {
|
||||
// if t == token {
|
||||
// token_found = true
|
||||
// break
|
||||
// }
|
||||
// }
|
||||
for _, t := range activeTokens {
|
||||
if t == token {
|
||||
tokenFound = true
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
// if token_found {
|
||||
// w.WriteHeader(http.StatusAccepted)
|
||||
// cb(req.URL.Path, 202, []byte(nil))
|
||||
// } else {
|
||||
// w.WriteHeader(http.StatusForbidden)
|
||||
// cb(req.URL.Path, 403, []byte(nil))
|
||||
// }
|
||||
// }
|
||||
if tokenFound {
|
||||
w.WriteHeader(http.StatusAccepted)
|
||||
cb(req.URL.Path, 202, nil)
|
||||
} else {
|
||||
w.WriteHeader(http.StatusForbidden)
|
||||
cb(req.URL.Path, 403, nil)
|
||||
}
|
||||
}
|
||||
|
||||
// wg.Done()
|
||||
// }))
|
||||
wg.Done()
|
||||
}))
|
||||
|
||||
// return server
|
||||
// }
|
||||
return server
|
||||
}
|
||||
|
||||
// func TestFakeSecureService(t *testing.T) {
|
||||
// var resp, token []byte
|
||||
func TestFakeSecureService(t *testing.T) {
|
||||
var resp, token []byte
|
||||
|
||||
// wg := new(sync.WaitGroup)
|
||||
wg := new(sync.WaitGroup)
|
||||
|
||||
// server := NewFakeSecureService(wg, func(path string, status int, resp []byte) {
|
||||
// })
|
||||
// defer server.Close()
|
||||
server := NewFakeSecureService(wg, func(path string, status int, resp []byte) {
|
||||
})
|
||||
defer server.Close()
|
||||
|
||||
// wg.Add(3)
|
||||
wg.Add(3)
|
||||
|
||||
// client := NewHTTPClient(server.URL, &HTTPClientConfig{Debug: true})
|
||||
// resp, _ = client.Get("/token")
|
||||
// token = proto.Body(resp)
|
||||
client := NewHTTPClient(&HTTPOutputConfig{}).Client
|
||||
rep, _ := client.Get(server.URL + "/token")
|
||||
resp, _ = httputil.DumpResponse(rep, true)
|
||||
token = proto.Body(resp)
|
||||
|
||||
// // Right token
|
||||
// resp, _ = client.Get("/secure?token=" + string(token))
|
||||
// if !bytes.Equal(proto.Status(resp), []byte("202")) {
|
||||
// t.Error("Valid token should return status 202:", string(proto.Status(resp)))
|
||||
// }
|
||||
// Right token
|
||||
rep, _ = client.Get(server.URL + "/secure?token=" + string(token))
|
||||
resp, _ = httputil.DumpResponse(rep, true)
|
||||
if !bytes.Equal(proto.Status(resp), []byte("202")) {
|
||||
t.Error("Valid token should return status 202:", string(proto.Status(resp)))
|
||||
}
|
||||
|
||||
// // Wrong tokens forbidden
|
||||
// resp, _ = client.Get("/secure?token=wrong")
|
||||
// if !bytes.Equal(proto.Status(resp), []byte("403")) {
|
||||
// t.Error("Wrong token should returns status 403:", string(proto.Status(resp)))
|
||||
// }
|
||||
// Wrong tokens forbidden
|
||||
rep, _ = client.Get(server.URL + "/secure?token=wrong")
|
||||
resp, _ = httputil.DumpResponse(rep, true)
|
||||
if !bytes.Equal(proto.Status(resp), []byte("403")) {
|
||||
t.Error("Wrong token should returns status 403:", string(proto.Status(resp)))
|
||||
}
|
||||
|
||||
// wg.Wait()
|
||||
// }
|
||||
wg.Wait()
|
||||
}
|
||||
|
||||
// func TestEchoMiddleware(t *testing.T) {
|
||||
// wg := new(sync.WaitGroup)
|
||||
func TestEchoMiddleware(t *testing.T) {
|
||||
wg := new(sync.WaitGroup)
|
||||
|
||||
// from := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
// w.Header().Set("Env", "prod")
|
||||
// w.Header().Set("RequestPath", r.URL.Path)
|
||||
// wg.Done()
|
||||
// }))
|
||||
// defer from.Close()
|
||||
from := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Env", "prod")
|
||||
w.Header().Set("RequestPath", r.URL.Path)
|
||||
wg.Done()
|
||||
}))
|
||||
defer from.Close()
|
||||
|
||||
// to := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
// w.Header().Set("Env", "test")
|
||||
// w.Header().Set("RequestPath", r.URL.Path)
|
||||
// wg.Done()
|
||||
// }))
|
||||
// defer to.Close()
|
||||
to := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Env", "test")
|
||||
w.Header().Set("RequestPath", r.URL.Path)
|
||||
wg.Done()
|
||||
}))
|
||||
defer to.Close()
|
||||
|
||||
// quit := make(chan int)
|
||||
// Catch traffic from one service
|
||||
fromAddr := strings.Replace(from.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
|
||||
conf := RAWInputConfig{
|
||||
Engine: capture.EnginePcap,
|
||||
Expire: testRawExpire,
|
||||
Protocol: ProtocolHTTP,
|
||||
TrackResponse: true,
|
||||
}
|
||||
input := NewRAWInput(fromAddr, conf)
|
||||
|
||||
// // Catch traffic from one service
|
||||
// fromAddr := strings.Replace(from.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
|
||||
// conf := RAWInputConfig{
|
||||
// engine: capture.EnginePcap,
|
||||
// expire: testRawExpire,
|
||||
// protocol: ProtocolHTTP,
|
||||
// trackResponse: true,
|
||||
// }
|
||||
// input := NewRAWInput(fromAddr, conf)
|
||||
// And redirect to another
|
||||
output := NewHTTPOutput(to.URL, &HTTPOutputConfig{})
|
||||
|
||||
// // And redirect to another
|
||||
// output := NewHTTPOutput(to.URL, &HTTPOutputConfig{Debug: false})
|
||||
plugins := &InOutPlugins{
|
||||
Inputs: []PluginReader{input},
|
||||
Outputs: []PluginWriter{output},
|
||||
}
|
||||
plugins.All = append(plugins.All, input, output)
|
||||
|
||||
// plugins := &InOutPlugins{
|
||||
// Inputs: []io.Reader{input},
|
||||
// Outputs: []io.Writer{output},
|
||||
// }
|
||||
// plugins.All = append(plugins.All, input, output)
|
||||
// Start Gor
|
||||
emitter := NewEmitter()
|
||||
emitter.Start(plugins, "echo -n && GOR_TEST=true && ./examples/middleware/echo.sh")
|
||||
|
||||
// // Start Gor
|
||||
// emitter := NewEmitter(quit)
|
||||
// go emitter.Start(plugins, "echo -n && GOR_TEST=true && ./examples/middleware/echo.sh")
|
||||
// Wait till middleware initialization
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
|
||||
// // Wait till middleware initialization
|
||||
// time.Sleep(100 * time.Millisecond)
|
||||
// Should receive 2 requests from original + 2 from replayed
|
||||
client := NewHTTPClient(output.(*HTTPOutput).config).Client
|
||||
|
||||
// // Should receive 2 requests from original + 2 from replayed
|
||||
// client := NewHTTPClient(from.URL, &HTTPClientConfig{Debug: false})
|
||||
for i := 0; i < 10; i++ {
|
||||
wg.Add(2)
|
||||
// Request should be echoed
|
||||
client.Get(to.URL + "/a")
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
client.Get(to.URL + "/b")
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
}
|
||||
|
||||
// for i := 0; i < 10; i++ {
|
||||
// wg.Add(2)
|
||||
// // Request should be echoed
|
||||
// client.Get("/a")
|
||||
// time.Sleep(5 * time.Millisecond)
|
||||
// client.Get("/b")
|
||||
// time.Sleep(5 * time.Millisecond)
|
||||
// }
|
||||
wg.Wait()
|
||||
emitter.Close()
|
||||
}
|
||||
|
||||
// wg.Wait()
|
||||
// emitter.Close()
|
||||
// }
|
||||
func TestTokenMiddleware(t *testing.T) {
|
||||
var resp, token []byte
|
||||
|
||||
// func TestTokenMiddleware(t *testing.T) {
|
||||
// var resp, token []byte
|
||||
wg := new(sync.WaitGroup)
|
||||
|
||||
// wg := new(sync.WaitGroup)
|
||||
from := NewFakeSecureService(wg, func(path string, status int, tok []byte) {
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
})
|
||||
defer from.Close()
|
||||
|
||||
// from := NewFakeSecureService(wg, func(path string, status int, tok []byte) {
|
||||
// time.Sleep(10 * time.Millisecond)
|
||||
// })
|
||||
// defer from.Close()
|
||||
to := NewFakeSecureService(wg, func(path string, status int, tok []byte) {
|
||||
switch path {
|
||||
case "/secure":
|
||||
if status != 202 {
|
||||
t.Error("Server should receive valid rewritten token")
|
||||
}
|
||||
}
|
||||
|
||||
// to := NewFakeSecureService(wg, func(path string, status int, tok []byte) {
|
||||
// switch path {
|
||||
// case "/secure":
|
||||
// if status != 202 {
|
||||
// t.Error("Server should receive valid rewritten token")
|
||||
// }
|
||||
// }
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
})
|
||||
defer to.Close()
|
||||
|
||||
// time.Sleep(10 * time.Millisecond)
|
||||
// })
|
||||
// defer to.Close()
|
||||
Settings.Middleware = "echo -n && GOR_TEST=true && go run ./examples/middleware/token_modifier.go"
|
||||
|
||||
// quit := make(chan int)
|
||||
fromAddr := strings.Replace(from.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
|
||||
conf := RAWInputConfig{
|
||||
Engine: capture.EnginePcap,
|
||||
Expire: testRawExpire,
|
||||
Protocol: ProtocolHTTP,
|
||||
TrackResponse: true,
|
||||
}
|
||||
// Catch traffic from one service
|
||||
input := NewRAWInput(fromAddr, conf)
|
||||
|
||||
// Settings.middleware = "echo -n && GOR_TEST=true && go run ./examples/middleware/token_modifier.go"
|
||||
// And redirect to another
|
||||
output := NewHTTPOutput(to.URL, &HTTPOutputConfig{})
|
||||
|
||||
// fromAddr := strings.Replace(from.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
|
||||
// conf := RAWInputConfig{
|
||||
// engine: capture.EnginePcap,
|
||||
// expire: testRawExpire,
|
||||
// protocol: ProtocolHTTP,
|
||||
// trackResponse: true,
|
||||
// }
|
||||
// // Catch traffic from one service
|
||||
// input := NewRAWInput(fromAddr, conf)
|
||||
plugins := &InOutPlugins{
|
||||
Inputs: []PluginReader{input},
|
||||
Outputs: []PluginWriter{output},
|
||||
}
|
||||
plugins.All = append(plugins.All, input, output)
|
||||
|
||||
// // And redirect to another
|
||||
// output := NewHTTPOutput(to.URL, &HTTPOutputConfig{Debug: true})
|
||||
// Start Gor
|
||||
emitter := NewEmitter()
|
||||
emitter.Start(plugins, Settings.Middleware)
|
||||
|
||||
// plugins := &InOutPlugins{
|
||||
// Inputs: []io.Reader{input},
|
||||
// Outputs: []io.Writer{output},
|
||||
// }
|
||||
// plugins.All = append(plugins.All, input, output)
|
||||
// Should receive 2 requests from original + 2 from replayed
|
||||
wg.Add(2)
|
||||
|
||||
// // Start Gor
|
||||
// emitter := NewEmitter(quit)
|
||||
// go emitter.Start(plugins, Settings.middleware)
|
||||
client := NewHTTPClient(&HTTPOutputConfig{}).Client
|
||||
|
||||
// // Should receive 2 requests from original + 2 from replayed
|
||||
// wg.Add(2)
|
||||
// Sending traffic to original service
|
||||
rep, _ := client.Get(to.URL + "/token")
|
||||
resp, _ = httputil.DumpResponse(rep, true)
|
||||
token = proto.Body(resp)
|
||||
|
||||
// client := NewHTTPClient(from.URL, &HTTPClientConfig{Debug: true})
|
||||
rep, _ = client.Get(to.URL + "/secure?token=" + string(token))
|
||||
resp, _ = httputil.DumpResponse(rep, true)
|
||||
if !bytes.Equal(proto.Status(resp), []byte("202")) {
|
||||
t.Error("Valid token should return 202:", proto.Status(resp))
|
||||
}
|
||||
|
||||
// // Sending traffic to original service
|
||||
// resp, _ = client.Get("/token")
|
||||
// token = proto.Body(resp)
|
||||
|
||||
// resp, _ = client.Get("/secure?token=" + string(token))
|
||||
// if !bytes.Equal(proto.Status(resp), []byte("202")) {
|
||||
// t.Error("Valid token should return 202:", proto.Status(resp))
|
||||
// }
|
||||
|
||||
// wg.Wait()
|
||||
// emitter.Close()
|
||||
// Settings.middleware = ""
|
||||
// }
|
||||
wg.Wait()
|
||||
emitter.Close()
|
||||
Settings.Middleware = ""
|
||||
}
|
||||
|
||||
+39
-44
@@ -1,7 +1,6 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"io"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
@@ -25,30 +24,26 @@ 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
|
||||
}
|
||||
|
||||
// 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 +84,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 +108,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 +124,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 +165,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
|
||||
}
|
||||
|
||||
+11
-5
@@ -1,7 +1,7 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
)
|
||||
|
||||
// DummyOutput used for debugging, prints all incoming requests
|
||||
@@ -15,10 +15,16 @@ func NewDummyOutput() (di *DummyOutput) {
|
||||
return
|
||||
}
|
||||
|
||||
func (i *DummyOutput) Write(data []byte) (int, error) {
|
||||
fmt.Println(string(data))
|
||||
|
||||
return len(data), nil
|
||||
// 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
|
||||
}
|
||||
|
||||
func (i *DummyOutput) String() string {
|
||||
|
||||
+12
-9
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+26
-29
@@ -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) {
|
||||
|
||||
+213
-236
@@ -1,56 +1,26 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"crypto/tls"
|
||||
"fmt"
|
||||
"io"
|
||||
"log"
|
||||
"math"
|
||||
"net/http"
|
||||
"net/http/httputil"
|
||||
"net/url"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/buger/goreplay/proto"
|
||||
"github.com/buger/goreplay/size"
|
||||
)
|
||||
|
||||
const initialDynamicWorkers = 10
|
||||
|
||||
type httpWorker struct {
|
||||
output *HTTPOutput
|
||||
client *HTTPClient
|
||||
lastActivity time.Time
|
||||
queue chan []byte
|
||||
stop chan bool
|
||||
}
|
||||
|
||||
func newHTTPWorker(output *HTTPOutput, queue chan []byte) *httpWorker {
|
||||
client := NewHTTPClient(output.address, &HTTPClientConfig{
|
||||
FollowRedirects: output.config.RedirectLimit,
|
||||
Debug: output.config.Debug,
|
||||
OriginalHost: output.config.OriginalHost,
|
||||
Timeout: output.config.Timeout,
|
||||
ResponseBufferSize: int(output.config.BufferSize),
|
||||
})
|
||||
|
||||
w := &httpWorker{client: client}
|
||||
if queue == nil {
|
||||
w.queue = make(chan []byte, 100)
|
||||
} else {
|
||||
w.queue = queue
|
||||
}
|
||||
w.stop = make(chan bool)
|
||||
|
||||
go func() {
|
||||
for {
|
||||
select {
|
||||
case payload := <-w.queue:
|
||||
output.sendRequest(client, payload)
|
||||
case <-w.stop:
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
return w
|
||||
}
|
||||
const (
|
||||
initialDynamicWorkers = 10
|
||||
readChunkSize = 64 * 1024
|
||||
maxResponseSize = 1073741824
|
||||
)
|
||||
|
||||
type response struct {
|
||||
payload []byte
|
||||
@@ -61,272 +31,279 @@ type response struct {
|
||||
|
||||
// HTTPOutputConfig struct for holding http output configuration
|
||||
type HTTPOutputConfig struct {
|
||||
RedirectLimit int `json:"output-http-redirect-limit"`
|
||||
|
||||
Stats bool `json:"output-http-stats"`
|
||||
WorkersMin int `json:"output-http-workers-min"`
|
||||
WorkersMax int `json:"output-http-workers"`
|
||||
StatsMs int `json:"output-http-stats-ms"`
|
||||
Workers int
|
||||
QueueLen int `json:"output-http-queue-len"`
|
||||
|
||||
ElasticSearch string `json:"output-http-elasticsearch"`
|
||||
|
||||
Timeout time.Duration `json:"output-http-timeout"`
|
||||
OriginalHost bool `json:"output-http-original-host"`
|
||||
BufferSize size.Size `json:"output-http-response-buffer"`
|
||||
|
||||
CompatibilityMode bool `json:"output-http-compatibility-mode"`
|
||||
|
||||
RequestGroup string
|
||||
|
||||
Debug bool `json:"output-http-debug"`
|
||||
|
||||
TrackResponses bool `json:"output-http-track-response"`
|
||||
TrackResponses bool `json:"output-http-track-response"`
|
||||
OriginalHost bool `json:"output-http-original-host"`
|
||||
RedirectLimit int `json:"output-http-redirect-limit"`
|
||||
WorkersMin int `json:"output-http-workers-min"`
|
||||
WorkersMax int `json:"output-http-workers"`
|
||||
QueueLen int `json:"output-http-queue-len"`
|
||||
ElasticSearch string `json:"output-http-elasticsearch"`
|
||||
Timeout time.Duration `json:"output-http-timeout"`
|
||||
WorkerTimeout time.Duration `json:"output-http-worker-timeout"`
|
||||
BufferSize size.Size `json:"output-http-response-buffer"`
|
||||
SkipVerify bool `json:"output-http-skip-verify"`
|
||||
rawURL string
|
||||
url *url.URL
|
||||
}
|
||||
|
||||
// HTTPOutput plugin manage pool of workers which send request to replayed server
|
||||
// By default workers pool is dynamic and starts with 10 workers
|
||||
// You can specify fixed number of workers using `--output-http-workers`
|
||||
// By default workers pool is dynamic and starts with 1 worker or workerMin workers
|
||||
// You can specify maximum number of workers using `--output-http-workers`
|
||||
type HTTPOutput struct {
|
||||
// Keep this as first element of struct because it guarantees 64bit
|
||||
// alignment. atomic.* functions crash on 32bit machines if operand is not
|
||||
// aligned at 64bit. See https://github.com/golang/go/issues/599
|
||||
activeWorkers int64
|
||||
|
||||
workerSessions map[string]*httpWorker
|
||||
|
||||
address string
|
||||
limit int
|
||||
queue chan []byte
|
||||
|
||||
responses chan response
|
||||
|
||||
needWorker chan int
|
||||
|
||||
config *HTTPOutputConfig
|
||||
|
||||
queueStats *GorStat
|
||||
|
||||
activeWorkers int32
|
||||
config *HTTPOutputConfig
|
||||
elasticSearch *ESPlugin
|
||||
|
||||
stop chan bool // Channel used only to indicate goroutine should shutdown
|
||||
client *HTTPClient
|
||||
stopWorker chan struct{}
|
||||
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)
|
||||
|
||||
o.address = address
|
||||
var err error
|
||||
config.url, err = url.Parse(address)
|
||||
if err != nil {
|
||||
log.Fatal(fmt.Sprintf("[OUTPUT-HTTP] parse HTTP output URL error[%q]", err))
|
||||
}
|
||||
if config.url.Scheme == "" {
|
||||
config.url.Scheme = "http"
|
||||
}
|
||||
config.rawURL = config.url.String()
|
||||
if config.Timeout < time.Millisecond*100 {
|
||||
config.Timeout = time.Second
|
||||
}
|
||||
if config.BufferSize <= 0 {
|
||||
config.BufferSize = 100 * 1024 // 100kb
|
||||
}
|
||||
if config.WorkersMin <= 0 {
|
||||
config.WorkersMin = 1
|
||||
}
|
||||
if config.WorkersMin > 1000 {
|
||||
config.WorkersMin = 1000
|
||||
}
|
||||
if config.WorkersMax <= 0 {
|
||||
config.WorkersMax = math.MaxInt32 // idealy so large
|
||||
}
|
||||
if config.WorkersMax < config.WorkersMin {
|
||||
config.WorkersMax = config.WorkersMin
|
||||
}
|
||||
if config.QueueLen <= 0 {
|
||||
config.QueueLen = 1000
|
||||
}
|
||||
if config.RedirectLimit < 0 {
|
||||
config.RedirectLimit = 0
|
||||
}
|
||||
if config.WorkerTimeout <= 0 {
|
||||
config.WorkerTimeout = time.Second * 2
|
||||
}
|
||||
o.config = config
|
||||
o.stop = make(chan bool)
|
||||
|
||||
if o.config.Stats {
|
||||
o.queueStats = NewGorStat("output_http", o.config.StatsMs)
|
||||
}
|
||||
|
||||
o.queue = make(chan []byte, o.config.QueueLen)
|
||||
o.responses = make(chan response, o.config.QueueLen)
|
||||
o.needWorker = make(chan int, 1)
|
||||
|
||||
// Initial workers count
|
||||
if o.config.WorkersMax == 0 {
|
||||
o.needWorker <- initialDynamicWorkers
|
||||
} else {
|
||||
o.needWorker <- o.config.WorkersMax
|
||||
o.queue = make(chan *Message, o.config.QueueLen)
|
||||
if o.config.TrackResponses {
|
||||
o.responses = make(chan *response, o.config.QueueLen)
|
||||
}
|
||||
// it should not be buffered to avoid races
|
||||
o.stopWorker = make(chan struct{})
|
||||
|
||||
if o.config.ElasticSearch != "" {
|
||||
o.elasticSearch = new(ESPlugin)
|
||||
o.elasticSearch.Init(o.config.ElasticSearch)
|
||||
}
|
||||
|
||||
if Settings.RecognizeTCPSessions {
|
||||
if !PRO {
|
||||
log.Fatal("Detailed TCP sessions work only with PRO license")
|
||||
}
|
||||
o.workerSessions = make(map[string]*httpWorker, 100)
|
||||
go o.sessionWorkerMaster()
|
||||
} else {
|
||||
go o.workerMaster()
|
||||
o.client = NewHTTPClient(o.config)
|
||||
o.activeWorkers += int32(o.config.WorkersMin)
|
||||
for i := 0; i < o.config.WorkersMin; i++ {
|
||||
go o.startWorker()
|
||||
}
|
||||
|
||||
go o.workerMaster()
|
||||
return o
|
||||
}
|
||||
|
||||
func (o *HTTPOutput) workerMaster() {
|
||||
for {
|
||||
newWorkers := <-o.needWorker
|
||||
atomic.AddInt64(&o.activeWorkers, int64(newWorkers))
|
||||
for i := 0; i < newWorkers; i++ {
|
||||
go o.startWorker()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (o *HTTPOutput) sessionWorkerMaster() {
|
||||
gc := time.Tick(time.Second)
|
||||
|
||||
for {
|
||||
select {
|
||||
case p := <-o.queue:
|
||||
id := payloadID(p)
|
||||
sessionID := string(id[0:20])
|
||||
worker, ok := o.workerSessions[sessionID]
|
||||
|
||||
if !ok {
|
||||
atomic.AddInt64(&o.activeWorkers, 1)
|
||||
worker = newHTTPWorker(o, nil)
|
||||
o.workerSessions[sessionID] = worker
|
||||
}
|
||||
|
||||
worker.queue <- p
|
||||
worker.lastActivity = time.Now()
|
||||
case <-gc:
|
||||
now := time.Now()
|
||||
|
||||
for id, w := range o.workerSessions {
|
||||
if !w.lastActivity.IsZero() && now.Sub(w.lastActivity) >= 120*time.Second {
|
||||
w.stop <- true
|
||||
delete(o.workerSessions, id)
|
||||
atomic.AddInt64(&o.activeWorkers, -1)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (o *HTTPOutput) startWorker() {
|
||||
client := NewHTTPClient(o.address, &HTTPClientConfig{
|
||||
FollowRedirects: o.config.RedirectLimit,
|
||||
Debug: o.config.Debug,
|
||||
OriginalHost: o.config.OriginalHost,
|
||||
Timeout: o.config.Timeout,
|
||||
ResponseBufferSize: int(o.config.BufferSize),
|
||||
CompatibilityMode: o.config.CompatibilityMode,
|
||||
})
|
||||
|
||||
var timer = time.NewTimer(o.config.WorkerTimeout)
|
||||
defer func() {
|
||||
// recover from panics caused by trying to send in
|
||||
// a closed chan(o.stopWorker)
|
||||
recover()
|
||||
}()
|
||||
defer timer.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-o.stop:
|
||||
return
|
||||
case data := <-o.queue:
|
||||
o.sendRequest(client, data)
|
||||
case <-time.After(2 * time.Second):
|
||||
// When dynamic scaling enabled workers die after 2s of inactivity
|
||||
if o.config.WorkersMin == o.config.WorkersMax {
|
||||
continue
|
||||
}
|
||||
default:
|
||||
<-timer.C
|
||||
}
|
||||
// rollback workers
|
||||
rollback:
|
||||
if atomic.LoadInt32(&o.activeWorkers) > int32(o.config.WorkersMin) && len(o.queue) < 1 {
|
||||
// close one worker
|
||||
o.stopWorker <- struct{}{}
|
||||
atomic.AddInt32(&o.activeWorkers, -1)
|
||||
goto rollback
|
||||
}
|
||||
timer.Reset(o.config.WorkerTimeout)
|
||||
}
|
||||
}
|
||||
|
||||
workersCount := int(atomic.LoadInt64(&o.activeWorkers))
|
||||
|
||||
// At least 1 startWorker should be alive
|
||||
if workersCount != 1 && workersCount > o.config.WorkersMin {
|
||||
atomic.AddInt64(&o.activeWorkers, -1)
|
||||
return
|
||||
}
|
||||
func (o *HTTPOutput) startWorker() {
|
||||
for {
|
||||
select {
|
||||
case <-o.stopWorker:
|
||||
return
|
||||
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 {
|
||||
o.queueStats.Write(len(o.queue))
|
||||
}
|
||||
|
||||
if !Settings.RecognizeTCPSessions && o.config.WorkersMax != o.config.WorkersMin {
|
||||
workersCount := int(atomic.LoadInt64(&o.activeWorkers))
|
||||
|
||||
if len(o.queue) > workersCount {
|
||||
extraWorkersReq := len(o.queue) - workersCount + 1
|
||||
maxWorkersAvailable := o.config.WorkersMax - workersCount
|
||||
if extraWorkersReq > maxWorkersAvailable {
|
||||
extraWorkersReq = maxWorkersAvailable
|
||||
}
|
||||
if extraWorkersReq > 0 {
|
||||
o.needWorker <- extraWorkersReq
|
||||
}
|
||||
if len(o.queue) > 0 {
|
||||
// try to start a new worker to serve
|
||||
if atomic.LoadInt32(&o.activeWorkers) < int32(o.config.WorkersMax) {
|
||||
go o.startWorker()
|
||||
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) {
|
||||
var resp response
|
||||
// PluginRead reads message from this plugin
|
||||
func (o *HTTPOutput) PluginRead() (*Message, error) {
|
||||
if !o.config.TrackResponses {
|
||||
return nil, ErrorStopped
|
||||
}
|
||||
var resp *response
|
||||
var msg Message
|
||||
select {
|
||||
case <-o.stop:
|
||||
return 0, ErrorStopped
|
||||
return nil, ErrorStopped
|
||||
case resp = <-o.responses:
|
||||
msg.Data = resp.payload
|
||||
}
|
||||
|
||||
Debug(3, "[OUTPUT-HTTP] Received response:", string(resp.payload))
|
||||
msg.Meta = payloadHeader(ReplayedResponsePayload, resp.uuid, resp.roundTripTime, resp.startedAt)
|
||||
|
||||
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, "[OUTPUT-HTTP] discarded", dis, "increase copy buffer size")
|
||||
}
|
||||
|
||||
return n, nil
|
||||
return &msg, nil
|
||||
}
|
||||
|
||||
func (o *HTTPOutput) sendRequest(client *HTTPClient, request []byte) {
|
||||
meta := payloadMeta(request)
|
||||
|
||||
Debug(2, fmt.Sprintf("[OUTPUT-HTTP] meta: %q", meta))
|
||||
|
||||
if len(meta) < 2 {
|
||||
func (o *HTTPOutput) sendRequest(client *HTTPClient, msg *Message) {
|
||||
if !isRequestPayload(msg.Meta) {
|
||||
return
|
||||
}
|
||||
uuid := meta[1]
|
||||
|
||||
body := payloadBody(request)
|
||||
if !proto.HasRequestTitle(body) {
|
||||
return
|
||||
}
|
||||
|
||||
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, "Error when sending ", err)
|
||||
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()}
|
||||
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)
|
||||
}
|
||||
}
|
||||
|
||||
func (o *HTTPOutput) String() string {
|
||||
return "HTTP output: " + o.address
|
||||
return "HTTP output: " + o.config.rawURL
|
||||
}
|
||||
|
||||
// Close closes the data channel so that data
|
||||
func (o *HTTPOutput) Close() error {
|
||||
close(o.stop)
|
||||
close(o.stopWorker)
|
||||
return nil
|
||||
}
|
||||
|
||||
// HTTPClient holds configurations for a single HTTP client
|
||||
type HTTPClient struct {
|
||||
config *HTTPOutputConfig
|
||||
Client *http.Client
|
||||
}
|
||||
|
||||
// NewHTTPClient returns new http client with check redirects policy
|
||||
func NewHTTPClient(config *HTTPOutputConfig) *HTTPClient {
|
||||
client := new(HTTPClient)
|
||||
client.config = config
|
||||
var transport *http.Transport
|
||||
client.Client = &http.Client{
|
||||
Timeout: client.config.Timeout,
|
||||
CheckRedirect: func(req *http.Request, via []*http.Request) error {
|
||||
if len(via) >= client.config.RedirectLimit {
|
||||
Debug(1, fmt.Sprintf("[HTTPCLIENT] maximum output-http-redirects[%d] reached!", client.config.RedirectLimit))
|
||||
return http.ErrUseLastResponse
|
||||
}
|
||||
lastReq := via[len(via)-1]
|
||||
resp := req.Response
|
||||
Debug(2, fmt.Sprintf("[HTTPCLIENT] HTTP redirects from %q to %q with %q", lastReq.Host, req.Host, resp.Status))
|
||||
return nil
|
||||
},
|
||||
}
|
||||
if config.SkipVerify {
|
||||
// clone to avoid modying global default RoundTripper
|
||||
transport = http.DefaultTransport.(*http.Transport).Clone()
|
||||
transport.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
|
||||
client.Client.Transport = transport
|
||||
}
|
||||
|
||||
return client
|
||||
}
|
||||
|
||||
// Send sends an http request using client create by NewHTTPClient
|
||||
func (c *HTTPClient) Send(data []byte) ([]byte, error) {
|
||||
var req *http.Request
|
||||
var resp *http.Response
|
||||
var err error
|
||||
|
||||
req, err = http.ReadRequest(bufio.NewReader(bytes.NewReader(data)))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
// we don't send CONNECT or OPTIONS request
|
||||
if req.Method == http.MethodConnect {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
if !c.config.OriginalHost {
|
||||
req.Host = c.config.url.Host
|
||||
}
|
||||
|
||||
req.URL = c.config.url
|
||||
// force connection to not be closed, which can affect the global client
|
||||
req.Close = false
|
||||
// it's an error if this is not equal to empty string
|
||||
req.RequestURI = ""
|
||||
|
||||
resp, err = c.Client.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if c.config.TrackResponses {
|
||||
return httputil.DumpResponse(resp, true)
|
||||
}
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
+61
-37
@@ -1,19 +1,16 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"io"
|
||||
"io/ioutil"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
_ "net/http/httputil"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestHTTPOutput(t *testing.T) {
|
||||
wg := new(sync.WaitGroup)
|
||||
quit := make(chan int)
|
||||
|
||||
input := NewTestInput()
|
||||
|
||||
@@ -39,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}
|
||||
|
||||
http_output := NewHTTPOutput(server.URL, &HTTPOutputConfig{Debug: true, 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{http_output, output},
|
||||
Inputs: []PluginReader{input},
|
||||
Outputs: []PluginWriter{httpOutput, output},
|
||||
}
|
||||
plugins.All = append(plugins.All, input, output, http_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()
|
||||
@@ -73,7 +70,6 @@ func TestHTTPOutput(t *testing.T) {
|
||||
|
||||
func TestHTTPOutputKeepOriginalHost(t *testing.T) {
|
||||
wg := new(sync.WaitGroup)
|
||||
quit := make(chan int)
|
||||
|
||||
input := NewTestInput()
|
||||
|
||||
@@ -86,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{Debug: false, OriginalHost: true})
|
||||
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)
|
||||
@@ -110,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) {
|
||||
@@ -118,15 +113,15 @@ func TestHTTPOutputSSL(t *testing.T) {
|
||||
}))
|
||||
|
||||
input := NewTestInput()
|
||||
output := NewHTTPOutput(server.URL, &HTTPOutputConfig{})
|
||||
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)
|
||||
@@ -140,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
|
||||
@@ -151,25 +145,27 @@ func TestHTTPOutputSessions(t *testing.T) {
|
||||
defer server.Close()
|
||||
|
||||
Settings.RecognizeTCPSessions = true
|
||||
output := NewHTTPOutput(server.URL, &HTTPOutputConfig{Debug: 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"))
|
||||
@@ -180,28 +176,56 @@ 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.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
wg.Done()
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
input := NewTestInput()
|
||||
output := NewHTTPOutput(server.URL, &HTTPOutputConfig{})
|
||||
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++ {
|
||||
wg.Add(1)
|
||||
input.EmitPOST()
|
||||
}
|
||||
|
||||
wg.Wait()
|
||||
emitter.Close()
|
||||
}
|
||||
|
||||
func BenchmarkHTTPOutputTLS(b *testing.B) {
|
||||
wg := new(sync.WaitGroup)
|
||||
|
||||
server := httptest.NewTLSServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
wg.Done()
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
input := NewTestInput()
|
||||
output := NewHTTPOutput(server.URL, &HTTPOutputConfig{SkipVerify: true, WorkersMax: 1})
|
||||
|
||||
plugins := &InOutPlugins{
|
||||
Inputs: []PluginReader{input},
|
||||
Outputs: []PluginWriter{output},
|
||||
}
|
||||
plugins.All = append(plugins.All, input, output)
|
||||
|
||||
emitter := NewEmitter()
|
||||
go emitter.Start(plugins, Settings.Middleware)
|
||||
|
||||
for i := 0; i < b.N; i++ {
|
||||
|
||||
+21
-24
@@ -2,11 +2,11 @@ package main
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"io"
|
||||
"log"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/buger/goreplay/byteutils"
|
||||
"github.com/buger/goreplay/proto"
|
||||
|
||||
"github.com/Shopify/sarama"
|
||||
@@ -22,13 +22,8 @@ type KafkaOutput struct {
|
||||
// KafkaOutputFrequency in milliseconds
|
||||
const KafkaOutputFrequency = 500
|
||||
|
||||
// NewKafkaOutput creates instance of kafka producer client.
|
||||
func NewKafkaOutput(address string, config *OutputKafkaConfig) io.Writer {
|
||||
return NewKafkaOutputWithTLS(address, config, nil)
|
||||
}
|
||||
|
||||
// NewKafkaOutputWithTLS creates instance of kafka producer client.
|
||||
func NewKafkaOutputWithTLS(address string, config *OutputKafkaConfig, tlsConfig *KafkaTLSConfig) io.Writer {
|
||||
// NewKafkaOutput creates instance of kafka producer client with TLS config
|
||||
func NewKafkaOutput(address string, config *OutputKafkaConfig, tlsConfig *KafkaTLSConfig) PluginWriter {
|
||||
c := NewKafkaConfig(tlsConfig)
|
||||
|
||||
var producer sarama.AsyncProducer
|
||||
@@ -67,31 +62,33 @@ 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(data)
|
||||
message = sarama.StringEncoder(byteutils.SliceToString(msg.Meta) + byteutils.SliceToString(msg.Data))
|
||||
} else {
|
||||
headers := make(map[string]string)
|
||||
proto.ParseHeaders([][]byte{data}, func(header []byte, value []byte) {
|
||||
headers[string(header)] = string(value)
|
||||
})
|
||||
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: string(proto.Path(req)),
|
||||
ReqType: string(meta[0]),
|
||||
ReqID: string(meta[1]),
|
||||
ReqTs: string(meta[2]),
|
||||
ReqMethod: string(proto.Method(req)),
|
||||
ReqBody: string(proto.Body(req)),
|
||||
ReqHeaders: headers,
|
||||
ReqURL: byteutils.SliceToString(proto.Path(req)),
|
||||
ReqType: byteutils.SliceToString(meta[0]),
|
||||
ReqID: byteutils.SliceToString(meta[1]),
|
||||
ReqTs: byteutils.SliceToString(meta[2]),
|
||||
ReqMethod: byteutils.SliceToString(proto.Method(req)),
|
||||
ReqBody: byteutils.SliceToString(proto.Body(req)),
|
||||
ReqHeaders: header,
|
||||
}
|
||||
jsonMessage, _ := json.Marshal(&kafkaMessage)
|
||||
message = sarama.StringEncoder(jsonMessage)
|
||||
message = sarama.StringEncoder(byteutils.SliceToString(jsonMessage))
|
||||
}
|
||||
|
||||
o.producer.Input() <- &sarama.ProducerMessage{
|
||||
|
||||
@@ -17,16 +17,16 @@ func TestOutputKafkaRAW(t *testing.T) {
|
||||
producer: producer,
|
||||
Topic: "test",
|
||||
UseJSON: false,
|
||||
})
|
||||
}, 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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -40,9 +40,9 @@ func TestOutputKafkaJSON(t *testing.T) {
|
||||
producer: producer,
|
||||
Topic: "test",
|
||||
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()
|
||||
|
||||
|
||||
+3
-2
@@ -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 {
|
||||
|
||||
+11
-9
@@ -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 {
|
||||
@@ -107,22 +108,23 @@ func (o *S3Output) onBufferUpdate(path string) {
|
||||
idx := getFileIndex(path)
|
||||
bucket, key := o.keyPath(idx)
|
||||
|
||||
file, _ := os.Open(path)
|
||||
// reader := bufio.NewReader(file)
|
||||
file, err := os.Open(path)
|
||||
if err != nil {
|
||||
Debug(0, fmt.Sprintf("[S3 Output] Failed to open file %q. err: %q", path, err))
|
||||
return
|
||||
}
|
||||
defer os.Remove(path)
|
||||
|
||||
_, err := svc.PutObject(&s3.PutObjectInput{
|
||||
_, err = svc.PutObject(&s3.PutObjectInput{
|
||||
Body: file,
|
||||
Bucket: aws.String(bucket),
|
||||
Key: aws.String(key),
|
||||
})
|
||||
if err != nil {
|
||||
log.Printf("[S3 Output] Failed to upload data to %s/%s, %s\n", bucket, key, err)
|
||||
os.Remove(path)
|
||||
Debug(0, fmt.Sprintf("[S3 Output] Failed to upload data to %q/%q, %q", bucket, key, err))
|
||||
return
|
||||
}
|
||||
|
||||
os.Remove(path)
|
||||
|
||||
if o.closeCh != nil {
|
||||
o.closeCh <- struct{}{}
|
||||
}
|
||||
|
||||
+51
-53
@@ -4,8 +4,6 @@ import (
|
||||
"crypto/tls"
|
||||
"fmt"
|
||||
"hash/fnv"
|
||||
"io"
|
||||
"log"
|
||||
"net"
|
||||
"time"
|
||||
)
|
||||
@@ -14,45 +12,36 @@ import (
|
||||
// Currently used for internal communication between listener and replay server
|
||||
// Can be used for transfering binary payloads like protocol buffers
|
||||
type TCPOutput struct {
|
||||
address string
|
||||
limit int
|
||||
buf []chan []byte
|
||||
bufStats *GorStat
|
||||
config *TCPOutputConfig
|
||||
address string
|
||||
limit int
|
||||
buf []chan *Message
|
||||
config *TCPOutputConfig
|
||||
workerIndex uint32
|
||||
stop chan struct{}
|
||||
}
|
||||
|
||||
// TCPOutputConfig tcp output configuration
|
||||
type TCPOutputConfig struct {
|
||||
Secure bool `json:"output-tcp-secure"`
|
||||
Sticky bool `json:"output-tcp-sticky"`
|
||||
Secure bool `json:"output-tcp-secure"`
|
||||
Sticky bool `json:"output-tcp-sticky"`
|
||||
SkipVerify bool `json:"output-tcp-skip-verify"`
|
||||
Workers int `json:"output-tcp-workers"`
|
||||
}
|
||||
|
||||
// NewTCPOutput constructor for TCPOutput
|
||||
// Initialize 10 workers which hold keep-alive connection
|
||||
func NewTCPOutput(address string, config *TCPOutputConfig) io.Writer {
|
||||
// Initialize X workers which hold keep-alive connection
|
||||
func NewTCPOutput(address string, config *TCPOutputConfig) PluginWriter {
|
||||
o := new(TCPOutput)
|
||||
|
||||
o.address = address
|
||||
o.config = config
|
||||
o.stop = make(chan struct{})
|
||||
|
||||
if Settings.OutputTCPStats {
|
||||
o.bufStats = NewGorStat("output_tcp", 5000)
|
||||
}
|
||||
|
||||
if o.config.Sticky {
|
||||
// create 10 buffers and send the buffer index to the worker
|
||||
o.buf = make([]chan []byte, 10)
|
||||
for i := 0; i < 10; i++ {
|
||||
o.buf[i] = make(chan []byte, 100)
|
||||
go o.worker(i)
|
||||
}
|
||||
} else {
|
||||
// create 1 buffer and send its index (0) to all workers
|
||||
o.buf = make([]chan []byte, 1)
|
||||
o.buf[0] = make(chan []byte, 1000)
|
||||
for i := 0; i < 10; i++ {
|
||||
go o.worker(0)
|
||||
}
|
||||
// create X buffers and send the buffer index to the worker
|
||||
o.buf = make([]chan *Message, o.config.Workers)
|
||||
for i := 0; i < o.config.Workers; i++ {
|
||||
o.buf[i] = make(chan *Message, 100)
|
||||
go o.worker(i)
|
||||
}
|
||||
|
||||
return o
|
||||
@@ -66,7 +55,7 @@ func (o *TCPOutput) worker(bufferIndex int) {
|
||||
break
|
||||
}
|
||||
|
||||
log.Println("Can't connect to aggregator instance, reconnecting in 1 second. Retries:", retries)
|
||||
Debug(1, fmt.Sprintf("Can't connect to aggregator instance, reconnecting in 1 second. Retries:%d", retries))
|
||||
time.Sleep(1 * time.Second)
|
||||
|
||||
conn, err = o.connect(o.address)
|
||||
@@ -74,19 +63,27 @@ func (o *TCPOutput) worker(bufferIndex int) {
|
||||
}
|
||||
|
||||
if retries > 0 {
|
||||
log.Println("Connected to aggregator instance after ", retries, " retries")
|
||||
Debug(2, fmt.Sprintf("Connected to aggregator instance after %d retries", retries))
|
||||
}
|
||||
|
||||
defer conn.Close()
|
||||
|
||||
for {
|
||||
data := <-o.buf[bufferIndex]
|
||||
conn.Write(data)
|
||||
_, err := conn.Write([]byte(payloadSeparator))
|
||||
var msg *Message
|
||||
select {
|
||||
case <-o.stop:
|
||||
return
|
||||
case 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 {
|
||||
log.Println("INFO: TCP output connection closed, reconnecting")
|
||||
o.buf[bufferIndex] <- data
|
||||
Debug(2, "INFO: TCP output connection closed, reconnecting")
|
||||
o.buf[bufferIndex] <- msg
|
||||
go o.worker(bufferIndex)
|
||||
break
|
||||
}
|
||||
@@ -95,36 +92,31 @@ func (o *TCPOutput) worker(bufferIndex int) {
|
||||
|
||||
func (o *TCPOutput) getBufferIndex(data []byte) int {
|
||||
if !o.config.Sticky {
|
||||
return 0
|
||||
o.workerIndex++
|
||||
return int(o.workerIndex) % o.config.Workers
|
||||
}
|
||||
|
||||
hasher := fnv.New32a()
|
||||
hasher.Write(payloadMeta(data)[1])
|
||||
return int(hasher.Sum32()) % 10
|
||||
return int(hasher.Sum32()) % o.config.Workers
|
||||
|
||||
}
|
||||
|
||||
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(msg.Data)
|
||||
o.buf[bufferIndex] <- msg
|
||||
|
||||
bufferIndex := o.getBufferIndex(data)
|
||||
o.buf[bufferIndex] <- newBuf
|
||||
|
||||
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) {
|
||||
if o.config.Secure {
|
||||
conn, err = tls.Dial("tcp", address, &tls.Config{})
|
||||
conn, err = tls.Dial("tcp", address, &tls.Config{InsecureSkipVerify: o.config.SkipVerify})
|
||||
} else {
|
||||
conn, err = net.Dial("tcp", address)
|
||||
}
|
||||
@@ -135,3 +127,9 @@ func (o *TCPOutput) connect(address string) (conn net.Conn, err error) {
|
||||
func (o *TCPOutput) String() string {
|
||||
return fmt.Sprintf("TCP output %s, limit: %d", o.address, o.limit)
|
||||
}
|
||||
|
||||
// Close closes the data channel so that data
|
||||
func (o *TCPOutput) Close() error {
|
||||
close(o.stop)
|
||||
return nil
|
||||
}
|
||||
|
||||
+27
-31
@@ -2,7 +2,6 @@ package main
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"io"
|
||||
"log"
|
||||
"net"
|
||||
"sync"
|
||||
@@ -12,24 +11,22 @@ import (
|
||||
|
||||
func TestTCPOutput(t *testing.T) {
|
||||
wg := new(sync.WaitGroup)
|
||||
quit := make(chan int)
|
||||
|
||||
listener := startTCP(func(data []byte) {
|
||||
wg.Done()
|
||||
})
|
||||
input := NewTestInput()
|
||||
output := NewTCPOutput(listener.Addr().String(), &TCPOutputConfig{})
|
||||
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 < 100; i++ {
|
||||
for i := 0; i < 10; i++ {
|
||||
wg.Add(1)
|
||||
input.EmitGET()
|
||||
}
|
||||
@@ -48,9 +45,9 @@ func startTCP(cb func([]byte)) net.Listener {
|
||||
go func() {
|
||||
for {
|
||||
conn, _ := listener.Accept()
|
||||
defer conn.Close()
|
||||
|
||||
go func() {
|
||||
go func(conn net.Conn) {
|
||||
defer conn.Close()
|
||||
reader := bufio.NewReader(conn)
|
||||
scanner := bufio.NewScanner(reader)
|
||||
scanner.Split(payloadScanner)
|
||||
@@ -58,7 +55,7 @@ func startTCP(cb func([]byte)) net.Listener {
|
||||
for scanner.Scan() {
|
||||
cb(scanner.Bytes())
|
||||
}
|
||||
}()
|
||||
}(conn)
|
||||
}
|
||||
}()
|
||||
|
||||
@@ -67,51 +64,50 @@ 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()
|
||||
})
|
||||
input := NewTestInput()
|
||||
output := NewTCPOutput(listener.Addr().String(), &TCPOutputConfig{})
|
||||
|
||||
plugins := &InOutPlugins{
|
||||
Inputs: []io.Reader{input},
|
||||
Outputs: []io.Writer{output},
|
||||
}
|
||||
plugins.All = append(plugins.All, input, output)
|
||||
|
||||
emitter := NewEmitter(quit)
|
||||
go emitter.Start(plugins, Settings.Middleware)
|
||||
|
||||
b.ResetTimer()
|
||||
input.data = make(chan []byte, b.N)
|
||||
for i := 0; i < b.N; i++ {
|
||||
wg.Add(1)
|
||||
input.EmitGET()
|
||||
}
|
||||
wg.Add(b.N)
|
||||
output := NewTCPOutput(listener.Addr().String(), &TCPOutputConfig{Workers: 10})
|
||||
|
||||
plugins := &InOutPlugins{
|
||||
Inputs: []PluginReader{input},
|
||||
Outputs: []PluginWriter{output},
|
||||
}
|
||||
|
||||
emitter := NewEmitter()
|
||||
// avoid counting above initialization
|
||||
b.ResetTimer()
|
||||
go emitter.Start(plugins, Settings.Middleware)
|
||||
|
||||
wg.Wait()
|
||||
emitter.Close()
|
||||
}
|
||||
|
||||
func TestStickyDisable(t *testing.T) {
|
||||
tcpOutput := TCPOutput{config: &TCPOutputConfig{Sticky: false}}
|
||||
tcpOutput := TCPOutput{config: &TCPOutputConfig{Sticky: false, Workers: 10}}
|
||||
|
||||
for i := 0; i < 1000; i++ {
|
||||
for i := 0; i < 10; i++ {
|
||||
index := tcpOutput.getBufferIndex(getTestBytes())
|
||||
if index != 0 {
|
||||
t.Errorf("Sticky is disable. Got: %d want 0", index)
|
||||
if index != (i+1)%10 {
|
||||
t.Errorf("Sticky is disable. Got: %d want %d", index, (i+1)%10)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestBufferDistribution(t *testing.T) {
|
||||
numberOfWorkers := 10
|
||||
numberOfMessages := 1000000
|
||||
numberOfMessages := 10000
|
||||
percentDistributionErrorRange := 20
|
||||
|
||||
buffer := make([]int, numberOfWorkers)
|
||||
tcpOutput := TCPOutput{config: &TCPOutputConfig{Sticky: true}}
|
||||
tcpOutput := TCPOutput{config: &TCPOutputConfig{Sticky: true, Workers: 10}}
|
||||
for i := 0; i < numberOfMessages; i++ {
|
||||
buffer[tcpOutput.getBufferIndex(getTestBytes())]++
|
||||
}
|
||||
|
||||
+30
-15
@@ -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
|
||||
|
||||
+2
-2
@@ -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 {
|
||||
|
||||
+1
-3
@@ -4,9 +4,7 @@ package proto
|
||||
|
||||
func Fuzz(data []byte) int {
|
||||
|
||||
ParseHeaders([][]byte{data}, func(header []byte, value []byte) bool {
|
||||
return true
|
||||
})
|
||||
ParseHeaders(data)
|
||||
|
||||
return 1
|
||||
}
|
||||
|
||||
+19
-23
@@ -21,7 +21,6 @@ import (
|
||||
"bytes"
|
||||
"net/http"
|
||||
"net/textproto"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/buger/goreplay/byteutils"
|
||||
@@ -112,14 +111,13 @@ func header(payload []byte, name []byte) (value []byte, headerStart, headerEnd,
|
||||
return
|
||||
}
|
||||
|
||||
// ParseHeaders Parsing headers from multiple payloads
|
||||
func ParseHeaders(payloads [][]byte, cb func(header []byte, value []byte)) {
|
||||
p := bytes.Join(payloads, nil)
|
||||
// ParseHeaders Parsing headers from the payload
|
||||
func ParseHeaders(p []byte) textproto.MIMEHeader {
|
||||
// trimming off the title of the request
|
||||
if HasRequestTitle(p) || HasResponseTitle(p) {
|
||||
if HasTitle(p) {
|
||||
headerStart := MIMEHeadersStartPos(p)
|
||||
if headerStart > len(p)-1 {
|
||||
return
|
||||
return nil
|
||||
}
|
||||
p = p[headerStart:]
|
||||
}
|
||||
@@ -127,17 +125,12 @@ func ParseHeaders(payloads [][]byte, cb func(header []byte, value []byte)) {
|
||||
if headerEnd > 1 {
|
||||
p = p[:headerEnd]
|
||||
}
|
||||
reader := textproto.NewReader(bufio.NewReader(bytes.NewBuffer(p)))
|
||||
reader := textproto.NewReader(bufio.NewReader(bytes.NewReader(p)))
|
||||
mime, err := reader.ReadMIMEHeader()
|
||||
if err != nil {
|
||||
return
|
||||
return nil
|
||||
}
|
||||
for k, v := range mime {
|
||||
for _, value := range v {
|
||||
cb([]byte(k), []byte(value))
|
||||
}
|
||||
}
|
||||
return
|
||||
return mime
|
||||
}
|
||||
|
||||
// Header returns header value, if header not found, value will be blank
|
||||
@@ -163,14 +156,16 @@ func SetHeader(payload, name, value []byte) []byte {
|
||||
// AddHeader takes http payload and appends new header to the start of headers section
|
||||
// Returns modified request payload
|
||||
func AddHeader(payload, name, value []byte) []byte {
|
||||
mimeStart := MIMEHeadersStartPos(payload)
|
||||
if mimeStart < 1 {
|
||||
return payload
|
||||
}
|
||||
header := make([]byte, len(name)+2+len(value)+2)
|
||||
copy(header[0:], name)
|
||||
copy(header[len(name):], HeaderDelim)
|
||||
copy(header[len(name)+2:], value)
|
||||
copy(header[len(header)-2:], CRLF)
|
||||
|
||||
mimeStart := MIMEHeadersStartPos(payload)
|
||||
|
||||
return byteutils.Insert(payload, mimeStart, header)
|
||||
}
|
||||
|
||||
@@ -330,8 +325,7 @@ const (
|
||||
|
||||
// HasResponseTitle reports whether this payload has an HTTP/1 response title
|
||||
func HasResponseTitle(payload []byte) bool {
|
||||
var s string
|
||||
byteutils.SliceToString(&payload, &s)
|
||||
s := byteutils.SliceToString(payload)
|
||||
if len(s) < MinResponseCount {
|
||||
return false
|
||||
}
|
||||
@@ -343,8 +337,11 @@ func HasResponseTitle(payload []byte) bool {
|
||||
if !(ok && major == 1 && (minor == 0 || minor == 1)) {
|
||||
return false
|
||||
}
|
||||
status, err := strconv.Atoi(s[VersionLen+1 : VersionLen+4])
|
||||
if err != nil {
|
||||
if s[VersionLen] != ' ' {
|
||||
return false
|
||||
}
|
||||
status, ok := atoI(payload[VersionLen+1:VersionLen+4], 10)
|
||||
if !ok || s[VersionLen+4] != ' ' {
|
||||
return false
|
||||
}
|
||||
statusText := http.StatusText(status)
|
||||
@@ -359,8 +356,7 @@ func HasResponseTitle(payload []byte) bool {
|
||||
|
||||
// HasRequestTitle reports whether this payload has an HTTP/1 request title
|
||||
func HasRequestTitle(payload []byte) bool {
|
||||
var s string
|
||||
byteutils.SliceToString(&payload, &s)
|
||||
s := byteutils.SliceToString(payload)
|
||||
if len(s) < MinRequestCount {
|
||||
return false
|
||||
}
|
||||
@@ -484,7 +480,7 @@ func atoI(s []byte, base int) (num int, ok bool) {
|
||||
return 0, false
|
||||
}
|
||||
v = int(hexTable[s[i]])
|
||||
if v >= base {
|
||||
if v >= base || (v == 0 && s[i] != '0') {
|
||||
return 0, false
|
||||
}
|
||||
num = (num * base) + v
|
||||
|
||||
+32
-47
@@ -2,10 +2,9 @@ package proto
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"net/textproto"
|
||||
"reflect"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestHeader(t *testing.T) {
|
||||
@@ -98,6 +97,10 @@ func TestSetHeader(t *testing.T) {
|
||||
if payload = SetHeader(payload, []byte("User-Agent"), []byte("Gor")); !bytes.Equal(payload, payloadAfter) {
|
||||
t.Error("Should add header if not found", string(payload))
|
||||
}
|
||||
invalidPayload := []byte("POST /post HTTP/1.1")
|
||||
if invalidPayload = SetHeader(invalidPayload, []byte("User-Agent"), []byte("Gor")); !bytes.Equal(invalidPayload, []byte("POST /post HTTP/1.1")) {
|
||||
t.Error("Should not modify payload if request is invalid", string(payload))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeleteHeader(t *testing.T) {
|
||||
@@ -122,16 +125,12 @@ func TestDeleteHeader(t *testing.T) {
|
||||
func TestParseHeaders(t *testing.T) {
|
||||
payload := [][]byte{[]byte("POST /post HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.or"), []byte("g\r\nUser-Ag"), []byte("ent:Chrome\r\n\r\n"), []byte("Fake-Header: asda")}
|
||||
|
||||
headers := make(map[string]string)
|
||||
headers := ParseHeaders(bytes.Join(payload, nil))
|
||||
|
||||
ParseHeaders(payload, func(header []byte, value []byte) {
|
||||
headers[string(header)] = string(value)
|
||||
})
|
||||
|
||||
expected := map[string]string{
|
||||
"Content-Length": "7",
|
||||
"Host": "www.w3.org",
|
||||
"User-Agent": "Chrome",
|
||||
expected := textproto.MIMEHeader{
|
||||
"Content-Length": []string{"7"},
|
||||
"Host": []string{"www.w3.org"},
|
||||
"User-Agent": []string{"Chrome"},
|
||||
}
|
||||
|
||||
if !reflect.DeepEqual(headers, expected) {
|
||||
@@ -146,8 +145,7 @@ func TestFuzzCrashers(t *testing.T) {
|
||||
}
|
||||
|
||||
for _, f := range crashers {
|
||||
ParseHeaders([][]byte{[]byte(f)}, func(header []byte, value []byte) {
|
||||
})
|
||||
ParseHeaders([]byte(f))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -156,17 +154,13 @@ func TestParseHeadersWithComplexUserAgent(t *testing.T) {
|
||||
// Parser should wait for \r\n
|
||||
payload := [][]byte{[]byte("POST /post HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.or"), []byte("g\r\nUser-Ag"), []byte("ent:Mozilla/5.0 (Windows NT 6.1; WOW64; Trident/7.0; rv:11.0) like Gecko\r\n\r\n"), []byte("Fake-Header: asda")}
|
||||
|
||||
headers := make(map[string]string)
|
||||
|
||||
ParseHeaders(payload, func(header []byte, value []byte) {
|
||||
headers[string(header)] = string(value)
|
||||
})
|
||||
headers := ParseHeaders(bytes.Join(payload, nil))
|
||||
|
||||
expected := map[string]string{
|
||||
"User-Agent": "Mozilla/5.0 (Windows NT 6.1; WOW64; Trident/7.0; rv:11.0) like Gecko",
|
||||
}
|
||||
|
||||
if expected["User-Agent"] != headers["User-Agent"] {
|
||||
if expected["User-Agent"] != headers["User-Agent"][0] {
|
||||
t.Errorf("Header 'User-Agent' expected '%s' and parsed: '%s'", expected["User-Agent"], headers["User-Agent"])
|
||||
}
|
||||
}
|
||||
@@ -176,11 +170,7 @@ func TestParseHeadersWithOrigin(t *testing.T) {
|
||||
// Parser should wait for \r\n
|
||||
payload := [][]byte{[]byte("POST /post HTTP/1.1\r\nContent-Length: 7\r\nHost: www.w3.or"), []byte("g\r\nReferrer: http://127.0.0.1:3000\r\nOrigi"), []byte("n: https://www.example.com\r\nUser-Ag"), []byte("ent:Mozilla/5.0 (Windows NT 6.1; WOW64; Trident/7.0; rv:11.0) like Gecko\r\n\r\n"), []byte("in:https://www.example.com\r\n\r\n"), []byte("Fake-Header: asda")}
|
||||
|
||||
headers := make(map[string]string)
|
||||
|
||||
ParseHeaders(payload, func(header []byte, value []byte) {
|
||||
headers[string(header)] = string(value)
|
||||
})
|
||||
headers := ParseHeaders(bytes.Join(payload, nil))
|
||||
|
||||
expected := map[string]string{
|
||||
"Origin": "https://www.example.com",
|
||||
@@ -188,15 +178,15 @@ func TestParseHeadersWithOrigin(t *testing.T) {
|
||||
"Referrer": "http://127.0.0.1:3000",
|
||||
}
|
||||
|
||||
if expected["Referrer"] != headers["Referrer"] {
|
||||
if expected["Referrer"] != headers["Referrer"][0] {
|
||||
t.Errorf("Header 'Referrer' expected '%s' and parsed: '%s'", expected["Referrer"], headers["Referrer"])
|
||||
}
|
||||
|
||||
if expected["Origin"] != headers["Origin"] {
|
||||
if expected["Origin"] != headers["Origin"][0] {
|
||||
t.Errorf("Header 'Origin' expected '%s' and parsed: '%s'", expected["Origin"], headers["Origin"])
|
||||
}
|
||||
|
||||
if expected["User-Agent"] != headers["User-Agent"] {
|
||||
if expected["User-Agent"] != headers["User-Agent"][0] {
|
||||
t.Errorf("Header 'User-Agent' expected '%s' and parsed: '%s'", expected["User-Agent"], headers["User-Agent"])
|
||||
}
|
||||
}
|
||||
@@ -328,7 +318,9 @@ func TestHasResponseTitle(t *testing.T) {
|
||||
"HTTP/1.1 100 Continue\r\n": true,
|
||||
"HTTP/1.1 \r\n": false,
|
||||
"HTTP/4.0 100Continue\r\n": false,
|
||||
"HTTP/4.0 100 Continue\r\n": false,
|
||||
"HTTP/1.0 10r Continue\r\n": false,
|
||||
"HTTP/1.1 200\r\n": false,
|
||||
"HTTP/1.1 200\r\nServer: Tengine\r\nContent-Length: 0\r\nConnection: close\r\n\r\n": false,
|
||||
}
|
||||
for k, v := range m {
|
||||
if HasResponseTitle([]byte(k)) != v {
|
||||
@@ -454,28 +446,21 @@ func TestHasFullPayload(t *testing.T) {
|
||||
}
|
||||
|
||||
func BenchmarkHasFullPayload(b *testing.B) {
|
||||
now := time.Now()
|
||||
payload := make([]byte, 0xfc00)
|
||||
for i := 0; i < 0xfc00; i++ {
|
||||
payload[i] = '1'
|
||||
var buf bytes.Buffer
|
||||
buf.Write([]byte("HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nTransfer-Encoding: chunked\r\n\r\n"))
|
||||
var chunk = []byte("1e\r\n111111111111111111111111111111\r\n")
|
||||
for i := 0; i < 5000; i++ {
|
||||
buf.Write(chunk)
|
||||
}
|
||||
buf.Write([]byte("0\r\n\r\n"))
|
||||
data := buf.Bytes()
|
||||
var ok bool
|
||||
data := []byte("HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nTransfer-Encoding: chunked\r\n\r\n")
|
||||
if ok = HasFullPayload(data); ok {
|
||||
b.Error("HasFullPayload should fail")
|
||||
return
|
||||
}
|
||||
b.ResetTimer() // ignores the upper initialization
|
||||
b.ReportMetric(float64(5000), "chunks/op")
|
||||
for i := 0; i < b.N; i++ {
|
||||
data = append(data, []byte(fmt.Sprintf("fc00\r\n%s\r\n", payload))...)
|
||||
if ok = HasFullPayload(data); ok {
|
||||
b.Error("HasFullPayload should fail")
|
||||
return
|
||||
}
|
||||
ok = HasFullPayload(data)
|
||||
}
|
||||
data = append(data, []byte("0\r\n\r\n")...)
|
||||
if ok = HasFullPayload(data); !ok {
|
||||
b.Error("HasFullPayload should pass")
|
||||
return
|
||||
if !ok {
|
||||
b.Fail()
|
||||
}
|
||||
b.Logf("%dKB chunks in %s", b.N*64, time.Since(now))
|
||||
}
|
||||
|
||||
+12
-10
@@ -66,25 +66,27 @@ 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)
|
||||
|
||||
if len(meta) < 2 {
|
||||
return
|
||||
}
|
||||
// id is encoded in hex, we need to revert to how it was
|
||||
id = make([]byte, 20)
|
||||
hex.Decode(id, meta[1])
|
||||
return
|
||||
return meta[1]
|
||||
}
|
||||
|
||||
func isOriginPayload(payload []byte) bool {
|
||||
switch payload[0] {
|
||||
case RequestPayload, ResponsePayload:
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
return payload[0] == RequestPayload || payload[0] == ResponsePayload
|
||||
}
|
||||
|
||||
func isRequestPayload(payload []byte) bool {
|
||||
|
||||
+47
-46
@@ -3,8 +3,10 @@ package main
|
||||
import (
|
||||
"flag"
|
||||
"fmt"
|
||||
"log"
|
||||
"os"
|
||||
"runtime"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
@@ -28,12 +30,10 @@ func (h *MultiOption) Set(value string) error {
|
||||
// AppSettings is the struct of main configuration
|
||||
type AppSettings struct {
|
||||
Verbose int `json:"verbose"`
|
||||
Stats bool `json:"stats"`
|
||||
ExitAfter time.Duration `json:"exit-after"`
|
||||
|
||||
SplitOutput bool `json:"split-output"`
|
||||
RecognizeTCPSessions bool `json:"recognize-tcp-sessions"`
|
||||
Pprof string `json:"http-pprof"`
|
||||
EmitterConfig
|
||||
Pprof string `json:"http-pprof"`
|
||||
|
||||
InputDummy MultiOption `json:"input-dummy"`
|
||||
OutputDummy MultiOption
|
||||
@@ -44,7 +44,6 @@ type AppSettings struct {
|
||||
InputTCPConfig TCPInputConfig
|
||||
OutputTCP MultiOption `json:"output-tcp"`
|
||||
OutputTCPConfig TCPOutputConfig
|
||||
OutputTCPStats bool `json:"output-tcp-stats"`
|
||||
|
||||
InputFile MultiOption `json:"input-file"`
|
||||
InputFileLoop bool `json:"input-file-loop"`
|
||||
@@ -81,11 +80,10 @@ func usage() {
|
||||
os.Exit(2)
|
||||
}
|
||||
|
||||
func init() {
|
||||
func initFlags() {
|
||||
flag.Usage = usage
|
||||
flag.StringVar(&Settings.Pprof, "http-pprof", "", "Enable profiling. Starts http server on specified port, exposing special /debug/pprof endpoint. Example: `:8181`")
|
||||
flag.IntVar(&Settings.Verbose, "verbose", 0, "set the level of verbosity, if greater than zero then it will turn on debug output")
|
||||
flag.BoolVar(&Settings.Stats, "stats", false, "Turn on queue stats output")
|
||||
|
||||
if DEMO == "" {
|
||||
flag.DurationVar(&Settings.ExitAfter, "exit-after", 0, "exit after specified duration")
|
||||
@@ -93,15 +91,12 @@ func init() {
|
||||
Settings.ExitAfter = 5 * time.Minute
|
||||
}
|
||||
|
||||
flag.BoolVar(&Settings.EmitterConfig.Stats, "meta-stats", false, "Turn on stats for meta data.")
|
||||
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")
|
||||
@@ -111,8 +106,9 @@ func init() {
|
||||
|
||||
flag.Var(&Settings.OutputTCP, "output-tcp", "Used for internal communication between Gor instances. Example: \n\t# Listen for requests on 80 port and forward them to other Gor instance on 28020 port\n\tgor --input-raw :80 --output-tcp replay.local:28020")
|
||||
flag.BoolVar(&Settings.OutputTCPConfig.Secure, "output-tcp-secure", false, "Use TLS secure connection. --input-file on another end should have TLS turned on as well.")
|
||||
flag.BoolVar(&Settings.OutputTCPConfig.SkipVerify, "output-tcp-skip-verify", false, "Don't verify hostname on TLS secure connection.")
|
||||
flag.BoolVar(&Settings.OutputTCPConfig.Sticky, "output-tcp-sticky", false, "Use Sticky connection. Request/Response with same ID will be sent to the same connection.")
|
||||
flag.BoolVar(&Settings.OutputTCPStats, "output-tcp-stats", false, "Report TCP output queue stats to console every 5 seconds.")
|
||||
flag.IntVar(&Settings.OutputTCPConfig.Workers, "output-tcp-workers", 10, "Number of parallel tcp connections, default is 10")
|
||||
|
||||
flag.Var(&Settings.InputFile, "input-file", "Read requests from file: \n\tgor --input-file ./requests.gor --output-http staging.com")
|
||||
flag.BoolVar(&Settings.InputFileLoop, "input-file-loop", false, "Loop input files, useful for performance testing.")
|
||||
@@ -143,7 +139,7 @@ func init() {
|
||||
flag.Var(&Settings.BufferSize, "input-raw-buffer-size", "Controls size of the OS buffer which holds packets until they dispatched. Default value depends by system: in Linux around 2MB. If you see big package drop, increase this value.")
|
||||
flag.BoolVar(&Settings.Promiscuous, "input-raw-promisc", false, "enable promiscuous mode")
|
||||
flag.BoolVar(&Settings.Monitor, "input-raw-monitor", false, "enable RF monitor mode")
|
||||
flag.BoolVar(&Settings.Stats, "input-raw-stats", false, "enable stats generator on raw TCP messages")
|
||||
flag.BoolVar(&Settings.RAWInputConfig.Stats, "input-raw-stats", false, "enable stats generator on raw TCP messages")
|
||||
|
||||
flag.StringVar(&Settings.Middleware, "middleware", "", "Used for modifying traffic using external command")
|
||||
|
||||
@@ -151,24 +147,20 @@ func init() {
|
||||
|
||||
/* outputHTTPConfig */
|
||||
flag.Var(&Settings.OutputHTTPConfig.BufferSize, "output-http-response-buffer", "HTTP response buffer size, all data after this size will be discarded.")
|
||||
flag.BoolVar(&Settings.OutputHTTPConfig.CompatibilityMode, "output-http-compatibility-mode", false, "Use standard Go client, instead of built-in implementation. Can be slower, but more compatible.")
|
||||
|
||||
flag.IntVar(&Settings.OutputHTTPConfig.WorkersMin, "output-http-workers-min", 0, "Gor uses dynamic worker scaling. Enter a number to set a minimum number of workers. default = 1.")
|
||||
flag.IntVar(&Settings.OutputHTTPConfig.WorkersMax, "output-http-workers", 0, "Gor uses dynamic worker scaling. Enter a number to set a maximum number of workers. default = 0 = unlimited.")
|
||||
flag.IntVar(&Settings.OutputHTTPConfig.QueueLen, "output-http-queue-len", 1000, "Number of requests that can be queued for output, if all workers are busy. default = 1000")
|
||||
|
||||
flag.BoolVar(&Settings.OutputHTTPConfig.SkipVerify, "output-http-skip-verify", false, "Don't verify hostname on TLS secure connection.")
|
||||
flag.DurationVar(&Settings.OutputHTTPConfig.WorkerTimeout, "output-http-worker-timeout", 2*time.Second, "Duration to rollback idle workers.")
|
||||
flag.IntVar(&Settings.OutputHTTPConfig.RedirectLimit, "output-http-redirects", 0, "Enable how often redirects should be followed.")
|
||||
flag.DurationVar(&Settings.OutputHTTPConfig.Timeout, "output-http-timeout", 5*time.Second, "Specify HTTP request/response timeout. By default 5s. Example: --output-http-timeout 30s")
|
||||
flag.BoolVar(&Settings.OutputHTTPConfig.TrackResponses, "output-http-track-response", false, "If turned on, HTTP output responses will be set to all outputs like stdout, file and etc.")
|
||||
|
||||
flag.BoolVar(&Settings.OutputHTTPConfig.Stats, "output-http-stats", false, "Report http output queue stats to console every N milliseconds. See output-http-stats-ms")
|
||||
flag.IntVar(&Settings.OutputHTTPConfig.StatsMs, "output-http-stats-ms", 5000, "Report http output queue stats to console every N milliseconds. default: 5000")
|
||||
flag.BoolVar(&Settings.OutputHTTPConfig.OriginalHost, "http-original-host", false, "Normally gor replaces the Host http header with the host supplied with --output-http. This option disables that behavior, preserving the original Host header.")
|
||||
flag.BoolVar(&Settings.OutputHTTPConfig.Debug, "output-http-debug", false, "Enables http debug output.")
|
||||
flag.StringVar(&Settings.OutputHTTPConfig.ElasticSearch, "output-http-elasticsearch", "", "Send request and response stats to ElasticSearch:\n\tgor --input-raw :8080 --output-http staging.com --output-http-elasticsearch 'es_host:api_port/index_name'")
|
||||
/* outputHTTPConfig */
|
||||
|
||||
flag.Var(&Settings.OutputBinary, "output-binary", "Forwards incoming binary payloads to given address.\n\t# Redirect all incoming requests to staging.com address \n\tgor --input-raw :80 --input-raw-protocol binary --output-binary staging.com:80")
|
||||
|
||||
/* outputBinaryConfig */
|
||||
flag.Var(&Settings.OutputBinaryConfig.BufferSize, "output-tcp-response-buffer", "TCP response buffer size, all data after this size will be discarded.")
|
||||
flag.IntVar(&Settings.OutputBinaryConfig.Workers, "output-binary-workers", 0, "Gor uses dynamic worker scaling by default. Enter a number to run a set number of workers.")
|
||||
@@ -187,36 +179,20 @@ func init() {
|
||||
flag.BoolVar(&Settings.InputKafkaConfig.UseJSON, "input-kafka-json-format", false, "If turned on, it will assume that messages coming in JSON format rather than GoReplay text format.")
|
||||
|
||||
flag.StringVar(&Settings.KafkaTLSConfig.CACert, "kafka-tls-ca-cert", "", "CA certificate for Kafka TLS Config:\n\tgor --input-raw :3000 --output-kafka-host '192.168.0.1:9092' --output-kafka-topic 'topic' --kafka-tls-ca-cert cacert.cer.pem --kafka-tls-client-cert client.cer.pem --kafka-tls-client-key client.key.pem")
|
||||
flag.StringVar(&Settings.KafkaTLSConfig.clientCert, "kafka-tls-client-cert", "", "Client certificate for Kafka TLS Config (mandatory with to kafka-tls-ca-cert and kafka-tls-client-key)")
|
||||
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.StringVar(&Settings.KafkaTLSConfig.ClientCert, "kafka-tls-client-cert", "", "Client certificate for Kafka TLS Config (mandatory with to kafka-tls-ca-cert and kafka-tls-client-key)")
|
||||
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
|
||||
@@ -236,16 +212,12 @@ func checkSettings() {
|
||||
if Settings.CopyBufferSize < 1 {
|
||||
Settings.CopyBufferSize.Set("5mb")
|
||||
}
|
||||
// libpcap has bug in mac os x. More info: https://github.com/buger/goreplay/issues/730
|
||||
if Settings.Expire == time.Second*2 && runtime.GOOS == "darwin" {
|
||||
Settings.Expire = time.Second
|
||||
}
|
||||
}
|
||||
|
||||
var previousDebugTime = time.Now()
|
||||
var debugMutex sync.Mutex
|
||||
|
||||
// Debug take an effect only if --verbose is greater than 0 specified
|
||||
// Debug take an effect only if --verbose greater than 0 is specified
|
||||
func Debug(level int, args ...interface{}) {
|
||||
if Settings.Verbose >= level {
|
||||
debugMutex.Lock()
|
||||
@@ -253,7 +225,36 @@ func Debug(level int, args ...interface{}) {
|
||||
now := time.Now()
|
||||
diff := now.Sub(previousDebugTime)
|
||||
previousDebugTime = now
|
||||
fmt.Printf("[DEBUG][elapsed %s]: ", diff)
|
||||
fmt.Println(args...)
|
||||
fmt.Fprintf(os.Stderr, "[DEBUG][elapsed %s]: ", diff)
|
||||
fmt.Fprintln(os.Stderr, args...)
|
||||
}
|
||||
}
|
||||
|
||||
func initGorDir() {
|
||||
var err error
|
||||
p := os.Getenv("GORDIR")
|
||||
if p == "" {
|
||||
if p, err = os.UserHomeDir(); err != nil {
|
||||
goto errLog
|
||||
}
|
||||
if !strings.Contains(p, "root") {
|
||||
p = filepath.Join(p, ".gor")
|
||||
} else {
|
||||
p = ".gor"
|
||||
}
|
||||
}
|
||||
if err = os.MkdirAll(p, os.ModeDir|os.ModePerm); err != nil {
|
||||
goto errLog
|
||||
}
|
||||
if err = os.Setenv("GORDIR", p); err != nil {
|
||||
goto errLog
|
||||
}
|
||||
return
|
||||
errLog:
|
||||
log.Fatalf("init GORDIR error:%q", err)
|
||||
}
|
||||
|
||||
func init() {
|
||||
initGorDir()
|
||||
initFlags()
|
||||
}
|
||||
|
||||
+1
-1
@@ -7,7 +7,7 @@ import (
|
||||
)
|
||||
|
||||
// Size represents size that implements flag.Var
|
||||
type Size int
|
||||
type Size int64
|
||||
|
||||
// the following regexes follow Go semantics https://golang.org/ref/spec#Letters_and_digits
|
||||
var (
|
||||
|
||||
@@ -0,0 +1,47 @@
|
||||
package csv
|
||||
|
||||
import (
|
||||
"encoding/csv"
|
||||
"os"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Writer check encoding/csv.Writer
|
||||
type Writer struct {
|
||||
*csv.Writer
|
||||
f *os.File
|
||||
}
|
||||
|
||||
// NewWriter returns an csv writer with file opened
|
||||
// error returned will be associated with opening the file
|
||||
func NewWriter(path string) (*Writer, error) {
|
||||
var wr Writer
|
||||
f, err := os.Create(path)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
wr.f = f
|
||||
wr.Writer = csv.NewWriter(f)
|
||||
return &wr, nil
|
||||
}
|
||||
|
||||
// WriteHeader will write infos about this docs in form of comments
|
||||
// it uses "#" to indicate that this should be ignored by the reader
|
||||
func (wr *Writer) WriteHeader(hdr string) error {
|
||||
_, err := wr.f.WriteString(formatHeader(hdr))
|
||||
return err
|
||||
}
|
||||
|
||||
// Close the underlying file
|
||||
func (wr *Writer) Close() error {
|
||||
return wr.f.Close()
|
||||
}
|
||||
|
||||
func formatHeader(hdr string) string {
|
||||
if !strings.HasSuffix(hdr, "\n") {
|
||||
hdr += "\n"
|
||||
}
|
||||
// add comments before every line
|
||||
hdr = strings.ReplaceAll("#"+hdr, "\n", "\n#")
|
||||
return hdr[:len(hdr)-1]
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
package csv
|
||||
|
||||
import (
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestFormatHeader(t *testing.T) {
|
||||
m := map[string]string{
|
||||
"hello": "#hello\n",
|
||||
"": "#\n",
|
||||
"hello\n": "#hello\n",
|
||||
"hello\nhello\n": "#hello\n#hello\n",
|
||||
"hello\nhello": "#hello\n#hello\n",
|
||||
"hello\nhello\nhello": "#hello\n#hello\n#hello\n",
|
||||
"hello\nhello\nhello\n": "#hello\n#hello\n#hello\n",
|
||||
"hello\nhello\r\n\r\nhello\n": "#hello\n#hello\r\n#\r\n#hello\n",
|
||||
}
|
||||
for k, v := range m {
|
||||
if vl := formatHeader(k); vl != v {
|
||||
t.Errorf("expected %q to equal %q", vl, v)
|
||||
}
|
||||
}
|
||||
}
|
||||
+124
@@ -0,0 +1,124 @@
|
||||
package stat
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ErrNotImplemented returned when impls requested are not implemnted
|
||||
var ErrNotImplemented = errors.New("method not implemented")
|
||||
|
||||
// RecordWriter represents a stat writer it should have all this method.
|
||||
// for rationale behind this interface check encoding/csv or
|
||||
type RecordWriter interface {
|
||||
Write(record []string) error
|
||||
Flush()
|
||||
Error() error
|
||||
}
|
||||
|
||||
const flushSpan = 10
|
||||
|
||||
// Writer represents a single instance of goreplay stats writer
|
||||
type Writer struct {
|
||||
name string
|
||||
records [flushSpan][]string // a temporary stats holder
|
||||
off int
|
||||
quit chan struct{}
|
||||
mu sync.Mutex
|
||||
w RecordWriter
|
||||
err error
|
||||
}
|
||||
|
||||
// NewWriter returns a new instance of stats associated with the writer "wr"
|
||||
// and "name" as the stats identity
|
||||
func NewWriter(w RecordWriter, name string) (wr *Writer) {
|
||||
wr = new(Writer)
|
||||
wr.name = name
|
||||
wr.w = w
|
||||
wr.quit = make(chan struct{})
|
||||
return
|
||||
}
|
||||
|
||||
// WriteHeader writes header of the documents to the underlying writer.
|
||||
// Header can be any kind of data like json unmarshal or file info.
|
||||
// underlying writer should sets a different way to represents the header.
|
||||
func (wr *Writer) WriteHeader(header string) error {
|
||||
if wHdr, ok := wr.w.(interface{ WriteHeader(header string) error }); ok {
|
||||
return wHdr.WriteHeader(header)
|
||||
}
|
||||
return ErrNotImplemented
|
||||
}
|
||||
|
||||
// SetReader sets another source of records reader to use along with (*Writer).Write([]string)
|
||||
// rd function will be called every "rate" duration time
|
||||
func (wr *Writer) SetReader(rd func() []string, rate time.Duration) {
|
||||
go func() {
|
||||
timer := time.NewTicker(rate)
|
||||
defer timer.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-wr.quit:
|
||||
return
|
||||
case <-timer.C:
|
||||
err := wr.Write(rd())
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func (wr *Writer) Write(record []string) error {
|
||||
wr.mu.Lock()
|
||||
defer wr.mu.Unlock()
|
||||
if wr.err != nil {
|
||||
return wr.err
|
||||
}
|
||||
if len(record) > 0 {
|
||||
wr.records[wr.off] = record
|
||||
wr.off++
|
||||
}
|
||||
return wr.flush(wr.off >= flushSpan)
|
||||
}
|
||||
|
||||
func (wr *Writer) String() string {
|
||||
return "stats: " + wr.name
|
||||
}
|
||||
|
||||
// Error reports any error that has occurred during a previous Write or Flush.
|
||||
func (wr *Writer) Error() error {
|
||||
wr.mu.Lock()
|
||||
wr.mu.Unlock()
|
||||
return wr.err
|
||||
}
|
||||
|
||||
// Close closes and flush all data of the stat
|
||||
func (wr *Writer) Close() error {
|
||||
wr.mu.Lock()
|
||||
defer wr.mu.Unlock()
|
||||
close(wr.quit)
|
||||
wr.err = wr.flush(true)
|
||||
if cl, ok := wr.w.(interface{ Close() error }); ok {
|
||||
wr.err = cl.Close()
|
||||
}
|
||||
return wr.err
|
||||
}
|
||||
|
||||
func (wr *Writer) flush(f bool) error {
|
||||
if !f {
|
||||
return nil
|
||||
}
|
||||
i := 0
|
||||
for i < wr.off {
|
||||
if wr.err = wr.w.Write(wr.records[i]); wr.err != nil {
|
||||
break
|
||||
}
|
||||
i++
|
||||
}
|
||||
wr.off -= i
|
||||
wr.w.Flush()
|
||||
wr.err = wr.w.Error()
|
||||
return wr.err
|
||||
}
|
||||
@@ -0,0 +1,33 @@
|
||||
package stat
|
||||
|
||||
import "testing"
|
||||
|
||||
type testWriter struct {
|
||||
f func([]string)
|
||||
fl func()
|
||||
}
|
||||
|
||||
func (t testWriter) WriteHeader(_ string) error { return nil }
|
||||
|
||||
func (t testWriter) Flush() { t.fl() }
|
||||
|
||||
func (t testWriter) Error() error { return nil }
|
||||
|
||||
func (t testWriter) Write(r []string) error {
|
||||
t.f(r)
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestTempSize(t *testing.T) {
|
||||
count := 0
|
||||
tW := testWriter{f: func(_ []string) { count++ }, fl: func() {
|
||||
if count != flushSpan {
|
||||
t.Errorf("should flush after %d records", flushSpan)
|
||||
}
|
||||
}}
|
||||
s := NewWriter(tW, "")
|
||||
defer s.Close()
|
||||
for i := 0; i < flushSpan; i++ {
|
||||
s.Write([]string{"1"})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"log"
|
||||
"os"
|
||||
"path/filepath"
|
||||
|
||||
"github.com/buger/goreplay/stat"
|
||||
"github.com/buger/goreplay/stat/csv"
|
||||
)
|
||||
|
||||
func statsWriteRecord(st *stat.Writer, record []string, plugin string) {
|
||||
if st == nil {
|
||||
return
|
||||
}
|
||||
if st.Error() != nil {
|
||||
Debug(1, plugin, " stats closed with error: ", st.Error())
|
||||
st.Close()
|
||||
st = nil
|
||||
return
|
||||
}
|
||||
st.Write(record)
|
||||
}
|
||||
|
||||
func initializeStatWriter(file string, hdr string) *stat.Writer {
|
||||
csvWr, err := csv.NewWriter(filepath.Join(os.Getenv("GORDIR"), file))
|
||||
if err != nil {
|
||||
log.Fatalf("init meta stat writer error: %q", err)
|
||||
}
|
||||
st := stat.NewWriter(csvWr, "meta")
|
||||
if hdr != "" {
|
||||
st.WriteHeader("type: GOR_STATS\ntitle: message meta")
|
||||
}
|
||||
return st
|
||||
}
|
||||
+105
-44
@@ -1,7 +1,7 @@
|
||||
package tcp
|
||||
|
||||
import (
|
||||
"crypto/sha1"
|
||||
"encoding/binary"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"sort"
|
||||
@@ -30,7 +30,7 @@ type Stats struct {
|
||||
type Message struct {
|
||||
packets []*Packet
|
||||
done chan bool
|
||||
data []byte
|
||||
pool *MessagePool
|
||||
Stats
|
||||
}
|
||||
|
||||
@@ -44,33 +44,46 @@ func NewMessage(srcAddr, dstAddr string, ipVersion uint8) (m *Message) {
|
||||
return
|
||||
}
|
||||
|
||||
// UUID the unique id of a TCP session it is not granted to be unique overtime
|
||||
// UUID returns the UUID of a TCP request and its response.
|
||||
func (m *Message) UUID() []byte {
|
||||
var src, dst string
|
||||
if m.IsIncoming {
|
||||
src = m.SrcAddr
|
||||
dst = m.DstAddr
|
||||
} else {
|
||||
src = m.DstAddr
|
||||
dst = m.SrcAddr
|
||||
var key uint64
|
||||
pckt := m.packets[0]
|
||||
|
||||
// check if response or request have generated the ID before.
|
||||
if m.pool.uuids != nil {
|
||||
lst := len(pckt.SrcIP) - 4
|
||||
if m.IsIncoming {
|
||||
key = uint64(pckt.SrcPort)<<48 | uint64(pckt.DstPort)<<32 |
|
||||
uint64(_uint32(pckt.SrcIP[lst:]))
|
||||
} else {
|
||||
key = uint64(pckt.DstPort)<<48 | uint64(pckt.SrcPort)<<32 |
|
||||
uint64(_uint32(pckt.DstIP[lst:]))
|
||||
}
|
||||
if uuidHex, ok := m.pool.uuids[key]; ok {
|
||||
delete(m.pool.uuids, key)
|
||||
return uuidHex
|
||||
}
|
||||
}
|
||||
|
||||
length := len(src) + len(dst)
|
||||
uuid := make([]byte, length)
|
||||
copy(uuid, src)
|
||||
copy(uuid[len(src):], dst)
|
||||
sha := sha1.Sum(uuid)
|
||||
uuid = make([]byte, 40)
|
||||
hex.Encode(uuid, sha[:])
|
||||
|
||||
return uuid
|
||||
id := make([]byte, 12, 12)
|
||||
binary.BigEndian.PutUint32(id, pckt.Seq)
|
||||
tStamp := m.End.UnixNano()
|
||||
binary.BigEndian.PutUint64(id[4:], uint64(tStamp))
|
||||
uuidHex := make([]byte, 24, 24)
|
||||
hex.Encode(uuidHex[:], id[:])
|
||||
if m.pool.uuids != nil {
|
||||
if len(m.pool.uuids) >= 1000 {
|
||||
m.pool.cleanUUIDs()
|
||||
}
|
||||
m.pool.uuids[key] = uuidHex
|
||||
}
|
||||
return uuidHex
|
||||
}
|
||||
|
||||
func (m *Message) add(pckt *Packet) {
|
||||
m.Length += len(pckt.Payload)
|
||||
m.LostData += int(pckt.Lost)
|
||||
m.packets = append(m.packets, pckt)
|
||||
m.data = append(m.data, pckt.Payload...)
|
||||
m.End = pckt.Timestamp
|
||||
}
|
||||
|
||||
@@ -81,7 +94,16 @@ func (m *Message) Packets() []*Packet {
|
||||
|
||||
// Data returns data in this message
|
||||
func (m *Message) Data() []byte {
|
||||
return m.data
|
||||
buf := make([]byte, m.Length, m.Length)
|
||||
var n, nn int
|
||||
for _, p := range m.packets {
|
||||
if len(p.Payload) < 1 {
|
||||
continue
|
||||
}
|
||||
nn = copy(buf[n:], p.Payload)
|
||||
n += nn
|
||||
}
|
||||
return buf
|
||||
}
|
||||
|
||||
// Sort a helper to sort packets
|
||||
@@ -101,17 +123,17 @@ type Debugger func(int, ...interface{})
|
||||
type HintEnd func(*Message) bool
|
||||
|
||||
// HintStart hints the pool to start the reassembling the message, see MessagePool.Start
|
||||
// when set, it will be used instead of checking SYN flag
|
||||
// when set, it will be called after checking SYN flag
|
||||
type HintStart func(*Packet) (IsIncoming, IsOutgoing bool)
|
||||
|
||||
// MessagePool holds data of all tcp messages in progress(still receiving/sending packets).
|
||||
// Incoming message is identified by its source port and address e.g: 127.0.0.1:45785.
|
||||
// Outgoing message is identified by server.addr and dst.addr e.g: localhost:80=internet:45785.
|
||||
// message is identified by its source port and dst port, and last 4bytes of src IP.
|
||||
type MessagePool struct {
|
||||
sync.Mutex
|
||||
debug Debugger
|
||||
maxSize size.Size // maximum message size, default 5mb
|
||||
pool map[string]*Message
|
||||
pool map[uint64]*Message
|
||||
uuids map[uint64][]byte
|
||||
handler Handler
|
||||
messageExpire time.Duration // the maximum time to wait for the final packet, minimum is 100ms
|
||||
End HintEnd
|
||||
@@ -131,7 +153,7 @@ func NewMessagePool(maxSize size.Size, messageExpire time.Duration, debugger Deb
|
||||
if pool.maxSize < 1 {
|
||||
pool.maxSize = 5 << 20
|
||||
}
|
||||
pool.pool = make(map[string]*Message)
|
||||
pool.pool = make(map[uint64]*Message)
|
||||
return pool
|
||||
}
|
||||
|
||||
@@ -139,26 +161,27 @@ func NewMessagePool(maxSize size.Size, messageExpire time.Duration, debugger Deb
|
||||
func (pool *MessagePool) Handler(packet gopacket.Packet) {
|
||||
var in, out bool
|
||||
pckt, err := ParsePacket(packet)
|
||||
if err != nil || pckt == nil {
|
||||
if err != nil {
|
||||
go pool.say(4, fmt.Sprintf("error decoding packet(%dBytes):%s\n", packet.Metadata().CaptureLength, err))
|
||||
return
|
||||
}
|
||||
if pckt == nil {
|
||||
return
|
||||
}
|
||||
pool.Lock()
|
||||
defer pool.Unlock()
|
||||
srcKey := pckt.Src()
|
||||
dstKey := srcKey + "=" + pckt.Dst()
|
||||
m, ok := pool.pool[srcKey]
|
||||
if !ok {
|
||||
m, ok = pool.pool[dstKey]
|
||||
}
|
||||
lst := len(pckt.SrcIP) - 4
|
||||
key := uint64(pckt.SrcPort)<<48 | uint64(pckt.DstPort)<<32 |
|
||||
uint64(_uint32(pckt.SrcIP[lst:]))
|
||||
m, ok := pool.pool[key]
|
||||
if pckt.RST {
|
||||
if ok {
|
||||
m.done <- true
|
||||
<-m.done
|
||||
}
|
||||
if m, ok = pool.pool[pckt.Dst()]; !ok {
|
||||
m, ok = pool.pool[pckt.Dst()+"="+srcKey]
|
||||
}
|
||||
key = uint64(pckt.DstPort)<<48 | uint64(pckt.SrcPort)<<32 |
|
||||
uint64(_uint32(pckt.DstIP[lst:]))
|
||||
m, ok = pool.pool[key]
|
||||
if ok {
|
||||
m.done <- true
|
||||
<-m.done
|
||||
@@ -179,25 +202,49 @@ func (pool *MessagePool) Handler(packet gopacket.Packet) {
|
||||
default:
|
||||
return
|
||||
}
|
||||
m = NewMessage(srcKey, pckt.Dst(), pckt.Version)
|
||||
m = NewMessage(pckt.Src(), pckt.Dst(), pckt.Version)
|
||||
m.IsIncoming = in
|
||||
key := srcKey
|
||||
if !m.IsIncoming {
|
||||
key = dstKey
|
||||
}
|
||||
pool.pool[key] = m
|
||||
m.Start = pckt.Timestamp
|
||||
m.pool = pool
|
||||
go pool.dispatch(key, m)
|
||||
pool.addPacket(m, pckt)
|
||||
}
|
||||
|
||||
func (pool *MessagePool) dispatch(key string, m *Message) {
|
||||
// MatchUUID instructs the pool to use same UUID for request and responses
|
||||
// this function should be called at initial stage of the pool
|
||||
func (pool *MessagePool) MatchUUID(match bool) {
|
||||
if match {
|
||||
pool.uuids = make(map[uint64][]byte)
|
||||
return
|
||||
}
|
||||
pool.uuids = nil
|
||||
}
|
||||
|
||||
// run GC on UUID map
|
||||
func (pool *MessagePool) cleanUUIDs() {
|
||||
var tStamp int64
|
||||
now := time.Now().UnixNano()
|
||||
for k, v := range pool.uuids {
|
||||
// there is a timestamp wrapped in every ID
|
||||
tStamp = int64(binary.BigEndian.Uint64(v[4:]))
|
||||
if time.Duration(now-tStamp) > pool.messageExpire {
|
||||
delete(pool.uuids, k)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (pool *MessagePool) dispatch(key uint64, m *Message) {
|
||||
select {
|
||||
case <-m.done:
|
||||
defer func() { m.done <- true }()
|
||||
defer func() { m.done <- true }() // signal that message was dispatched
|
||||
case <-time.After(pool.messageExpire):
|
||||
pool.Lock()
|
||||
defer pool.Unlock()
|
||||
// avoid dispathing message twice
|
||||
if _, ok := pool.pool[key]; !ok {
|
||||
return
|
||||
}
|
||||
m.TimedOut = true
|
||||
}
|
||||
delete(pool.pool, key)
|
||||
@@ -212,10 +259,20 @@ func (pool *MessagePool) addPacket(m *Message, pckt *Packet) {
|
||||
}
|
||||
m.add(pckt)
|
||||
switch {
|
||||
|
||||
// if one of this cases matches, we dispatch the message
|
||||
case trunc >= 0:
|
||||
case pckt.FIN:
|
||||
case pool.End != nil && pool.End(m):
|
||||
default:
|
||||
if pool.End != nil {
|
||||
if pool.End(m) {
|
||||
break
|
||||
}
|
||||
} else if m.Length > 0 && len(pckt.Payload) < 1 {
|
||||
// generally, data are chunked and sent sequentially in packets, if one packets is empty we assume
|
||||
// that it's just an acknowledgement packet after the final chunk.
|
||||
break
|
||||
}
|
||||
return
|
||||
}
|
||||
m.done <- true
|
||||
@@ -228,3 +285,7 @@ func (pool *MessagePool) say(level int, args ...interface{}) {
|
||||
pool.debug(level, args...)
|
||||
}
|
||||
}
|
||||
|
||||
func _uint32(b []byte) uint32 {
|
||||
return binary.BigEndian.Uint32(b)
|
||||
}
|
||||
|
||||
+21
-43
@@ -21,8 +21,11 @@ type Packet struct {
|
||||
gopacket.LinkLayer
|
||||
|
||||
// IP Header
|
||||
gopacket.NetworkLayer
|
||||
Version uint8 // Ip version
|
||||
SrcIP net.IP
|
||||
DstIP net.IP
|
||||
IHL uint8
|
||||
Length uint16
|
||||
|
||||
// TCP Segment Header
|
||||
*layers.TCP
|
||||
@@ -35,10 +38,12 @@ type Packet struct {
|
||||
// ParsePacket parse raw packets
|
||||
func ParsePacket(packet gopacket.Packet) (pckt *Packet, err error) {
|
||||
// early check of error
|
||||
if packet == nil {
|
||||
return
|
||||
}
|
||||
defer func() {
|
||||
if packet.ErrorLayer() != nil {
|
||||
err = packet.ErrorLayer().Error()
|
||||
println(err.Error())
|
||||
return
|
||||
}
|
||||
}()
|
||||
@@ -55,11 +60,17 @@ func ParsePacket(packet gopacket.Packet) (pckt *Packet, err error) {
|
||||
|
||||
// parsing network layer
|
||||
if net4, ok := packet.NetworkLayer().(*layers.IPv4); ok {
|
||||
pckt.NetworkLayer = net4
|
||||
pckt.Version = 4
|
||||
pckt.SrcIP = net4.SrcIP
|
||||
pckt.DstIP = net4.DstIP
|
||||
pckt.IHL = net4.IHL * 4
|
||||
pckt.Length = net4.Length
|
||||
} else if net6, ok := packet.NetworkLayer().(*layers.IPv6); ok {
|
||||
pckt.NetworkLayer = net6
|
||||
pckt.Version = 6
|
||||
pckt.SrcIP = net6.SrcIP
|
||||
pckt.DstIP = net6.DstIP
|
||||
pckt.IHL = 40
|
||||
pckt.Length = net6.Length
|
||||
} else {
|
||||
pckt = nil
|
||||
return
|
||||
@@ -75,56 +86,23 @@ func ParsePacket(packet gopacket.Packet) (pckt *Packet, err error) {
|
||||
pckt.DataOffset *= 4
|
||||
|
||||
// calculating lost data
|
||||
headerSize := int(uint32(pckt.DataOffset) + uint32(pckt.IHL()))
|
||||
headerSize := int(uint32(pckt.DataOffset) + uint32(pckt.IHL))
|
||||
if pckt.Version == 6 {
|
||||
headerSize -= 40 // in ipv6 the length of payload doesn't include the IPheader size
|
||||
}
|
||||
pckt.Lost = pckt.Length() - uint16(headerSize+len(pckt.Payload))
|
||||
pckt.Lost = pckt.Length - uint16(headerSize+len(pckt.Payload))
|
||||
|
||||
return
|
||||
}
|
||||
|
||||
// Src returns the source socket of a packet
|
||||
func (pckt *Packet) Src() string {
|
||||
return fmt.Sprintf("%s:%d", pckt.SrcIP(), pckt.SrcPort)
|
||||
return fmt.Sprintf("%s:%d", pckt.SrcIP, pckt.SrcPort)
|
||||
}
|
||||
|
||||
// Dst returns destination socket
|
||||
func (pckt *Packet) Dst() string {
|
||||
return fmt.Sprintf("%s:%d", pckt.DstIP(), pckt.DstPort)
|
||||
}
|
||||
|
||||
// SrcIP returns source IP address
|
||||
func (pckt *Packet) SrcIP() net.IP {
|
||||
if pckt.Version == 4 {
|
||||
return pckt.NetworkLayer.(*layers.IPv4).SrcIP
|
||||
}
|
||||
return pckt.NetworkLayer.(*layers.IPv6).SrcIP
|
||||
}
|
||||
|
||||
// DstIP returns destination IP address
|
||||
func (pckt *Packet) DstIP() net.IP {
|
||||
if pckt.Version == 4 {
|
||||
return pckt.NetworkLayer.(*layers.IPv4).DstIP
|
||||
}
|
||||
return pckt.NetworkLayer.(*layers.IPv6).DstIP
|
||||
}
|
||||
|
||||
// IHL returns IP header length in bytes
|
||||
func (pckt *Packet) IHL() uint8 {
|
||||
if l, ok := pckt.NetworkLayer.(*layers.IPv4); ok {
|
||||
return l.IHL * 4
|
||||
}
|
||||
// on IPV6 it's constant, https://en.wikipedia.org/wiki/IPv6_packet#Fixed_header
|
||||
return 40
|
||||
}
|
||||
|
||||
// Length returns the total length of the packet(IP header, TCP header and the actual data)
|
||||
func (pckt *Packet) Length() uint16 {
|
||||
if l, ok := pckt.NetworkLayer.(*layers.IPv4); ok {
|
||||
return l.Length
|
||||
}
|
||||
return pckt.NetworkLayer.(*layers.IPv6).Length
|
||||
return fmt.Sprintf("%s:%d", pckt.DstIP, pckt.DstPort)
|
||||
}
|
||||
|
||||
// SYNOptions returns MSS and windowscale of syn packets
|
||||
@@ -205,8 +183,8 @@ Lost Data: %d`,
|
||||
pckt.LinkInfo(),
|
||||
pckt.Src(),
|
||||
pckt.Dst(),
|
||||
pckt.IHL(),
|
||||
pckt.Length(),
|
||||
pckt.IHL,
|
||||
pckt.Length,
|
||||
pckt.Seq,
|
||||
pckt.Ack,
|
||||
pckt.DataOffset,
|
||||
|
||||
+82
-69
@@ -3,7 +3,6 @@ package tcp
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
@@ -52,9 +51,6 @@ func TestMessageParserWithHint(t *testing.T) {
|
||||
pool.Start = func(pckt *Packet) (bool, bool) {
|
||||
return proto.HasRequestTitle(pckt.Payload), proto.HasResponseTitle(pckt.Payload)
|
||||
}
|
||||
pool.End = func(m *Message) bool {
|
||||
return proto.HasFullPayload(m.Data())
|
||||
}
|
||||
packets := GetPackets(1, 30, nil)
|
||||
packets[0].Data()[14:][20:][13] = 2 // SYN flag
|
||||
packets[10].Data()[14:][20:][13] = 2 // SYN flag
|
||||
@@ -76,8 +72,8 @@ func TestMessageParserWithHint(t *testing.T) {
|
||||
return
|
||||
case m = <-mssg:
|
||||
}
|
||||
if len(m.packets) != 7 {
|
||||
t.Errorf("expected to have 7 packets got %d", len(m.packets))
|
||||
if len(m.packets) != 8 {
|
||||
t.Errorf("expected to have 8 packets got %d", len(m.packets))
|
||||
}
|
||||
if !bytes.HasSuffix(m.Data(), []byte("\n7\r\nNetwork\r\n0\r\n\r\n")) {
|
||||
t.Errorf("expected to %q to have suffix %q", m.Data(), []byte("\n7\r\nNetwork\r\n0\r\n\r\n"))
|
||||
@@ -89,8 +85,8 @@ func TestMessageParserWithHint(t *testing.T) {
|
||||
return
|
||||
case m = <-mssg:
|
||||
}
|
||||
if len(m.packets) != 7 {
|
||||
t.Errorf("expected to have 7 packets got %d", len(m.packets))
|
||||
if len(m.packets) != 8 {
|
||||
t.Errorf("expected to have 8 packets got %d", len(m.packets))
|
||||
}
|
||||
if !bytes.HasSuffix(m.Data(), []byte("Network")) {
|
||||
t.Errorf("expected to %q to have suffix %q", m.Data(), []byte("Network"))
|
||||
@@ -102,8 +98,8 @@ func TestMessageParserWithHint(t *testing.T) {
|
||||
return
|
||||
case m = <-mssg:
|
||||
}
|
||||
if len(m.packets) != 6 {
|
||||
t.Errorf("expected to have 6 packets got %d", len(m.packets))
|
||||
if len(m.packets) != 2 {
|
||||
t.Errorf("expected to have 2 packets got %d", len(m.packets))
|
||||
}
|
||||
if !bytes.HasSuffix(m.Data(), []byte("Content-Length: 0\r\n\r")) {
|
||||
t.Errorf("expected to %q to have suffix %q", m.Data(), []byte("Content-Length: 0\r\n\r"))
|
||||
@@ -192,88 +188,105 @@ func TestMessageTimeoutReached(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestMessageUUID(t *testing.T) {
|
||||
m1 := &Message{}
|
||||
m1.IsIncoming = true
|
||||
m1.SrcAddr = "src"
|
||||
m1.DstAddr = "dst"
|
||||
m2 := &Message{}
|
||||
m2.SrcAddr = "dst"
|
||||
m2.DstAddr = "src"
|
||||
if string(m1.UUID()) != string(m2.UUID()) {
|
||||
t.Errorf("expected %s, to equal %s", m1.UUID(), m2.UUID())
|
||||
packets := GetPackets(1, 10, nil)
|
||||
packets[0].Data()[14:][20:][13] = 2 // SYN flag
|
||||
packets[4].Data()[14:][20:][13] = 1 // FIN flag
|
||||
packets[5].Data()[14:][20:][13] = 2 // SYN flag
|
||||
packets[9].Data()[14:][20:][13] = 1 // FIN flag
|
||||
var uuid, uuid1 []byte
|
||||
pool := NewMessagePool(0, 0, nil, func(msg *Message) {
|
||||
if len(uuid) == 0 {
|
||||
uuid = msg.UUID()
|
||||
return
|
||||
}
|
||||
uuid1 = msg.UUID()
|
||||
})
|
||||
pool.MatchUUID(true)
|
||||
for _, p := range packets {
|
||||
pool.Handler(p)
|
||||
}
|
||||
|
||||
if string(uuid) != string(uuid1) {
|
||||
t.Errorf("expected %s, to equal %s", uuid, uuid1)
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkPacketParseAndSort(b *testing.B) {
|
||||
if b.N < 3 {
|
||||
return
|
||||
func BenchmarkMessageUUID(b *testing.B) {
|
||||
packets := GetPackets(1, 5, nil)
|
||||
packets[0].Data()[14:][20:][13] = 2 // SYN flag
|
||||
packets[4].Data()[14:][20:][13] = 1 // FIN flag
|
||||
var uuid []byte
|
||||
var msg *Message
|
||||
pool := NewMessagePool(0, 0, nil, func(m *Message) {
|
||||
msg = m
|
||||
})
|
||||
pool.MatchUUID(true)
|
||||
for _, p := range packets {
|
||||
pool.Handler(p)
|
||||
}
|
||||
now := time.Now()
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
uuid = msg.UUID()
|
||||
}
|
||||
_ = uuid
|
||||
}
|
||||
|
||||
func BenchmarkPacketParseAndSort(b *testing.B) {
|
||||
m := new(Message)
|
||||
m.packets = make([]*Packet, b.N)
|
||||
for i, v := range GetPackets(1, b.N, nil) {
|
||||
m.packets = make([]*Packet, 100)
|
||||
for i, v := range GetPackets(1, 100, nil) {
|
||||
m.packets[i], _ = ParsePacket(v)
|
||||
}
|
||||
m.Sort()
|
||||
b.Logf("%d packets in %s", b.N, time.Since(now))
|
||||
b.ResetTimer()
|
||||
for i := 0; i < b.N; i++ {
|
||||
m.Sort()
|
||||
}
|
||||
}
|
||||
|
||||
func BenchmarkMessageParserWithoutHint(b *testing.B) {
|
||||
var mssg = make(chan *Message, 1)
|
||||
if b.N < 3 {
|
||||
return
|
||||
}
|
||||
now := time.Now()
|
||||
n := b.N
|
||||
packets := GetPackets(1, n, nil)
|
||||
packets[0].Data()[14:][20:][13] = 2 // SYN flag
|
||||
packets[b.N-1].Data()[14:][20:][13] = 1 // FIN flag
|
||||
var chunk = []byte("111111111111111111111111111111")
|
||||
packets := GetPackets(1, 1000, chunk)
|
||||
packets[0].Data()[14:][20:][13] = 2 // SYN flag
|
||||
packets[1000-1].Data()[14:][20:][13] = 1 // FIN flag
|
||||
p := NewMessagePool(1<<20, time.Second*2, nil, func(m *Message) {
|
||||
b.Logf("%d/%d packets in %s", len(m.packets), n, time.Since(now))
|
||||
mssg <- m
|
||||
})
|
||||
for _, v := range packets {
|
||||
p.Handler(v)
|
||||
b.ResetTimer()
|
||||
b.ReportMetric(float64(1000), "packets/op")
|
||||
for i := 0; i < b.N; i++ {
|
||||
for _, v := range packets {
|
||||
p.Handler(v)
|
||||
}
|
||||
<-mssg
|
||||
}
|
||||
<-mssg
|
||||
}
|
||||
|
||||
func BenchmarkMessageParserWithHint(b *testing.B) {
|
||||
if b.N < 3 {
|
||||
return
|
||||
var buf [1003][]byte
|
||||
var chunk = []byte("1e\r\n111111111111111111111111111111\r\n")
|
||||
buf[0] = []byte("HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\nTransfer-Encoding: chunked\r\n\r\n")
|
||||
for i := 1; i < 1000; i++ {
|
||||
buf[i] = chunk
|
||||
}
|
||||
now := time.Now()
|
||||
n := b.N
|
||||
buf[1001] = []byte("0\r\n\r\n")
|
||||
packets := make([]gopacket.Packet, len(buf))
|
||||
for i := 0; i < len(buf); i++ {
|
||||
packets[i] = GetPackets(uint32(i+10), 1, buf[i])[0]
|
||||
}
|
||||
packets[1002] = GetPackets(1020, 1, nil)[0]
|
||||
var mssg = make(chan *Message, 1)
|
||||
payload := make([]byte, 0xfc00)
|
||||
for i := 0; i < 0xfc00; i++ {
|
||||
payload[i] = '1'
|
||||
}
|
||||
pool := NewMessagePool(1<<30, time.Second*10, nil, func(m *Message) { mssg <- m })
|
||||
pool.Start = func(pckt *Packet) (bool, bool) {
|
||||
return proto.HasRequestTitle(pckt.Payload), proto.HasResponseTitle(pckt.Payload)
|
||||
return false, proto.HasResponseTitle(pckt.Payload)
|
||||
}
|
||||
pool.End = func(m *Message) bool {
|
||||
return proto.HasFullPayload(m.Data())
|
||||
}
|
||||
pool.Handler(GetPackets(1, 1, []byte("POST / HTTP/1.1\r\nContent-Type: text/plain\r\nTransfer-Encoding: chunked\r\n\r\n"))[0])
|
||||
i := 0
|
||||
var d []byte
|
||||
for {
|
||||
select {
|
||||
case m := <-mssg:
|
||||
b.Logf("%d/%d packets, %dbytes, truncated: %v, timedout: %v in %s", len(m.packets), n, m.Length, m.Truncated, m.TimedOut, time.Since(now))
|
||||
return
|
||||
default:
|
||||
if i > n-2 {
|
||||
break
|
||||
} else if i < n-2 {
|
||||
d = []byte(fmt.Sprintf("fc00\r\n%s\r\n", payload))
|
||||
} else {
|
||||
d = []byte("0\r\n\r\n")
|
||||
}
|
||||
pool.Handler(GetPackets(1, i+2, d)[0])
|
||||
i++
|
||||
b.ResetTimer()
|
||||
b.ReportMetric(float64(len(packets)), "packets/op")
|
||||
b.ReportMetric(float64(1000), "chunks/op")
|
||||
for i := 0; i < b.N; i++ {
|
||||
for j := range packets {
|
||||
pool.Handler(packets[j])
|
||||
}
|
||||
<-mssg
|
||||
}
|
||||
}
|
||||
|
||||
+8
-9
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
+9
-8
@@ -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"
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user