From 674ffc620241cd0c2562ba7a917fa9ce322a463a Mon Sep 17 00:00:00 2001 From: Leonid Bugaev Date: Sun, 28 Jun 2015 00:16:19 +0500 Subject: [PATCH] Start adding tests for traffic modifier --- input_raw_test.go | 2 +- output_http_test.go | 10 ++-- plugins.go | 1 - traffic_modifier.go | 115 +++++++++++++++++++-------------------- traffic_modifier_test.go | 74 +++++++++++++++++++++++++ 5 files changed, 137 insertions(+), 65 deletions(-) create mode 100644 traffic_modifier_test.go diff --git a/input_raw_test.go b/input_raw_test.go index b1bcca6..2eeb6d9 100644 --- a/input_raw_test.go +++ b/input_raw_test.go @@ -13,7 +13,7 @@ func TestRAWInput(t *testing.T) { wg := new(sync.WaitGroup) quit := make(chan int) - listener := startHTTP(func(req *http.Request) {}) + listener := startHTTP(func(w http.ResponseWriter, req *http.Request) {}) input := NewRAWInput(listener.Addr().String()) output := NewTestOutput(func(data []byte) { diff --git a/output_http_test.go b/output_http_test.go index 2432d5c..b24f442 100644 --- a/output_http_test.go +++ b/output_http_test.go @@ -12,9 +12,9 @@ import ( "time" ) -func startHTTP(cb func(*http.Request)) net.Listener { +func startHTTP(cb func(http.ResponseWriter, *http.Request)) net.Listener { handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - cb(r) + cb(w, r) }) listener, _ := net.Listen("tcp", ":0") @@ -54,7 +54,7 @@ func TestHTTPOutput(t *testing.T) { headers := HTTPHeaders{HTTPHeader{"User-Agent", "Gor"}} methods := HTTPMethods{"GET", "PUT", "POST"} - listener := startHTTP(func(req *http.Request) { + listener := startHTTP(func(w http.ResponseWriter, req *http.Request) { if req.Header.Get("User-Agent") != "Gor" { t.Error("Wrong header") } @@ -104,7 +104,7 @@ func TestHTTPOutputChunkedEncoding(t *testing.T) { headers := HTTPHeaders{HTTPHeader{"User-Agent", "Gor"}} methods := HTTPMethods{"GET", "PUT", "POST"} - listener := startHTTP(func(req *http.Request) { + listener := startHTTP(func(w http.ResponseWriter, req *http.Request) { defer req.Body.Close() body, _ := ioutil.ReadAll(req.Body) @@ -140,7 +140,7 @@ func BenchmarkHTTPOutput(b *testing.B) { headers := HTTPHeaders{HTTPHeader{"User-Agent", "Gor"}} methods := HTTPMethods{"GET", "PUT", "POST"} - listener := startHTTP(func(req *http.Request) { + listener := startHTTP(func(w http.ResponseWriter, req *http.Request) { time.Sleep(50 * time.Millisecond) wg.Done() }) diff --git a/plugins.go b/plugins.go index ecafbfa..62bce40 100644 --- a/plugins.go +++ b/plugins.go @@ -62,7 +62,6 @@ func registerPlugin(constructor interface{}, options ...interface{}) { Plugins.Inputs = append(Plugins.Inputs, plugin_wrapper.(io.Reader)) } - if _, ok := plugin.(io.Writer); ok { Plugins.Outputs = append(Plugins.Outputs, plugin_wrapper.(io.Writer)) } diff --git a/traffic_modifier.go b/traffic_modifier.go index ddc92d1..7a34bb9 100644 --- a/traffic_modifier.go +++ b/traffic_modifier.go @@ -1,97 +1,96 @@ package main import ( - "fmt" - "log" - "io" - "os/exec" - "os" - "bufio" - "encoding/hex" + "bufio" + "encoding/hex" + "fmt" + "io" + "log" + "os" + "os/exec" ) type TrafficModifier struct { - plugin interface{} - command string + plugin interface{} + command string - data chan []byte + data chan []byte - Stdin io.Writer - Stdout io.Reader + Stdin io.Writer + Stdout io.Reader } func NewTrafficModifier(plugin interface{}, command string) io.Reader { - m := new(TrafficModifier) - m.plugin = plugin - m.command = command - m.data = make(chan []byte) + m := new(TrafficModifier) + m.plugin = plugin + m.command = command + m.data = make(chan []byte) - cmd := exec.Command(command) + cmd := exec.Command(command) - m.Stdout, _ = cmd.StdoutPipe() - m.Stdin, _ = cmd.StdinPipe() - cmd.Stderr = os.Stderr + m.Stdout, _ = cmd.StdoutPipe() + m.Stdin, _ = cmd.StdinPipe() + cmd.Stderr = os.Stderr - go m.copy(m.Stdin, m.plugin.(io.Reader)) - go m.read(m.Stdout) + go m.copy(m.Stdin, m.plugin.(io.Reader)) + go m.read(m.Stdout) - go func(){ - err := cmd.Start() + go func() { + err := cmd.Start() - if (err != nil) { - log.Fatal(err) - } - }() + if err != nil { + log.Fatal(err) + } + }() - defer cmd.Wait() + defer cmd.Wait() - return m + return m } func (m *TrafficModifier) copy(to io.Writer, from io.Reader) { - buf := make([]byte, 5*1024*1024) - dst := make([]byte, len(buf)*2) + buf := make([]byte, 5*1024*1024) + dst := make([]byte, len(buf)*2) - for { - nr, _ := from.Read(buf) - if nr > 0 && len(buf) > nr { - hex.Encode(dst, buf[0:nr]) - to.Write(dst[0:nr*2]) - to.Write([]byte("\r\n")) - } - } + for { + nr, _ := from.Read(buf) + if nr > 0 && len(buf) > nr { + hex.Encode(dst, buf[0:nr]) + to.Write(dst[0 : nr*2]) + to.Write([]byte("\r\n")) + } + } } func (m *TrafficModifier) read(from io.Reader) { - buf := make([]byte, 5*1024*1024) + buf := make([]byte, 5*1024*1024) - scanner := bufio.NewScanner(from) + scanner := bufio.NewScanner(from) - for scanner.Scan() { - bytes := scanner.Bytes() - hex.Decode(buf, bytes) + for scanner.Scan() { + bytes := scanner.Bytes() + hex.Decode(buf, bytes) - Debug("Received:", buf[0:len(bytes)/2]) + Debug("Received:", buf[0:len(bytes)/2]) - m.data <- buf[0:len(bytes)/2] - } + m.data <- buf[0 : len(bytes)/2] + } - if err := scanner.Err(); err != nil { - fmt.Fprintln(os.Stderr, "Traffic modifier command failed:", err) - } + if err := scanner.Err(); err != nil { + fmt.Fprintln(os.Stderr, "Traffic modifier command failed:", err) + } - return + return } func (m *TrafficModifier) Read(data []byte) (int, error) { - Debug("Trying to read channel!") - buf := <- m.data - copy(data, buf) + Debug("Trying to read channel!") + buf := <-m.data + copy(data, buf) - return len(buf), nil + return len(buf), nil } - func (m *TrafficModifier) String() string { - return fmt.Sprintf("Modifying traffic for %s using '%s' command", m.plugin, m.command) + return fmt.Sprintf("Modifying traffic for %s using '%s' command", m.plugin, m.command) } diff --git a/traffic_modifier_test.go b/traffic_modifier_test.go new file mode 100644 index 0000000..547876e --- /dev/null +++ b/traffic_modifier_test.go @@ -0,0 +1,74 @@ +package main + +import ( + _ "bufio" + "bytes" + "crypto/rand" + _ "io" + "io/ioutil" + _ "log" + _ "net" + "net/http" + "sync" + "testing" +) + +// Simple service that generate token on request, and require this token for accesing to secure area +func NewFakeSecureService(wg *sync.WaitGroup) string { + active_tokens := make([][]byte, 0) + + listener := startHTTP(func(w http.ResponseWriter, req *http.Request) { + switch req.URL.Path { + case "/token": + // Generate random token + token_length := 10 + token := make([]byte, token_length) + rand.Read(token) + + w.Write(token) + + active_tokens = append(active_tokens, token) + case "/secure": + token := []byte(req.URL.Query().Get("token")) + + for _, t := range active_tokens { + if bytes.Equal(t, token) { + w.WriteHeader(http.StatusAccepted) + } else { + w.WriteHeader(http.StatusForbidden) + } + } + } + + wg.Done() + }) + + return "http://" + listener.Addr().String() +} + +func TestFakeSecureService(t *testing.T) { + var resp *http.Response + + wg := new(sync.WaitGroup) + + addr := NewFakeSecureService(wg) + + wg.Add(3) + + resp, _ = http.Get(addr + "/token") + token, _ := ioutil.ReadAll(resp.Body) + + // Right token + resp, _ = http.Get(addr + "/secure?token=" + string(token)) + if resp.StatusCode != http.StatusAccepted { + t.Error("Valid token should returns wrong status:", resp.StatusCode) + } + + // Wrong tokens forbidden + resp, _ = http.Get(addr + "/secure?token=wrong") + if resp.StatusCode != http.StatusForbidden { + t.Error("Wrong tokens should be forbidden, instead:", resp.StatusCode) + } + + wg.Wait() +}