Merge pull request #25 from slawosz/save-request-to-reply-later

Save request to file for later replay

.gitignore

@@ -0,0 +1 @@
+*.swp

README.md

@@ -2,17 +2,17 @@
 
 ## About
 
-Gor is a simple http traffic replication tool written in Go. 
+Gor is a simple http traffic replication tool written in Go.
 Its main goal is to replay traffic from production servers to staging and dev environments.
 
 
-Now you can test your code on real user sessions in an automated and repeatable fashion.  
+Now you can test your code on real user sessions in an automated and repeatable fashion.
 **No more falling down in production!**
 
 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. 
+and sends it to the replay server or saves to file.
 The replay server forwards traffic to a given address.
 
 
@@ -23,9 +23,9 @@ The replay server forwards traffic to a given address.
 
 ```bash
 # Run on servers where you want to catch traffic. You can run it on each `web` machine.
-sudo gor listen -p 80 -r replay.server.local:28020 
+sudo gor listen -p 80 -r replay.server.local:28020
 
-# Replay server (replay.server.local). 
+# Replay server (replay.server.local).
 gor replay -f http://staging.server -p 28020
 ```
 
@@ -54,6 +54,16 @@ You can forward traffic to multiple endpoints. Just separate the addresses by co
 ```
 gor replay -f "http://staging.server|10,http://dev.server|5"
 ```
+### Saving requests to file
+You can save request to save to file to replay multiple time, or in different network:
+```
+gor listen -p 8080 -file requests.gor
+```
+
+And replaying:
+```
+gor replay -f "http://staging.server" -file requests.gor
+```
 
 ## Additional help
 ```
@@ -80,7 +90,7 @@ https://github.com/buger/gor/releases
 
 ## Building from source
 1. Setup standard Go environment http://golang.org/doc/code.html and ensure that $GOPATH environment variable properly set.
-2. `go get github.com/buger/gor`. 
+2. `go get github.com/buger/gor`.
 3. `cd $GOPATH/src/github.com/buger/gor`
 4. `go build gor.go` to get binary, or `go run gor.go` to build and run (useful for development)
 

integration_test.go

@@ -28,6 +28,7 @@ type Env struct {
 
 	ReplayLimit   int
 	ListenerLimit int
+	ForwardPort   int
 }
 
 func (e *Env) start() (p int) {
@@ -196,3 +197,106 @@ func TestListenerRateLimit(t *testing.T) {
 		t.Error("It should forward only 3 requests with rate-limiting", processed)
 	}
 }
+
+func (e *Env) startFileListener() (p int) {
+	p = 50000 + envs*10
+
+	e.ForwardPort = p + 2
+	go e.startHTTP(p, http.HandlerFunc(e.ListenHandler))
+	go e.startHTTP(p+2, http.HandlerFunc(e.ReplayHandler))
+	go e.startFileUsingListener(p, p+1)
+
+	// Time to start http and gor instances
+	time.Sleep(time.Millisecond * 100)
+
+	envs++
+
+	return
+}
+
+func (e *Env) startFileUsingListener(port int, replayPort int) {
+	listener.Settings.Verbose = e.Verbose
+	listener.Settings.Address = "127.0.0.1"
+	listener.Settings.FileToReplyPath = "integration_request.gor"
+	listener.Settings.Port = port
+
+	if e.ListenerLimit != 0 {
+		listener.Settings.ReplayAddress += "|" + strconv.Itoa(e.ListenerLimit)
+	}
+
+	listener.Run()
+}
+
+func (e *Env) startFileUsingReplay() {
+	replay.Settings.Verbose = e.Verbose
+	replay.Settings.FileToReplyPath = "integration_request.gor"
+	replay.Settings.ForwardAddress = "127.0.0.1:" + strconv.Itoa(e.ForwardPort)
+
+	if e.ReplayLimit != 0 {
+		replay.Settings.ForwardAddress += "|" + strconv.Itoa(e.ReplayLimit)
+	}
+
+	replay.Run()
+}
+
+func TestSavingRequestToFileAndReplyThem(t *testing.T) {
+	var request *http.Request
+	processed := make(chan int)
+
+	listenHandler := func(w http.ResponseWriter, r *http.Request) {
+		http.Error(w, "OK", http.StatusNotFound)
+	}
+
+	requestsCount := 0
+	var replayedRequests []*http.Request
+	replayHandler := func(w http.ResponseWriter, r *http.Request) {
+		requestsCount++
+
+		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)
+
+		replayedRequests = append(replayedRequests, r)
+		if t.Failed() {
+			fmt.Println("\nReplayed:", r, "\nOriginal:", request)
+		}
+
+		if requestsCount > 1 {
+			processed <- 1
+		}
+	}
+
+	env := &Env{
+		Verbose:       true,
+		ListenHandler: listenHandler,
+		ReplayHandler: replayHandler,
+	}
+
+	p := env.startFileListener()
+
+	request = getRequest(p)
+
+	for i := 0; i < 2; i++ {
+		go func() {
+			_, err := http.DefaultClient.Do(request)
+
+			if err != nil {
+				t.Error("Can't make request", err)
+			}
+		}()
+	}
+
+	// TODO: wait until gor will process response, should be kind of flag/semaphore
+	time.Sleep(time.Millisecond * 700)
+	go env.startFileUsingReplay()
+
+	select {
+	case <-processed:
+	case <-time.After(2 * time.Second):
+		for _, value := range replayedRequests {
+			fmt.Println(value)
+		}
+		t.Error("Timeout error")
+	}
+}

listener/listener.go

@@ -44,7 +44,26 @@ func Run() {
 	}
 
 	fmt.Println("Listening for HTTP traffic on", Settings.Address+":"+strconv.Itoa(Settings.Port))
-	fmt.Println("Forwarding requests to replay server:", Settings.ReplayAddress, "Limit:", Settings.ReplayLimit)
+
+	var messageLogger *log.Logger
+
+	if Settings.FileToReplyPath != "" {
+
+    file, err := os.OpenFile(Settings.FileToReplyPath, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0660)
+    defer file.Close()
+
+    if err != nil {
+      log.Fatal("Cannot open file %q. Error: %s", Settings.FileToReplyPath, err)
+    }
+
+    messageLogger = log.New(file, "", 0)
+	}
+
+  if messageLogger == nil {
+    fmt.Println("Forwarding requests to replay server:", Settings.ReplayAddress, "Limit:", Settings.ReplayLimit)
+  } else {
+    fmt.Println("Saving requests to file", Settings.FileToReplyPath)
+  }
 
 	// Sniffing traffic from given address
 	listener := RAWTCPListen(Settings.Address, Settings.Port)
@@ -52,6 +71,7 @@ func Run() {
 	currentTime := time.Now().UnixNano()
 	currentRPS := 0
 
+
 	for {
 		// Receiving TCPMessage object
 		m := listener.Receive()
@@ -69,7 +89,20 @@ func Run() {
 			currentRPS++
 		}
 
-		go sendMessage(m)
+		if messageLogger != nil {
+			go func() {
+				messageBuffer := new(bytes.Buffer)
+				messageWriter := bufio.NewWriter(messageBuffer)
+
+				fmt.Fprintf(messageWriter, "%v\n", time.Now().UnixNano())
+				fmt.Fprintf(messageWriter, "%s", string(m.Bytes()))
+
+				messageWriter.Flush()
+				messageLogger.Println(messageBuffer.String())
+			}()
+    } else {
+      go sendMessage(m)
+		}
 	}
 }
 

listener/raw_tcp_listener.go

@@ -26,19 +26,19 @@ type RAWTCPListener struct {
 }
 
 // RAWTCPListen creates a listener to capture traffic from RAW_SOCKET
-func RAWTCPListen(addr string, port int) (listener *RAWTCPListener) {
-	listener = &RAWTCPListener{}
+func RAWTCPListen(addr string, port int) (rawListener *RAWTCPListener) {
+	rawListener = &RAWTCPListener{}
 
-	listener.c_packets = make(chan *TCPPacket, 100)
-	listener.c_messages = make(chan *TCPMessage, 100)
-	listener.c_del_message = make(chan *TCPMessage, 100)
-	listener.messages = make(map[uint32]*TCPMessage)
+	rawListener.c_packets = make(chan *TCPPacket, 100)
+	rawListener.c_messages = make(chan *TCPMessage, 100)
+	rawListener.c_del_message = make(chan *TCPMessage, 100)
+	rawListener.messages = make(map[uint32]*TCPMessage)
 
-	listener.addr = addr
-	listener.port = port
+	rawListener.addr = addr
+	rawListener.port = port
 
-	go listener.listen()
-	go listener.readRAWSocket()
+	go rawListener.listen()
+	go rawListener.readRAWSocket()
 
 	return
 }

listener/settings.go

@@ -19,7 +19,8 @@ type ListenerSettings struct {
 	Port    int
 	Address string
 
-	ReplayAddress string
+	ReplayAddress   string
+	FileToReplyPath string
 
 	ReplayLimit int
 
@@ -50,5 +51,7 @@ func init() {
 	replayAddress := flag.String("r", defaultReplayAddress, "Address of replay server.")
 	Settings.ReplayServer(*replayAddress)
 
+	flag.StringVar(&Settings.FileToReplyPath, "file", "", "File to store captured requests")
+
 	flag.BoolVar(&Settings.Verbose, "verbose", false, "Log requests")
 }

replay/replay.go

@@ -25,64 +25,136 @@
 package replay
 
 import (
-	"bufio"
-	"bytes"
-	"io"
-	"log"
-	"net"
-	"net/http"
+  "bufio"
+  "bytes"
+  "io"
+  "log"
+  "net"
+  "net/http"
+  "time"
 )
 
+
 const bufSize = 4096
 
+type ReplayManager struct {
+  reqFactory *RequestFactory
+}
+
+func NewReplayManager() (rm *ReplayManager) {
+  rm = &ReplayManager{}
+  rm.reqFactory = NewRequestFactory()
+
+  return
+}
+
 // Debug enables logging only if "--verbose" flag passed
 func Debug(v ...interface{}) {
-	if Settings.Verbose {
-		log.Println(v...)
-	}
+  if Settings.Verbose {
+    log.Println(v...)
+  }
 }
 
 // ParseRequest in []byte returns a http request or an error
 func ParseRequest(data []byte) (request *http.Request, err error) {
-	buf := bytes.NewBuffer(data)
-	reader := bufio.NewReader(buf)
+  buf := bytes.NewBuffer(data)
+  reader := bufio.NewReader(buf)
 
-	request, err = http.ReadRequest(reader)
-	return
+  request, err = http.ReadRequest(reader)
+
+  if err != nil {
+    log.Fatal("Can not parse request", string(data), err)
+  }
+
+  return
 }
 
 // Run acts as `main` function of replay
 // Replay server listen to UDP traffic from Listeners
 // Each request processed by RequestFactory
 func Run() {
-	listener, err := net.Listen("tcp", Settings.Address)
+  rm := NewReplayManager()
 
-	log.Println("Starting replay server at:", Settings.Address)
+  if Settings.FileToReplyPath != "" {
+    rm.RunReplayFromFile()
+  } else {
+    rm.RunReplayFromNetwork()
+  }
+}
 
-	if err != nil {
-		log.Fatal("Can't start:", err)
-	}
+func (self *ReplayManager) RunReplayFromFile() {
+  TotalResponsesCount = 0
 
-	for _, host := range Settings.ForwardedHosts() {
-		log.Println("Forwarding requests to:", host.Url, "limit:", host.Limit)
-	}
+  log.Println("Starting file reply")
+  requests, err := parseReplyFile()
 
-	requestFactory := NewRequestFactory()
+  if err != nil {
+    log.Fatal("Can't parse request: ", err)
+  }
 
-	for {
-		conn, err := listener.Accept()
+  var lastTimestamp int64
 
-		if err != nil {
-			log.Println("Error while Accept()", err)
-			continue
-		}
+  if len(requests) > 0 {
+    lastTimestamp = requests[0].Timestamp
+  }
 
-		go handleConnection(conn, requestFactory)
-	}
+  requestsToReplay := 0
+
+	hosts := Settings.ForwardedHosts()
+  for _, host := range hosts {
+    if host.Limit > 0 {
+      requestsToReplay += host.Limit
+    } else {
+      requestsToReplay += len(requests)
+    }
+  }
+
+  for _, request := range requests {
+
+    parsedReq, err := ParseRequest(request.Request)
+
+    if err != nil {
+      log.Fatal("Can't parse request...:", err)
+    }
+
+    time.Sleep(time.Duration(request.Timestamp - lastTimestamp))
+
+    self.sendRequestToReplay(parsedReq)
+    lastTimestamp = request.Timestamp
+  }
+
+  for requestsToReplay > TotalResponsesCount {
+    time.Sleep(time.Second)
+  }
 
 }
 
-func handleConnection(conn net.Conn, rf *RequestFactory) error {
+func (self *ReplayManager) RunReplayFromNetwork() {
+  listener, err := net.Listen("tcp", Settings.Address)
+
+  log.Println("Starting replay server at:", Settings.Address)
+
+  if err != nil {
+    log.Fatal("Can't start:", err)
+  }
+
+  for _, host := range Settings.ForwardedHosts() {
+    log.Println("Forwarding requests to:", host.Url, "limit:", host.Limit)
+  }
+
+  for {
+    conn, err := listener.Accept()
+
+    if err != nil {
+      log.Println("Error while Accept()", err)
+      continue
+    }
+
+    go self.handleConnection(conn)
+  }
+}
+
+func (self *ReplayManager) handleConnection(conn net.Conn) error {
 	defer conn.Close()
 
 	var read = true
@@ -113,9 +185,13 @@ func handleConnection(conn net.Conn, rf *RequestFactory) error {
 		} else {
 			Debug("Adding request", request)
 
-			rf.Add(request)
+      self.sendRequestToReplay(request)
 		}
 	}()
 
 	return nil
 }
+
+func (self *ReplayManager) sendRequestToReplay(req *http.Request) {
+  self.reqFactory.Add(req)
+}

replay/replay_file_parser.go

@@ -0,0 +1,78 @@
+package replay
+
+import (
+  "bufio"
+  "log"
+  "os"
+  "bytes"
+  "strconv"
+
+  "fmt"
+)
+
+type ParsedRequest struct {
+  Request []byte
+  Timestamp int64
+}
+
+func (self ParsedRequest) String() string {
+  return fmt.Sprintf("Request: %v, timestamp: %v", string(self.Request), self.Timestamp)
+}
+
+func parseReplyFile() (requests []ParsedRequest, err error) {
+  requests, err = readLines(Settings.FileToReplyPath)
+
+  if err != nil {
+    log.Fatalf("readLines: %s", err)
+  }
+
+  return
+}
+
+// readLines reads a whole file into memory
+// and returns a slice of its lines.
+func readLines(path string) (requests []ParsedRequest, err error) {
+  file, err := os.Open(path)
+
+  if err != nil {
+    return nil, err
+  }
+  defer file.Close()
+
+  scanner := bufio.NewScanner(file)
+  scanner.Split(scanLinesFunc)
+
+  for scanner.Scan() {
+    if len(scanner.Text()) > 5 {
+      i := bytes.IndexByte(scanner.Bytes(), '\n')
+      timestamp, _ := strconv.Atoi(string(scanner.Bytes()[:i]))
+      pr := ParsedRequest{scanner.Bytes()[i + 1:], int64(timestamp)}
+
+      requests = append(requests, pr)
+    }
+  }
+
+  return requests, scanner.Err()
+}
+
+// scanner spliting logic
+func scanLinesFunc(data []byte, atEOF bool) (advance int, token []byte, err error) {
+	if atEOF && len(data) == 0 {
+		return 0, nil, nil
+	}
+
+  delimiter := []byte{'\r', '\n', '\r', '\n', '\n'}
+
+  // We have a http request end: \r\n\r\n
+	if i := bytes.Index(data, delimiter); i >= 0 {
+		return (i + len(delimiter)), data[0:(i + len(delimiter))], nil
+	}
+
+	// If we're at EOF, we have a final, non-terminated line. Return it.
+	if atEOF {
+		return len(data), data, nil
+	}
+
+	// Request more data.
+	return 0, nil, nil
+}

replay/request_factory.go

@@ -60,8 +60,6 @@ func (f *RequestFactory) sendRequest(host *ForwardHost, request *http.Request) {
 	request.RequestURI = ""
 	request.URL, _ = url.ParseRequestURI(URL)
 
-	Debug("Sending request:", host.Url, request)
-
 	resp, err := client.Do(request)
 
 	if err == nil {
@@ -96,6 +94,7 @@ func (f *RequestFactory) handleRequests() {
 			resp.host.Stat.IncResp(resp)
 		}
 	}
+
 }
 
 // Add request to channel for further processing

replay/request_stats.go

@@ -4,6 +4,8 @@ import (
 	"time"
 )
 
+var TotalResponsesCount int
+
 // RequestStat stores in context of current timestamp
 type RequestStat struct {
 	timestamp int64
@@ -31,6 +33,7 @@ func (s *RequestStat) IncReq() {
 // IncResp is called after response
 func (s *RequestStat) IncResp(resp *HttpResponse) {
 	s.Touch()
+  TotalResponsesCount++
 
 	if resp.err != nil {
 		s.Errors++

replay/settings.go

@@ -21,7 +21,10 @@ type ReplaySettings struct {
 	Host string
 
 	Address        string
-	ForwardAddress string
+
+	ForwardAddress  string
+
+	FileToReplyPath string
 
 	Verbose bool
 }
@@ -79,5 +82,7 @@ func init() {
 	Settings.SetAddress()
 	flag.StringVar(&Settings.ForwardAddress, "f", defaultForwardAddress, "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.FileToReplyPath, "file", "", "File to replay captured requests from")
+
 	flag.BoolVar(&Settings.Verbose, "verbose", false, "Log requests")
 }