Skip to content

Commit 06d129d

Browse files
bugerbuger
committed
Add support for injecting real ip (#296)
1 parent a3aeca2 commit 06d129d

File tree

9 files changed

+47
-28
lines changed

9 files changed

+47
-28
lines changed

‎gor.go

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -7,15 +7,15 @@ import (
77
"fmt"
88
"io"
99
"log"
10+
"net/http"
11+
"net/http/httputil"
1012
"os"
1113
"os/signal"
1214
"runtime"
1315
_ "runtime/debug"
1416
"runtime/pprof"
1517
"syscall"
1618
"time"
17-
"net/http"
18-
"net/http/httputil"
1919
)
2020

2121
var (
@@ -25,11 +25,11 @@ var (
2525
)
2626

2727
func loggingMiddleware(next http.Handler) http.Handler {
28-
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
29-
rb, _ := httputil.DumpRequest(r, false)
30-
log.Println(string(rb))
31-
next.ServeHTTP(w, r)
32-
})
28+
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
29+
rb, _ := httputil.DumpRequest(r, false)
30+
log.Println(string(rb))
31+
next.ServeHTTP(w, r)
32+
})
3333
}
3434

3535
func main() {

‎input_raw.go

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package main
22

33
import (
4+
"github.com/buger/gor/proto"
45
raw "github.com/buger/gor/raw_socket_listener"
56
"log"
67
"net"
@@ -14,6 +15,7 @@ type RAWInput struct {
1415
expire time.Duration
1516
quit chan bool
1617
engine int
18+
realIPHeader []byte
1719
trackResponse bool
1820
listener *raw.Listener
1921
}
@@ -25,12 +27,13 @@ const (
2527
)
2628

2729
// NewRAWInput constructor for RAWInput. Accepts address with port as argument.
28-
func NewRAWInput(address string, engine int, trackResponse bool, expire time.Duration) (i *RAWInput) {
30+
func NewRAWInput(address string, engine int, trackResponse bool, expire time.Duration, realIPHeader string) (i *RAWInput) {
2931
i = new(RAWInput)
3032
i.data = make(chan *raw.TCPMessage)
3133
i.address = address
3234
i.expire = expire
3335
i.engine = engine
36+
i.realIPHeader = []byte(realIPHeader)
3437
i.quit = make(chan bool)
3538
i.trackResponse = trackResponse
3639

@@ -48,6 +51,9 @@ func (i *RAWInput) Read(data []byte) (int, error) {
4851

4952
if msg.IsIncoming {
5053
header = payloadHeader(RequestPayload, msg.UUID(), msg.Start.UnixNano())
54+
if len(i.realIPHeader) > 0 {
55+
buf = proto.SetHeader(buf, i.realIPHeader, []byte(msg.IP().String()))
56+
}
5157
} else {
5258
header = payloadHeader(ResponsePayload, msg.UUID(), msg.End.UnixNano()-msg.AssocMessage.Start.UnixNano())
5359
}

‎input_raw_test.go

Lines changed: 12 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package main
22

33
import (
44
"bytes"
5+
"github.com/buger/gor/proto"
56
"io"
67
"io/ioutil"
78
"log"
@@ -22,7 +23,7 @@ import (
2223

2324
const testRawExpire = time.Millisecond * 200
2425

25-
func TestRAWInput(t *testing.T) {
26+
func TestRAWInputIPv4(t *testing.T) {
2627
wg := new(sync.WaitGroup)
2728
quit := make(chan int)
2829

@@ -42,11 +43,15 @@ func TestRAWInput(t *testing.T) {
4243

4344
var respCounter, reqCounter int64
4445

45-
input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire)
46+
input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "X-Real-IP")
4647
defer input.Close()
4748

4849
output := NewTestOutput(func(data []byte) {
4950
if data[0] == '1' {
51+
body := payloadBody(data)
52+
if len(proto.Header(body, []byte("X-Real-IP"))) == 0 {
53+
t.Error("Should have X-Real-IP header", string(body))
54+
}
5055
atomic.AddInt64(&reqCounter, 1)
5156
} else {
5257
atomic.AddInt64(&respCounter, 1)
@@ -97,7 +102,7 @@ func TestRAWInputIPv6(t *testing.T) {
97102

98103
var respCounter, reqCounter int64
99104

100-
input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire)
105+
input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "")
101106
defer input.Close()
102107

103108
output := NewTestOutput(func(data []byte) {
@@ -148,7 +153,7 @@ func TestInputRAW100Expect(t *testing.T) {
148153

149154
originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
150155

151-
input := NewRAWInput(originAddr, EnginePcap, true, time.Second)
156+
input := NewRAWInput(originAddr, EnginePcap, true, time.Second, "")
152157
defer input.Close()
153158

154159
// We will use it to get content of raw HTTP request
@@ -211,7 +216,7 @@ func TestInputRAWChunkedEncoding(t *testing.T) {
211216
}))
212217

213218
originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
214-
input := NewRAWInput(originAddr, EnginePcap, true, time.Second)
219+
input := NewRAWInput(originAddr, EnginePcap, true, time.Second, "")
215220
defer input.Close()
216221

217222
replay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -275,7 +280,7 @@ func TestInputRAWLargePayload(t *testing.T) {
275280
}))
276281
originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
277282

278-
input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire)
283+
input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "")
279284
defer input.Close()
280285

281286
replay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
@@ -320,7 +325,7 @@ func BenchmarkRAWInput(b *testing.B) {
320325

321326
var respCounter, reqCounter int64
322327

323-
input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire)
328+
input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "")
324329
defer input.Close()
325330

326331
output := NewTestOutput(func(data []byte) {

‎middleware_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -117,7 +117,7 @@ func TestEchoMiddleware(t *testing.T) {
117117

118118
// Catch traffic from one service
119119
fromAddr := strings.Replace(from.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
120-
input := NewRAWInput(fromAddr, EnginePcap, true, testRawExpire)
120+
input := NewRAWInput(fromAddr, EnginePcap, true, testRawExpire, "")
121121
defer input.Close()
122122

123123
// And redirect to another
@@ -179,7 +179,7 @@ func TestTokenMiddleware(t *testing.T) {
179179

180180
fromAddr := strings.Replace(from.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
181181
// Catch traffic from one service
182-
input := NewRAWInput(fromAddr, EnginePcap, true, testRawExpire)
182+
input := NewRAWInput(fromAddr, EnginePcap, true, testRawExpire, "")
183183
defer input.Close()
184184

185185
// And redirect to another

‎output_file.go

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -104,19 +104,19 @@ func withoutIndex(s string) string {
104104
type sortByFileIndex []string
105105

106106
func (s sortByFileIndex) Len() int {
107-
return len(s)
107+
return len(s)
108108
}
109109

110110
func (s sortByFileIndex) Swap(i, j int) {
111-
s[i], s[j] = s[j], s[i]
111+
s[i], s[j] = s[j], s[i]
112112
}
113113

114114
func (s sortByFileIndex) Less(i, j int) bool {
115-
if withoutIndex(s[i]) == withoutIndex(s[j]) {
116-
return getFileIndex(s[i]) < getFileIndex(s[j])
117-
}
115+
if withoutIndex(s[i]) == withoutIndex(s[j]) {
116+
return getFileIndex(s[i]) < getFileIndex(s[j])
117+
}
118118

119-
return s[i] < s[j]
119+
return s[i] < s[j]
120120
}
121121

122122
func (o *FileOutput) filename() string {
@@ -225,4 +225,4 @@ func (o *FileOutput) Close() {
225225
}
226226
o.file.Close()
227227
}
228-
}
228+
}

‎output_file_test.go

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,12 @@ import (
55
"io"
66
"math/rand"
77
"os"
8+
"reflect"
9+
"sort"
810
"sync"
911
"sync/atomic"
1012
"testing"
1113
"time"
12-
"sort"
13-
"reflect"
1414
)
1515

1616
func TestFileOutput(t *testing.T) {
@@ -270,4 +270,4 @@ func TestFileOutputSort(t *testing.T) {
270270
if !reflect.DeepEqual(files, expected) {
271271
t.Error("Should properly sort file names using indexes", files, expected)
272272
}
273-
}
273+
}

‎plugins.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ func InitPlugins() {
101101
}
102102

103103
for _, options := range Settings.inputRAW {
104-
registerPlugin(NewRAWInput, options, engine, Settings.inputRAWTrackResponse, time.Duration(0))
104+
registerPlugin(NewRAWInput, options, engine, Settings.inputRAWTrackResponse, time.Duration(0), Settings.inputRAWRealIPHeader)
105105
}
106106

107107
for _, options := range Settings.inputTCP {

‎raw_socket_listener/tcp_message.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import (
99
"log"
1010
"strconv"
1111
"time"
12+
"net"
1213
)
1314

1415
var _ = log.Println
@@ -267,3 +268,7 @@ func (t *TCPMessage) UpdateResponseAck() uint32 {
267268
func (t *TCPMessage) ID() tcpID {
268269
return t.packets[0].ID
269270
}
271+
272+
func (t *TCPMessage) IP() net.IP {
273+
return net.IP(t.packets[0].Addr)
274+
}

‎settings.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ type AppSettings struct {
4747
inputRAW MultiOption
4848
inputRAWEngine string
4949
inputRAWTrackResponse bool
50+
inputRAWRealIPHeader string
5051

5152
middleware string
5253

@@ -102,6 +103,8 @@ func init() {
102103

103104
flag.StringVar(&Settings.inputRAWEngine, "input-raw-engine", "libpcap", "Intercept traffic using `libpcap` (default), and `raw_socket`")
104105

106+
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")
107+
105108
flag.StringVar(&Settings.middleware, "middleware", "", "Used for modifying traffic using external command")
106109

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

0 commit comments

Comments
 (0)