Proper conn closing and settings refactoring

This commit is contained in:
Leonid Bugaev
2013-06-22 18:16:23 +06:00
parent 5910409379
commit 76e5981f49
6 changed files with 59 additions and 72 deletions
+36 -32
View File
@@ -12,31 +12,25 @@ import (
"net"
"net/http"
"os"
"time"
"strconv"
)
// Enable debug logging only if "--verbose" flag passed
func Debug(v ...interface{}) {
if Settings.verbose {
if Settings.Verbose {
log.Println(v...)
}
}
func ReplayServer() net.Conn {
func ReplayServer() (conn net.Conn, err error) {
// Connection to reaplay server
conn, err := net.Dial("tcp", Settings.ReplayServer())
conn, err = net.Dial("tcp", Settings.ReplayAddress)
if err != nil {
log.Println("Connection error ", err, Settings.ReplayServer())
log.Println("Reconnecting to replay server in 10 seconds")
time.Sleep(10 * time.Second)
return ReplayServer()
log.Println("Connection error ", err, Settings.ReplayAddress)
}
log.Println("Connected to replay server:", Settings.ReplayServer())
return conn
return
}
// Because its sub-program, Run acts as `main`
@@ -47,39 +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())
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()
conn := ReplayServer()
// For debugging purpose
// Usually request parsing happens in replay part
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)
}
}
if err != nil {
log.Println("Failed to send message. Replay server not respond.")
return
} else {
defer conn.Close()
}
// For debugging purpose
// Usually request parsing happens in replay part
if Settings.Verbose {
buf := bytes.NewBuffer(m.Bytes())
reader := bufio.NewReader(buf)
_, err := conn.Write(m.Bytes())
request, err := http.ReadRequest(reader)
if err != nil {
log.Println("Error while sending requests", err)
Debug("Error while parsing request:", err, string(m.Bytes()))
} else {
request.ParseMultipartForm(32 << 20)
Debug("Forwarding request:", request)
}
}
_, err = conn.Write(m.Bytes())
conn.Close()
if err != nil {
log.Println("Error while sending requests", err)
}
}
+1 -1
View File
@@ -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:
+8 -22
View File
@@ -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")
}
+1 -1
View File
@@ -37,7 +37,7 @@ const bufSize = 4096
// Enable debug logging only if "--verbose" flag passed
func Debug(v ...interface{}) {
if Settings.verbose {
if Settings.Verbose {
log.Println(v...)
}
}
+1 -3
View File
@@ -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)
+12 -13
View File
@@ -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")
}