From 77751af316823972726c4749c17b2010e0e89081 Mon Sep 17 00:00:00 2001
From: Leonid Bugaev <leonsbox@gmail.com>
Date: Mon, 5 Aug 2013 20:42:16 +0400
Subject: [PATCH] Added listener limit option

---
 integration_test.go  | 25 ++++++++++++++++++++-----
 listener/listener.go | 23 +++++++++++++++++++++--
 listener/settings.go | 14 ++++++++++++++
 3 files changed, 55 insertions(+), 7 deletions(-)

diff --git a/integration_test.go b/integration_test.go
index 612fd12..c5c44e9 100644
--- a/integration_test.go
+++ b/integration_test.go
@@ -29,7 +29,8 @@ type Env struct {
 	ListenHandler http.HandlerFunc
 	ReplayHandler http.HandlerFunc
 
-	ReplayLimit int
+	ReplayLimit   int
+	ListenerLimit int
 }
 
 func (e *Env) start() (p int) {
@@ -53,6 +54,11 @@ func (e *Env) startListener(port int, replayPort int) {
 	listener.Settings.Address = "127.0.0.1"
 	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()
 }
 
@@ -127,7 +133,7 @@ func TestReplay(t *testing.T) {
 	}
 }
 
-func rateLimitEnv(limit int) int32 {
+func rateLimitEnv(replayLimit int, listenerLimit int) int32 {
 	var processed int32
 
 	listenHandler := func(w http.ResponseWriter, r *http.Request) {
@@ -142,7 +148,8 @@ func rateLimitEnv(limit int) int32 {
 	env := &Env{
 		ListenHandler: listenHandler,
 		ReplayHandler: replayHandler,
-		ReplayLimit:   limit,
+		ReplayLimit:   replayLimit,
+		ListenerLimit: listenerLimit,
 	}
 
 	p := env.start()
@@ -158,7 +165,7 @@ func rateLimitEnv(limit int) int32 {
 }
 
 func TestWithoutReplayRateLimit(t *testing.T) {
-	processed := rateLimitEnv(0)
+	processed := rateLimitEnv(0, 0)
 
 	if processed != 10 {
 		t.Error("It should forward all requests without rate-limiting", processed)
@@ -166,9 +173,17 @@ func TestWithoutReplayRateLimit(t *testing.T) {
 }
 
 func TestReplayRateLimit(t *testing.T) {
-	processed := rateLimitEnv(5)
+	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..d81b957 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,33 @@ 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++
+		}
+
+		fmt.Println(currentRPS, Settings.ReplayLimit)
+
 		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