Merge pull request #10 from buger/use_tcp

Used TCP instead of UDP

README.md

@@ -10,11 +10,11 @@ Now you can test your code on real user sessions in an automated and repeatable
 Gor consists of 2 parts: listener and replay servers.
 
 The listener server catches http traffic from a given port in real-time
-and sends it to the replay server via UDP. 
+and sends it to the replay server. 
 The replay server forwards traffic to a given address.
 
 
-![Diagram](http://i.imgur.com/zZCFPCY.png)
+![Diagram](http://i.imgur.com/9mqj2SK.png)
 
 
 ## Basic example

listener/listener.go

@@ -12,15 +12,27 @@ import (
 	"net"
 	"net/http"
 	"os"
+	"strconv"
 )
 
 // Enable debug logging only if "--verbose" flag passed
 func Debug(v ...interface{}) {
-	if Settings.verbose {
+	if Settings.Verbose {
 		log.Println(v...)
 	}
 }
 
+func ReplayServer() (conn net.Conn, err error) {
+	// Connection to reaplay server
+	conn, err = net.Dial("tcp", Settings.ReplayAddress)
+
+	if err != nil {
+		log.Println("Connection error ", err, Settings.ReplayAddress)
+	}
+
+	return
+}
+
 // Because its sub-program, Run acts as `main`
 func Run() {
 	if os.Getuid() != 0 {
@@ -29,40 +41,49 @@ func Run() {
 		os.Exit(1)
 	}
 
-	fmt.Println("Listening for HTTP traffic on", Settings.Address())
-	fmt.Println("Forwarding requests to replay server:", Settings.ReplayServer())
-
-	// Connection to reaplay server
-	serverAddr, err := net.ResolveUDPAddr("udp4", Settings.ReplayServer())
-	conn, err := net.DialUDP("udp", nil, serverAddr)
-
-	if err != nil {
-		log.Fatal("Connection error", err)
-	}
+	fmt.Println("Listening for HTTP traffic on", Settings.Address+":"+strconv.Itoa(Settings.Port))
+	fmt.Println("Forwarding requests to replay server:", Settings.ReplayAddress)
 
 	// Sniffing traffic from given address
-	listener := RAWTCPListen(Settings.address, Settings.port)
+	listener := RAWTCPListen(Settings.Address, Settings.Port)
 
 	for {
 		// Receiving TCPMessage object
 		m := listener.Receive()
 
-		if Settings.verbose {
-			buf := bytes.NewBuffer(m.Bytes())
-			reader := bufio.NewReader(buf)
+		go sendMessage(m)
+	}
+}
 
-			request, err := http.ReadRequest(reader)
+func sendMessage(m *TCPMessage) {
+	conn, err := ReplayServer()
 
-			if err != nil {
-				Debug("Error while parsing request:", err, string(m.Bytes()))
-			} else {
-				request.ParseMultipartForm(32 << 20)
-				Debug("Forwarding request:", request)
-			}
-		}
-
-		conn.Write(m.Bytes())
+	if err != nil {
+		log.Println("Failed to send message. Replay server not respond.")
+		return
+	} else {
+		defer conn.Close()
 	}
 
-	conn.Close()
+	// For debugging purpose
+	// Usually request parsing happens in replay part
+	if Settings.Verbose {
+		buf := bytes.NewBuffer(m.Bytes())
+		reader := bufio.NewReader(buf)
+
+		request, err := http.ReadRequest(reader)
+
+		if err != nil {
+			Debug("Error while parsing request:", err, string(m.Bytes()))
+		} else {
+			request.ParseMultipartForm(32 << 20)
+			Debug("Forwarding request:", request)
+		}
+	}
+
+	_, err = conn.Write(m.Bytes())
+
+	if err != nil {
+		log.Println("Error while sending requests", err)
+	}
 }

listener/raw_tcp_listener.go

@@ -46,8 +46,8 @@ func (t *RAWTCPListener) listen() {
 		select {
 		// If message ready for deletion it means that its also complete or expired by timeout
 		case message := <-t.c_del_message:
-			t.deleteMessage(message)
 			t.c_messages <- message
+			t.deleteMessage(message)
 
 		// We need to use channgels to process each packet to avoid data races
 		case packet := <-t.c_packets:
@@ -120,7 +120,6 @@ func (t *RAWTCPListener) readRAWSocket() {
 
 					// To avoid socket locking processing packet in new goroutine
 					go func(buf []byte) {
-						log.Println("Received packet", string(new_buf))
 						packet := NewTCPPacket(new_buf)
 						t.c_packets <- packet
 					}(new_buf)
@@ -147,7 +146,6 @@ func (t *RAWTCPListener) processTCPPacket(packet *TCPPacket) {
 	if message == nil {
 		// We sending c_del_message channel, so message object can communicate with Listener and notify it if message completed
 		message = NewTCPMessage(packet.Ack, t.c_del_message)
-		Debug("Adding message")
 
 		t.messages = append(t.messages, message)
 	}

listener/settings.go

@@ -3,8 +3,6 @@ package listener
 import (
 	"flag"
 	"os"
-	"strconv"
-	"strings"
 )
 
 const (
@@ -15,38 +13,26 @@ const (
 )
 
 type ListenerSettings struct {
-	port    int
-	address string
+	Port    int
+	Address string
 
-	replayAddress string
+	ReplayAddress string
 
-	verbose bool
+	Verbose bool
 }
 
 var Settings ListenerSettings = ListenerSettings{}
 
-func (s *ListenerSettings) ReplayServer() string {
-	if !strings.Contains(s.replayAddress, ":") {
-		return s.replayAddress + ":28020"
-	}
-
-	return s.replayAddress
-}
-
-func (s *ListenerSettings) Address() string {
-	return s.address + ":" + strconv.Itoa(s.port)
-}
-
 func init() {
 	if len(os.Args) < 2 || os.Args[1] != "listen" {
 		return
 	}
 
-	flag.IntVar(&Settings.port, "p", defaultPort, "Specify the http server port whose traffic you want to capture")
+	flag.IntVar(&Settings.Port, "p", defaultPort, "Specify the http server port whose traffic you want to capture")
 
-	flag.StringVar(&Settings.address, "ip", defaultAddress, "Specifi IP address to listen")
+	flag.StringVar(&Settings.Address, "ip", defaultAddress, "Specifi IP address to listen")
 
-	flag.StringVar(&Settings.replayAddress, "r", defaultReplayAddress, "Address of replay server.")
+	flag.StringVar(&Settings.ReplayAddress, "r", defaultReplayAddress, "Address of replay server.")
 
-	flag.BoolVar(&Settings.verbose, "verbose", false, "Log requests")
+	flag.BoolVar(&Settings.Verbose, "verbose", false, "Log requests")
 }

replay/replay.go

@@ -27,16 +27,17 @@ package replay
 import (
 	"bufio"
 	"bytes"
+	"io"
 	"log"
 	"net"
 	"net/http"
 )
 
-const bufSize = 1024 * 10
+const bufSize = 4096
 
 // Enable debug logging only if "--verbose" flag passed
 func Debug(v ...interface{}) {
-	if Settings.verbose {
+	if Settings.Verbose {
 		log.Println(v...)
 	}
 }
@@ -53,22 +54,14 @@ func ParseRequest(data []byte) (request *http.Request, err error) {
 // Replay server listen to UDP traffic from Listeners
 // Each request processed by RequestFactory
 func Run() {
-	var buf [bufSize]byte
+	listener, err := net.Listen("tcp", Settings.Address())
 
-	addr, err := net.ResolveUDPAddr("udp", Settings.Address())
-	if err != nil {
-		log.Fatal("Can't start:", err)
-	}
-
-	conn, err := net.ListenUDP("udp", addr)
 	log.Println("Starting replay server at:", Settings.Address())
 
 	if err != nil {
 		log.Fatal("Can't start:", err)
 	}
 
-	defer conn.Close()
-
 	for _, host := range Settings.ForwardedHosts() {
 		log.Println("Forwarding requests to:", host.Url, "limit:", host.Limit)
 	}
@@ -76,23 +69,53 @@ func Run() {
 	requestFactory := NewRequestFactory()
 
 	for {
-		n, _, err := conn.ReadFromUDP(buf[0:])
+		conn, err := listener.Accept()
 
 		if err != nil {
+			log.Println("Error while Accept()", err)
 			continue
 		}
 
-		if n > 0 {
-			if n > bufSize {
-				Debug("Too large udp packet", bufSize)
-			}
-
-			if request, err := ParseRequest(buf[0:n]); err != nil {
-				Debug("Error while parsing request", err, buf[0:n])
-			} else {
-				requestFactory.Add(request)
-			}
-		}
+		go handleConnection(conn, requestFactory)
 	}
 
 }
+
+func handleConnection(conn net.Conn, rf *RequestFactory) error {
+	defer conn.Close()
+
+	var read = true
+	var response []byte
+	var buf []byte
+
+	buf = make([]byte, bufSize)
+
+	for read {
+		log.Println("Start reading")
+		n, err := conn.Read(buf)
+		log.Println("Reading", n, err)
+
+		switch err {
+		case io.EOF:
+			read = false
+		case nil:
+			response = append(response, buf[0:n]...)
+			if n < bufSize {
+				read = false
+			}
+		default:
+			read = false
+		}
+	}
+
+	log.Println("Response", string(response))
+
+	if request, err := ParseRequest(response); err != nil {
+		Debug("Error while parsing request", err, response)
+	} else {
+		Debug("Adding request", request)
+		rf.Add(request)
+	}
+
+	return nil
+}

replay/request_factory.go

@@ -47,9 +47,7 @@ func (f *RequestFactory) sendRequest(host *ForwardHost, request *http.Request) {
 	request.RequestURI = ""
 	request.URL, _ = url.ParseRequestURI(URL)
 
-	if Settings.verbose {
-		Debug("Sending request:", host.Url, request)
-	}
+	Debug("Sending request:", host.Url, request)
 
 	resp, err := client.Do(request)
 

replay/settings.go

@@ -1,10 +1,10 @@
 package replay
 
 import (
+	"flag"
+	"os"
 	"strconv"
 	"strings"
-	"os"
-	"flag"
 )
 
 type ForwardHost struct {
@@ -15,12 +15,12 @@ type ForwardHost struct {
 }
 
 type ReplaySettings struct {
-	port int
-	host string
+	Port int
+	Host string
 
-	forwardAddress string
+	ForwardAddress string
 
-	verbose bool
+	Verbose bool
 }
 
 // ForwardedHosts implements forwardAddress syntax support for multiple hosts (coma separated), and rate limiting by specifing "|maxRps" after host name.
@@ -30,7 +30,7 @@ type ReplaySettings struct {
 func (r *ReplaySettings) ForwardedHosts() (hosts []*ForwardHost) {
 	hosts = make([]*ForwardHost, 0, 10)
 
-	for _, address := range strings.Split(r.forwardAddress, ",") {
+	for _, address := range strings.Split(r.ForwardAddress, ",") {
 		host_info := strings.Split(address, "|")
 
 		if strings.Index(host_info[0], "http") == -1 {
@@ -52,12 +52,11 @@ func (r *ReplaySettings) ForwardedHosts() (hosts []*ForwardHost) {
 
 // Helper to return address with port, e.g.: 127.0.0.1:28020
 func (r *ReplaySettings) Address() string {
-	return r.host + ":" + strconv.Itoa(r.port)
+	return r.Host + ":" + strconv.Itoa(r.Port)
 }
 
 var Settings ReplaySettings = ReplaySettings{}
 
-
 func init() {
 	if len(os.Args) < 2 || os.Args[1] != "replay" {
 		return
@@ -70,11 +69,11 @@ func init() {
 		defaultAddress = "http://localhost:8080"
 	)
 
-	flag.IntVar(&Settings.port, "p", defaultPort, "specify port number")
+	flag.IntVar(&Settings.Port, "p", defaultPort, "specify port number")
 
-	flag.StringVar(&Settings.host, "ip", defaultHost, "ip addresses to listen on")
+	flag.StringVar(&Settings.Host, "ip", defaultHost, "ip addresses to listen on")
 
-	flag.StringVar(&Settings.forwardAddress, "f", defaultAddress, "http address to forward traffic.\n\tYou can limit requests per second by adding `|num` after address.\n\tIf you have multiple addresses with different limits. For example: http://staging.example.com|100,http://dev.example.com|10")
+	flag.StringVar(&Settings.ForwardAddress, "f", defaultAddress, "http address to forward traffic.\n\tYou can limit requests per second by adding `|num` after address.\n\tIf you have multiple addresses with different limits. For example: http://staging.example.com|100,http://dev.example.com|10")
 
-	flag.BoolVar(&Settings.verbose, "verbose", false, "Log requests")
+	flag.BoolVar(&Settings.Verbose, "verbose", false, "Log requests")
 }