diff --git a/listener/listener.go b/listener/listener.go index 14bb951..9965aa9 100644 --- a/listener/listener.go +++ b/listener/listener.go @@ -7,6 +7,7 @@ package listener import ( "bufio" "bytes" + "encoding/gob" "fmt" "log" "net" @@ -16,6 +17,11 @@ import ( "time" ) +type ParsedRequest struct { + Timestamp int64 + Request []byte +} + // Debug enables logging only if "--verbose" flag passed func Debug(v ...interface{}) { if Settings.Verbose { @@ -49,7 +55,7 @@ func Run() { fmt.Println("Listening for HTTP traffic on", Settings.Address+":"+strconv.Itoa(Settings.Port)) - var messageLogger *log.Logger + var fileEnc *gob.Encoder if Settings.FileToReplayPath != "" { @@ -60,13 +66,10 @@ func Run() { log.Fatal("Cannot open file %q. Error: %s", Settings.FileToReplayPath, err) } - messageLogger = log.New(file, "", 0) - } - - if messageLogger == nil { - fmt.Println("Forwarding requests to replay server:", Settings.ReplayAddress, "Limit:", Settings.ReplayLimit) - } else { + fileEnc = gob.NewEncoder(file) fmt.Println("Saving requests to file", Settings.FileToReplayPath) + } else { + fmt.Println("Forwarding requests to replay server:", Settings.ReplayAddress, "Limit:", Settings.ReplayLimit) } // Sniffing traffic from given address @@ -92,16 +95,10 @@ func Run() { currentRPS++ } - if messageLogger != nil { + if Settings.FileToReplayPath != "" { 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()) + message := ParsedRequest{time.Now().UnixNano(), m.Bytes()} + fileEnc.Encode(message) }() } else { go sendMessage(m) diff --git a/replay/replay_file_parser.go b/replay/replay_file_parser.go index d185400..216d5d2 100644 --- a/replay/replay_file_parser.go +++ b/replay/replay_file_parser.go @@ -1,79 +1,57 @@ package replay import ( - "bufio" - "log" - "os" - "bytes" - "strconv" + "bytes" + "encoding/gob" + "io" + "io/ioutil" + "log" - "fmt" + "fmt" ) type ParsedRequest struct { - Request []byte - Timestamp int64 + Timestamp int64 + Request []byte } func (self ParsedRequest) String() string { - return fmt.Sprintf("Request: %v, timestamp: %v", string(self.Request), self.Timestamp) + return fmt.Sprintf("Request: %v, timestamp: %v", string(self.Request), self.Timestamp) } func parseReplayFile() (requests []ParsedRequest, err error) { - requests, err = readLines(Settings.FileToReplayPath) + requests, err = readLines(Settings.FileToReplayPath) - if err != nil { - log.Fatalf("readLines: %s", err) - } + if err != nil { + log.Fatalf("readLines: %s", err) + } - return + return } // readLines reads a whole file into memory -// and returns a slice of its lines. +// and returns a slice of request+timestamps. func readLines(path string) (requests []ParsedRequest, err error) { - file, err := os.Open(path) + file, err := ioutil.ReadFile(path) - if err != nil { - return nil, err - } - defer file.Close() + if err != nil { + return nil, err + } - scanner := bufio.NewScanner(file) - scanner.Split(scanLinesFunc) + fileBuf := bytes.NewBuffer(file) + fileDec := gob.NewDecoder(fileBuf) - for scanner.Scan() { - if len(scanner.Text()) > 5 { - buf := append([]byte(nil), scanner.Bytes()...) - i := bytes.IndexByte(buf, '\n') - timestamp, _ := strconv.Atoi(string(buf[:i])) - pr := ParsedRequest{buf[i + 1:], int64(timestamp)} + for err == nil { + var reqBuf ParsedRequest + err = fileDec.Decode(&reqBuf) - requests = append(requests, pr) - } - } + if err == io.EOF { + err = nil + break + } - 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 + requests = append(requests, reqBuf) } - // Request more data. - return 0, nil, nil + return requests, err }