diff --git a/emitter.go b/emitter.go index 3fdbac4..8237f6f 100644 --- a/emitter.go +++ b/emitter.go @@ -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) diff --git a/examples/echo_modifier.rb b/examples/echo_modifier.rb index e3a489d..a20beb8 100755 --- a/examples/echo_modifier.rb +++ b/examples/echo_modifier.rb @@ -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 \ No newline at end of file diff --git a/examples/echo_modifier.sh b/examples/echo_modifier.sh index 1d56cdc..22c7386 100755 --- a/examples/echo_modifier.sh +++ b/examples/echo_modifier.sh @@ -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; diff --git a/middleware.go b/middleware.go index 5b20f77..dc387eb 100644 --- a/middleware.go +++ b/middleware.go @@ -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) } diff --git a/middleware_test.go b/middleware_test.go index 7ce251a..e069d23 100644 --- a/middleware_test.go +++ b/middleware_test.go @@ -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 = "" +} \ No newline at end of file diff --git a/plugins.go b/plugins.go index adfc242..321317e 100644 --- a/plugins.go +++ b/plugins.go @@ -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)) }