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

First middleware tests! :dance:

Browse Source
This commit is contained in:
Leonid Bugaev
2015-07-17 20:58:31 +05:00
parent f812120dfc
commit 5e6f7e03a7
6 changed files with 120 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
emitter.go
+28 -3
View File
@@ -3,11 +3,36 @@ package main
import (
"io"
"time"
"crypto/rand"
)
func uuid() []byte {
b := make([]byte, 16)
rand.Read(b)
return b
}
func Start(stop chan int) {
for _, in := range Plugins.Inputs {
go CopyMulty(in, Plugins.Outputs...)
if Settings.middleware != "" {
middleware := NewMiddleware(Settings.middleware)
for _, in := range Plugins.Inputs {
middleware.ReadFrom(in)
}
// We going only to read responses, so using same ReadFrom method
for _, out := range Plugins.Outputs {
if r, ok := out.(io.Reader); ok {
middleware.ReadFrom(r)
}
}
go CopyMulty(middleware, Plugins.Outputs...)
} else {
for _, in := range Plugins.Inputs {
go CopyMulty(in, Plugins.Outputs...)
}
}
for {
@@ -31,7 +56,7 @@ func CopyMulty(src io.Reader, writers ...io.Writer) (err error) {
if nr > 0 && len(buf) > nr {
payload := buf[0:nr]
Debug("[EMITTER] Received payload:", string(payload))
Debug("[EMITTER] input:", string(payload))
if modifier != nil {
payload = modifier.Rewrite(payload)
examples/echo_modifier.rb
+3 -3
View File
@@ -10,7 +10,7 @@ while data = STDIN.gets
STDOUT.puts encoded
STDERR.puts "[DEBUG] Original data: #{data}"
STDERR.puts "[DEBUG] Decoded request: #{decoded}"
STDERR.puts "[DEBUG] Encoded data: #{encoded}"
STDERR.puts "[DEBUG][MIDDLEWARE] Original data: #{data}"
STDERR.puts "[DEBUG][MIDDLEWARE] Decoded request: #{decoded}"
STDERR.puts "[DEBUG][MIDDLEWARE] Encoded data: #{encoded}"
end
examples/echo_modifier.sh
+3 -3
View File
@@ -4,7 +4,7 @@ while read line; do
encoded=$(echo "$decoded" | xxd -p | tr -d "\\n")
echo "$encoded"
>&2 echo "[DEBUG] Original data: $line"
>&2 echo "[DEBUG] Decoded request: $decoded"
>&2 echo "[DEBUG] Encoded data: $encoded"
>&2 echo "[DEBUG][MIDDLEWARE] Original data: $line"
>&2 echo "[DEBUG][MIDDLEWARE] Decoded request: $decoded"
>&2 echo "[DEBUG][MIDDLEWARE] Encoded data: $encoded"
done;
middleware.go
+10 -9
View File
@@ -12,7 +12,6 @@ import (
)
type Middleware struct {
plugin interface{}
command string
data chan []byte
@@ -21,11 +20,10 @@ type Middleware struct {
Stdout io.Reader
}
func NewMiddleware(plugin interface{}, command string) io.Reader {
func NewMiddleware(command string) *Middleware {
m := new(Middleware)
m.plugin = plugin
m.command = command
m.data = make(chan []byte)
m.data = make(chan []byte, 1000)
commands := strings.Split(command, " ")
cmd := exec.Command(commands[0], commands[1:]...)
@@ -34,7 +32,6 @@ func NewMiddleware(plugin interface{}, command string) io.Reader {
m.Stdin, _ = cmd.StdinPipe()
cmd.Stderr = os.Stderr
go m.copy(m.Stdin, m.plugin.(io.Reader))
go m.read(m.Stdout)
go func() {
@@ -43,13 +40,17 @@ func NewMiddleware(plugin interface{}, command string) io.Reader {
if err != nil {
log.Fatal(err)
}
cmd.Wait()
}()
defer cmd.Wait()
return m
}
func (m *Middleware) ReadFrom(plugin io.Reader) {
go m.copy(m.Stdin, plugin)
}
func (m *Middleware) copy(to io.Writer, from io.Reader) {
buf := make([]byte, 5*1024*1024)
dst := make([]byte, len(buf)*2)
@@ -59,7 +60,7 @@ func (m *Middleware) copy(to io.Writer, from io.Reader) {
if nr > 0 && len(buf) > nr {
hex.Encode(dst, buf[0:nr])
to.Write(dst[0 : nr*2])
to.Write([]byte("\r\n"))
to.Write([]byte("\n"))
}
}
}
@@ -94,5 +95,5 @@ func (m *Middleware) Read(data []byte) (int, error) {
}
func (m *Middleware) String() string {
return fmt.Sprintf("Modifying traffic for %s using '%s' command", m.plugin, m.command)
return fmt.Sprintf("Modifying traffic using '%s' command", m.command)
}
middleware_test.go
+76 -9
View File
@@ -9,14 +9,18 @@ import (
"testing"
"strings"
"github.com/buger/gor/proto"
"net/http/httptest"
"encoding/hex"
"time"
)
type fakeServiceCb func(string, int, []byte)
// Simple service that generate token on request, and require this token for accesing to secure area
func NewFakeSecureService(wg *sync.WaitGroup) string {
func NewFakeSecureService(wg *sync.WaitGroup, cb fakeServiceCb) string {
active_tokens := make([]string, 0)
listener := startHTTP(func(w http.ResponseWriter, req *http.Request) {
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
Debug("Received request: " + req.URL.String())
switch req.URL.Path {
@@ -29,6 +33,8 @@ func NewFakeSecureService(wg *sync.WaitGroup) string {
active_tokens = append(active_tokens, token)
w.Write([]byte(token))
cb(req.URL.Path, 200, []byte(token))
case "/secure":
token := req.URL.Query().Get("token")
token_found := false
@@ -42,15 +48,17 @@ func NewFakeSecureService(wg *sync.WaitGroup) string {
if token_found {
w.WriteHeader(http.StatusAccepted)
cb(req.URL.Path, 202, []byte(nil))
} else {
w.WriteHeader(http.StatusForbidden)
cb(req.URL.Path, 403, []byte(nil))
}
}
wg.Done()
})
}))
address := strings.Replace(listener.Addr().String(), "[::]", "127.0.0.1", -1)
address := strings.Replace(server.Listener.Addr().String(), "[::]", "127.0.0.1", -1)
return address
}
@@ -59,7 +67,9 @@ func TestFakeSecureService(t *testing.T) {
wg := new(sync.WaitGroup)
addr := NewFakeSecureService(wg)
addr := NewFakeSecureService(wg, func(path string, status int, resp []byte){
})
wg.Add(3)
@@ -82,13 +92,66 @@ func TestFakeSecureService(t *testing.T) {
wg.Wait()
}
func TestMiddleware(t *testing.T) {
func TestEchoMiddleware(t *testing.T) {
wg := new(sync.WaitGroup)
from := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
wg.Done()
}))
to := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
wg.Done()
}))
quit := make(chan int)
// Catch traffic from one service
input := NewRAWInput(from.Listener.Addr().String())
// And redirect to another
output := NewHTTPOutput(to.URL, &HTTPOutputConfig{})
Plugins.Inputs = []io.Reader{input}
Plugins.Outputs = []io.Writer{output}
Settings.middleware = "./examples/echo_modifier.sh"
// Start Gor
go Start(quit)
time.Sleep(time.Millisecond)
// Should receive 2 requests from original + 2 from replayed
wg.Add(4)
client := NewHTTPClient(from.URL, &HTTPClientConfig{Debug: true})
// Request should be echoed
client.Get("/")
client.Get("/")
wg.Wait()
close(quit)
Settings.middleware = ""
}
func TestTokenMiddleware(t *testing.T) {
var resp, token []byte
wg := new(sync.WaitGroup)
from := NewFakeSecureService(wg)
to := NewFakeSecureService(wg)
from := NewFakeSecureService(wg, func(path string, status int, tok []byte){
})
to := NewFakeSecureService(wg, func(path string, status int, tok []byte){
switch path {
case "/token":
if bytes.Equal(token, tok) {
t.Error("Tokens should not match")
}
case "/secure":
if status != 202 {
// t.Error("Server should receive valid rewritten token")
}
}
})
quit := make(chan int)
@@ -100,10 +163,13 @@ func TestMiddleware(t *testing.T) {
Plugins.Inputs = []io.Reader{input}
Plugins.Outputs = []io.Writer{output}
// Settings.middleware = "./examples/echo_modifier.sh"
// Start Gor
go Start(quit)
time.Sleep(time.Millisecond)
// Should receive 2 requests from original + 2 from replayed
wg.Add(4)
@@ -120,4 +186,5 @@ func TestMiddleware(t *testing.T) {
wg.Wait()
close(quit)
}
Settings.middleware = ""
}
plugins.go
-7
View File
@@ -9,12 +9,8 @@ import (
type InOutPlugins struct {
Inputs []io.Reader
Outputs []io.Writer
Middleware []Middleware
}
type ReaderOrWriter interface{}
var Plugins *InOutPlugins = new(InOutPlugins)
func extractLimitOptions(options string) (string, string) {
@@ -56,9 +52,6 @@ func registerPlugin(constructor interface{}, options ...interface{}) {
}
if _, ok := plugin.(io.Reader); ok {
if len(Settings.middleware) > 0 {
plugin_wrapper = NewMiddleware(plugin_wrapper, Settings.middleware)
}
Plugins.Inputs = append(Plugins.Inputs, plugin_wrapper.(io.Reader))
}