github/goreplay
4
0
Fork 0
mirror of https://github.com/buger/goreplay.git synced 2024-04-21 12:32:02 +00:00
Code Issues Releases Wiki Activity

Start adding tests for traffic modifier

Browse Source
This commit is contained in:
Leonid Bugaev
2015-06-28 00:16:19 +05:00
parent 4f03f81d04
commit 674ffc6202
5 changed files with 137 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
input_raw_test.go
+1 -1
View File
@@ -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) {
output_http_test.go
+5 -5
View File
@@ -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()
})
plugins.go
-1
View File
@@ -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))
}
traffic_modifier.go
+57 -58
View File
@@ -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)
}
traffic_modifier_test.go
+74
View File
@@ -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()
}