From 06d129d20d6fc90c8e42ad6319af55472c19025f Mon Sep 17 00:00:00 2001 From: Leonid Bugaev Date: Thu, 9 Jun 2016 19:06:59 +0600 Subject: [PATCH] Add support for injecting real ip (#296) --- gor.go | 14 +++++++------- input_raw.go | 8 +++++++- input_raw_test.go | 19 ++++++++++++------- middleware_test.go | 4 ++-- output_file.go | 14 +++++++------- output_file_test.go | 6 +++--- plugins.go | 2 +- raw_socket_listener/tcp_message.go | 5 +++++ settings.go | 3 +++ 9 files changed, 47 insertions(+), 28 deletions(-) diff --git a/gor.go b/gor.go index 6a5f794a1..8788a7457 100644 --- a/gor.go +++ b/gor.go @@ -7,6 +7,8 @@ import ( "fmt" "io" "log" + "net/http" + "net/http/httputil" "os" "os/signal" "runtime" @@ -14,8 +16,6 @@ import ( "runtime/pprof" "syscall" "time" - "net/http" - "net/http/httputil" ) var ( @@ -25,11 +25,11 @@ var ( ) func loggingMiddleware(next http.Handler) http.Handler { - return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - rb, _ := httputil.DumpRequest(r, false) - log.Println(string(rb)) - next.ServeHTTP(w, r) - }) + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + rb, _ := httputil.DumpRequest(r, false) + log.Println(string(rb)) + next.ServeHTTP(w, r) + }) } func main() { diff --git a/input_raw.go b/input_raw.go index 92fb789a6..c67cd8cd9 100644 --- a/input_raw.go +++ b/input_raw.go @@ -1,6 +1,7 @@ package main import ( + "github.com/buger/gor/proto" raw "github.com/buger/gor/raw_socket_listener" "log" "net" @@ -14,6 +15,7 @@ type RAWInput struct { expire time.Duration quit chan bool engine int + realIPHeader []byte trackResponse bool listener *raw.Listener } @@ -25,12 +27,13 @@ const ( ) // NewRAWInput constructor for RAWInput. Accepts address with port as argument. -func NewRAWInput(address string, engine int, trackResponse bool, expire time.Duration) (i *RAWInput) { +func NewRAWInput(address string, engine int, trackResponse bool, expire time.Duration, realIPHeader string) (i *RAWInput) { i = new(RAWInput) i.data = make(chan *raw.TCPMessage) i.address = address i.expire = expire i.engine = engine + i.realIPHeader = []byte(realIPHeader) i.quit = make(chan bool) i.trackResponse = trackResponse @@ -48,6 +51,9 @@ func (i *RAWInput) Read(data []byte) (int, error) { if msg.IsIncoming { header = payloadHeader(RequestPayload, msg.UUID(), msg.Start.UnixNano()) + if len(i.realIPHeader) > 0 { + buf = proto.SetHeader(buf, i.realIPHeader, []byte(msg.IP().String())) + } } else { header = payloadHeader(ResponsePayload, msg.UUID(), msg.End.UnixNano()-msg.AssocMessage.Start.UnixNano()) } diff --git a/input_raw_test.go b/input_raw_test.go index 7929c9248..b8fe1b114 100644 --- a/input_raw_test.go +++ b/input_raw_test.go @@ -2,6 +2,7 @@ package main import ( "bytes" + "github.com/buger/gor/proto" "io" "io/ioutil" "log" @@ -22,7 +23,7 @@ import ( const testRawExpire = time.Millisecond * 200 -func TestRAWInput(t *testing.T) { +func TestRAWInputIPv4(t *testing.T) { wg := new(sync.WaitGroup) quit := make(chan int) @@ -42,11 +43,15 @@ func TestRAWInput(t *testing.T) { var respCounter, reqCounter int64 - input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire) + input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "X-Real-IP") defer input.Close() 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)) + } atomic.AddInt64(&reqCounter, 1) } else { atomic.AddInt64(&respCounter, 1) @@ -97,7 +102,7 @@ func TestRAWInputIPv6(t *testing.T) { var respCounter, reqCounter int64 - input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire) + input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "") defer input.Close() output := NewTestOutput(func(data []byte) { @@ -148,7 +153,7 @@ func TestInputRAW100Expect(t *testing.T) { originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1) - input := NewRAWInput(originAddr, EnginePcap, true, time.Second) + input := NewRAWInput(originAddr, EnginePcap, true, time.Second, "") defer input.Close() // We will use it to get content of raw HTTP request @@ -211,7 +216,7 @@ func TestInputRAWChunkedEncoding(t *testing.T) { })) originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1) - input := NewRAWInput(originAddr, EnginePcap, true, time.Second) + input := NewRAWInput(originAddr, EnginePcap, true, time.Second, "") defer input.Close() replay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -275,7 +280,7 @@ func TestInputRAWLargePayload(t *testing.T) { })) originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1) - input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire) + input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "") defer input.Close() replay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { @@ -320,7 +325,7 @@ func BenchmarkRAWInput(b *testing.B) { var respCounter, reqCounter int64 - input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire) + input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "") defer input.Close() output := NewTestOutput(func(data []byte) { diff --git a/middleware_test.go b/middleware_test.go index 69ca99db2..beaed15ac 100644 --- a/middleware_test.go +++ b/middleware_test.go @@ -117,7 +117,7 @@ func TestEchoMiddleware(t *testing.T) { // Catch traffic from one service fromAddr := strings.Replace(from.Listener.Addr().String(), "[::]", "127.0.0.1", -1) - input := NewRAWInput(fromAddr, EnginePcap, true, testRawExpire) + input := NewRAWInput(fromAddr, EnginePcap, true, testRawExpire, "") defer input.Close() // And redirect to another @@ -179,7 +179,7 @@ func TestTokenMiddleware(t *testing.T) { fromAddr := strings.Replace(from.Listener.Addr().String(), "[::]", "127.0.0.1", -1) // Catch traffic from one service - input := NewRAWInput(fromAddr, EnginePcap, true, testRawExpire) + input := NewRAWInput(fromAddr, EnginePcap, true, testRawExpire, "") defer input.Close() // And redirect to another diff --git a/output_file.go b/output_file.go index 9e99c94ae..5b797d8c7 100644 --- a/output_file.go +++ b/output_file.go @@ -104,19 +104,19 @@ func withoutIndex(s string) string { type sortByFileIndex []string func (s sortByFileIndex) Len() int { - return len(s) + return len(s) } func (s sortByFileIndex) Swap(i, j int) { - s[i], s[j] = s[j], s[i] + s[i], s[j] = s[j], s[i] } func (s sortByFileIndex) Less(i, j int) bool { - if withoutIndex(s[i]) == withoutIndex(s[j]) { - return getFileIndex(s[i]) < getFileIndex(s[j]) - } + if withoutIndex(s[i]) == withoutIndex(s[j]) { + return getFileIndex(s[i]) < getFileIndex(s[j]) + } - return s[i] < s[j] + return s[i] < s[j] } func (o *FileOutput) filename() string { @@ -225,4 +225,4 @@ func (o *FileOutput) Close() { } o.file.Close() } -} \ No newline at end of file +} diff --git a/output_file_test.go b/output_file_test.go index 765045c45..ec7935101 100644 --- a/output_file_test.go +++ b/output_file_test.go @@ -5,12 +5,12 @@ import ( "io" "math/rand" "os" + "reflect" + "sort" "sync" "sync/atomic" "testing" "time" - "sort" - "reflect" ) func TestFileOutput(t *testing.T) { @@ -270,4 +270,4 @@ func TestFileOutputSort(t *testing.T) { if !reflect.DeepEqual(files, expected) { t.Error("Should properly sort file names using indexes", files, expected) } -} \ No newline at end of file +} diff --git a/plugins.go b/plugins.go index e85dee966..531ff544b 100644 --- a/plugins.go +++ b/plugins.go @@ -101,7 +101,7 @@ func InitPlugins() { } for _, options := range Settings.inputRAW { - registerPlugin(NewRAWInput, options, engine, Settings.inputRAWTrackResponse, time.Duration(0)) + registerPlugin(NewRAWInput, options, engine, Settings.inputRAWTrackResponse, time.Duration(0), Settings.inputRAWRealIPHeader) } for _, options := range Settings.inputTCP { diff --git a/raw_socket_listener/tcp_message.go b/raw_socket_listener/tcp_message.go index 71c14aeb2..f533ad61c 100644 --- a/raw_socket_listener/tcp_message.go +++ b/raw_socket_listener/tcp_message.go @@ -9,6 +9,7 @@ import ( "log" "strconv" "time" + "net" ) var _ = log.Println @@ -267,3 +268,7 @@ func (t *TCPMessage) UpdateResponseAck() uint32 { func (t *TCPMessage) ID() tcpID { return t.packets[0].ID } + +func (t *TCPMessage) IP() net.IP { + return net.IP(t.packets[0].Addr) +} \ No newline at end of file diff --git a/settings.go b/settings.go index 72a8c770a..702dcfc20 100644 --- a/settings.go +++ b/settings.go @@ -47,6 +47,7 @@ type AppSettings struct { inputRAW MultiOption inputRAWEngine string inputRAWTrackResponse bool + inputRAWRealIPHeader string middleware string @@ -102,6 +103,8 @@ func init() { flag.StringVar(&Settings.inputRAWEngine, "input-raw-engine", "libpcap", "Intercept traffic using `libpcap` (default), and `raw_socket`") + flag.StringVar(&Settings.inputRAWRealIPHeader, "input-raw-realip-header", "", "If not blank, injects header with given name and real IP value to the request payload. Usually this header should be named: X-Real-IP") + flag.StringVar(&Settings.middleware, "middleware", "", "Used for modifying traffic using external command") flag.Var(&Settings.inputHTTP, "input-http", "Read requests from HTTP, should be explicitly sent from your application:\n\t# Listen for http on 9000\n\tgor --input-http :9000 --output-http staging.com")