Add support for injecting real ip (#296)

gor.go

@@ -7,15 +7,15 @@ import (
 	"fmt"
 	"io"
 	"log"
+	"net/http"
+	"net/http/httputil"
 	"os"
 	"os/signal"
 	"runtime"
 	_ "runtime/debug"
 	"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() {

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())
 	}

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) {

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

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()
 	}
-}
+}

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)
 	}
-}
+}

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 {

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)
+}

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")