diff --git a/README.md b/README.md index e95d5fb..225c2d3 100644 --- a/README.md +++ b/README.md @@ -32,7 +32,7 @@ gor replay -f http://staging.server -p 28020 ## Advanced use ### Rate limiting -The replay server supports rate limiting. It can be useful if you want +Both replay and listener supports rate limiting. It can be useful if you want forward only part of production traffic and not overload your staging environment. You can specify your desired requests per second using the "|" operator after the server address: @@ -42,6 +42,12 @@ environment. You can specify your desired requests per second using the gor replay -f "http://staging.server|10" ``` +``` +# replay server will not get more than 10 requests per second +# useful for high-load environments +gor listen -p 8080 -r "replay.server.local:28020|10" +``` + ### Forward to multiple addresses You can forward traffic to multiple endpoints. Just separate the addresses by comma. diff --git a/integration_test.go b/integration_test.go index 5f97347..c5c44e9 100644 --- a/integration_test.go +++ b/integration_test.go @@ -8,32 +8,79 @@ import ( "time" + "fmt" "net/http" "strconv" + + "sync/atomic" ) -func startListener() { - listener.Settings.Verbose = true +func isEqual(t *testing.T, a interface{}, b interface{}) { + if a != b { + t.Error("Original and Replayed request not match\n", a, "!=", b) + } +} + +var envs int + +type Env struct { + Verbose bool + + ListenHandler http.HandlerFunc + ReplayHandler http.HandlerFunc + + ReplayLimit int + ListenerLimit int +} + +func (e *Env) start() (p int) { + p = 50000 + envs*10 + + go e.startHTTP(p, http.HandlerFunc(e.ListenHandler)) + go e.startHTTP(p+2, http.HandlerFunc(e.ReplayHandler)) + go e.startListener(p, p+1) + go e.startReplay(p+1, p+2) + + // Time to start http and gor instances + time.Sleep(time.Millisecond * 100) + + envs++ + + return +} + +func (e *Env) startListener(port int, replayPort int) { + listener.Settings.Verbose = e.Verbose listener.Settings.Address = "127.0.0.1" - listener.Settings.ReplayAddress = "127.0.0.1:50001" - listener.Settings.Port = 50000 - go listener.Run() + listener.Settings.ReplayAddress = "127.0.0.1:" + strconv.Itoa(replayPort) + listener.Settings.Port = port + + if e.ListenerLimit != 0 { + listener.Settings.ReplayAddress += "|" + strconv.Itoa(e.ListenerLimit) + } + + listener.Run() } -func startReplay() { - replay.Settings.Verbose = true +func (e *Env) startReplay(port int, forwardPort int) { + replay.Settings.Verbose = e.Verbose replay.Settings.Host = "127.0.0.1" - replay.Settings.ForwardAddress = "127.0.0.1:50002" - replay.Settings.Port = 50001 - go replay.Run() + replay.Settings.ForwardAddress = "127.0.0.1:" + strconv.Itoa(forwardPort) + replay.Settings.Port = port + + if e.ReplayLimit != 0 { + replay.Settings.ForwardAddress += "|" + strconv.Itoa(e.ReplayLimit) + } + + replay.Run() } -func startHTTP(port int, handler http.Handler) { - go http.ListenAndServe(":"+strconv.Itoa(port), handler) +func (e *Env) startHTTP(port int, handler http.Handler) { + http.ListenAndServe(":"+strconv.Itoa(port), handler) } -func getRequest() *http.Request { - req, _ := http.NewRequest("GET", "http://localhost:50000/test", nil) +func getRequest(port int) *http.Request { + req, _ := http.NewRequest("GET", "http://localhost:"+strconv.Itoa(port)+"/test", nil) ck1 := new(http.Cookie) ck1.Name = "test" ck1.Value = "value" @@ -43,36 +90,35 @@ func getRequest() *http.Request { return req } -func TestIntegration(t *testing.T) { - request := getRequest() - - listenHandler := func(w http.ResponseWriter, r *http.Request) { - http.Error(w, "404 page not found", http.StatusNotFound) - } - startHTTP(50000, http.HandlerFunc(listenHandler)) - - startListener() - startReplay() - +func TestReplay(t *testing.T) { + var request *http.Request received := make(chan int) - replayHandler := func(w http.ResponseWriter, r *http.Request) { - equal := func(a interface{}, b interface{}) { - if a != b { - t.Error("Original and Replayed request not match\n", a, "!=", b, "\nReplayed:", r, "\nOriginal:", request) - } - } + listenHandler := func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "OK", http.StatusNotFound) + } - equal(r.URL.Path, request.URL.Path) - equal(r.Cookies()[0].Value, request.Cookies()[0].Value) + replayHandler := func(w http.ResponseWriter, r *http.Request) { + isEqual(t, r.URL.Path, request.URL.Path) + isEqual(t, r.Cookies()[0].Value, request.Cookies()[0].Value) http.Error(w, "404 page not found", http.StatusNotFound) + if t.Failed() { + fmt.Println("\nReplayed:", r, "\nOriginal:", request) + } + received <- 1 } - startHTTP(50002, http.HandlerFunc(replayHandler)) - time.Sleep(time.Millisecond * 100) + env := &Env{ + Verbose: true, + ListenHandler: listenHandler, + ReplayHandler: replayHandler, + } + p := env.start() + + request = getRequest(p) _, err := http.DefaultClient.Do(request) @@ -85,6 +131,59 @@ func TestIntegration(t *testing.T) { case <-time.After(time.Second): t.Error("Timeout error") } +} + +func rateLimitEnv(replayLimit int, listenerLimit int) int32 { + var processed int32 + + listenHandler := func(w http.ResponseWriter, r *http.Request) { + http.Error(w, "OK", http.StatusAccepted) + } + + replayHandler := func(w http.ResponseWriter, r *http.Request) { + atomic.AddInt32(&processed, 1) + http.Error(w, "OK", http.StatusAccepted) + } + + env := &Env{ + ListenHandler: listenHandler, + ReplayHandler: replayHandler, + ReplayLimit: replayLimit, + ListenerLimit: listenerLimit, + } + + p := env.start() + req := getRequest(p) + + for i := 0; i < 10; i++ { + http.DefaultClient.Do(req) + } time.Sleep(time.Millisecond * 500) + + return processed +} + +func TestWithoutReplayRateLimit(t *testing.T) { + processed := rateLimitEnv(0, 0) + + if processed != 10 { + t.Error("It should forward all requests without rate-limiting", processed) + } +} + +func TestReplayRateLimit(t *testing.T) { + processed := rateLimitEnv(5, 0) + + if processed != 5 { + t.Error("It should forward only 5 requests with rate-limiting", processed) + } +} + +func TestListenerRateLimit(t *testing.T) { + processed := rateLimitEnv(0, 3) + + if processed != 3 { + t.Error("It should forward only 3 requests with rate-limiting", processed) + } } diff --git a/listener/listener.go b/listener/listener.go index 2733563..221ea14 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -13,6 +13,7 @@ import ( "net/http" "os" "strconv" + "time" ) // Enable debug logging only if "--verbose" flag passed @@ -24,7 +25,7 @@ func Debug(v ...interface{}) { func ReplayServer() (conn net.Conn, err error) { // Connection to reaplay server - conn, err = net.Dial("tcp", Settings.ReplayAddress) + conn, err = net.Dial("tcp", Settings.ReplayServer()) if err != nil { log.Println("Connection error ", err, Settings.ReplayAddress) @@ -42,15 +43,31 @@ func Run() { } fmt.Println("Listening for HTTP traffic on", Settings.Address+":"+strconv.Itoa(Settings.Port)) - fmt.Println("Forwarding requests to replay server:", Settings.ReplayAddress) + fmt.Println("Forwarding requests to replay server:", Settings.ReplayServer(), "Limit:", Settings.ReplayLimit) // Sniffing traffic from given address listener := RAWTCPListen(Settings.Address, Settings.Port) + currentTime := time.Now().UnixNano() + currentRPS := 0 + for { // Receiving TCPMessage object m := listener.Receive() + if Settings.ReplayLimit != 0 { + if (time.Now().UnixNano() - currentTime) > time.Second.Nanoseconds() { + currentTime = time.Now().UnixNano() + currentRPS = 0 + } + + if currentRPS >= Settings.ReplayLimit { + break + } + + currentRPS++ + } + go sendMessage(m) } } diff --git a/listener/settings.go b/listener/settings.go index 840dbfd..4d6cd1c 100644 --- a/listener/settings.go +++ b/listener/settings.go @@ -3,6 +3,8 @@ package listener import ( "flag" "os" + "strconv" + "strings" ) const ( @@ -18,11 +20,23 @@ type ListenerSettings struct { ReplayAddress string + ReplayLimit int + Verbose bool } var Settings ListenerSettings = ListenerSettings{} +func (s *ListenerSettings) ReplayServer() string { + host_info := strings.Split(s.ReplayAddress, "|") + + if len(host_info) > 1 { + s.ReplayLimit, _ = strconv.Atoi(host_info[1]) + } + + return host_info[0] +} + func init() { if len(os.Args) < 2 || os.Args[1] != "listen" { return