Middleware for custom request rewrite logic

Dockerfile

@@ -2,8 +2,10 @@ FROM google/golang
 
 RUN cd /goroot/src/ && GOOS=linux GOARCH=386 ./make.bash --no-clean
 
+RUN apt-get update && apt-get install ruby vim-common -y
+
 WORKDIR /gopath/src/github.com/buger/gor/
 
 ADD . /gopath/src/github.com/buger/gor/
 
-RUN go get
+RUN go get

Makefile

@@ -1,4 +1,4 @@
-SOURCE = emitter.go gor.go gor_stat.go input_dummy.go input_file.go input_raw.go input_tcp.go limiter.go output_dummy.go output_file.go input_http.go output_http.go output_tcp.go plugins.go settings.go test_input.go elasticsearch.go http_modifier.go http_modifier_settings.go http_client.go
+SOURCE = emitter.go gor.go gor_stat.go input_dummy.go input_file.go input_raw.go input_tcp.go limiter.go output_dummy.go output_file.go input_http.go output_http.go output_tcp.go plugins.go settings.go test_input.go elasticsearch.go http_modifier.go http_modifier_settings.go http_client.go traffic_modifier.go
 
 SOURCE_PATH = /gopath/src/github.com/buger/gor/
 
@@ -31,7 +31,7 @@ dbench:
 
 # Used mainly for debugging, because docker container do not have access to parent machine ports
 drun:
-	docker run -v `pwd`:$(SOURCE_PATH) -t -i gor go run $(SOURCE) --input-dummy=0 --output-http="http://localhost:9000"  --verbose -h
+	docker run -v `pwd`:/gopath/src/gor -t -i gor go run $(SOURCE) --input-modifier="bash ./examples/echo_modifier.sh" --input-dummy=0 --output-http="http://localhost:9000"  --verbose
 
-dbash: 
-	docker run -v `pwd`:$(SOURCE_PATH) -t -i gor /bin/bash
+dbash:
+	docker run -v `pwd`:$(SOURCE_PATH) -t -i gor /bin/bash

emitter.go

@@ -27,6 +27,7 @@ func CopyMulty(src io.Reader, writers ...io.Writer) (err error) {
 
 	for {
 		nr, er := src.Read(buf)
+
 		if nr > 0 && len(buf) > nr {
 			payload := buf[0:nr]
 

examples/echo_modifier.rb

@@ -0,0 +1,16 @@
+#!/usr/bin/env ruby
+# encoding: utf-8
+while data = STDIN.gets
+    next unless data
+    data = data.chomp
+
+    decoded = [data].pack("H*")
+    encoded = decoded.unpack("H*").first
+
+    STDOUT.puts encoded
+
+
+    STDERR.puts "[DEBUG] Original data: #{data}"
+    STDERR.puts "[DEBUG] Decoded request: #{decoded}"
+    STDERR.puts "[DEBUG] Encoded data: #{encoded}"
+end

examples/echo_modifier.sh

@@ -0,0 +1,10 @@
+#!/usr/bin/env bash
+while read line; do
+    decoded=$(echo "$line" | xxd -r -p)
+    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"
+done;

input_dummy.go

@@ -30,7 +30,7 @@ func (i *DummyInput) emit() {
 	for {
 		select {
 		case <-ticker.C:
-			i.data <- []byte("GET / HTTP/1.1\r\n\r\n")
+			i.data <- []byte("POST /pub/WWW/å HTTP/1.1\nHost: www.w3.org\r\n\r\na=1&b=2")
 		}
 	}
 }

input_raw.go

@@ -32,6 +32,8 @@ func (i *RAWInput) Read(data []byte) (int, error) {
 func (i *RAWInput) listen(address string) {
 	address = strings.Replace(address, "[::]", "127.0.0.1", -1)
 
+	Debug("Listening for traffic on: " + address)
+
 	host, port, err := net.SplitHostPort(address)
 
 	if err != nil {

input_raw_test.go

@@ -18,7 +18,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) {
@@ -50,7 +50,7 @@ func TestInputRAW100Expect(t *testing.T) {
 	file_content, _ := ioutil.ReadFile("README.md")
 
 	// Origing and Replay server initialization
-	origin := startHTTP(func(req *http.Request) {
+	origin := startHTTP(func(w http.ResponseWriter, req *http.Request) {
 		defer req.Body.Close()
 		ioutil.ReadAll(req.Body)
 
@@ -69,7 +69,7 @@ func TestInputRAW100Expect(t *testing.T) {
 		wg.Done()
 	})
 
-	listener := startHTTP(func(req *http.Request) {
+	listener := startHTTP(func(w http.ResponseWriter, req *http.Request) {
 		defer req.Body.Close()
 		body, _ := ioutil.ReadAll(req.Body)
 
@@ -107,7 +107,7 @@ func TestInputRAWChunkedEncoding(t *testing.T) {
 	file_content, _ := ioutil.ReadFile("README.md")
 
 	// Origing and Replay server initialization
-	origin := startHTTP(func(req *http.Request) {
+	origin := startHTTP(func(w http.ResponseWriter, req *http.Request) {
 		defer req.Body.Close()
 		ioutil.ReadAll(req.Body)
 
@@ -118,7 +118,7 @@ func TestInputRAWChunkedEncoding(t *testing.T) {
 
 	input := NewRAWInput(origin_address)
 
-	listener := startHTTP(func(req *http.Request) {
+	listener := startHTTP(func(w http.ResponseWriter, req *http.Request) {
 		defer req.Body.Close()
 		body, _ := ioutil.ReadAll(req.Body)
 

output_http.go

@@ -29,6 +29,7 @@ type HTTPOutput struct {
 	address string
 	limit   int
 	queue   chan []byte
+	responses   chan []byte
 
 	needWorker chan int
 
@@ -40,7 +41,6 @@ type HTTPOutput struct {
 }
 
 func NewHTTPOutput(address string, config *HTTPOutputConfig) io.Writer {
-
 	o := new(HTTPOutput)
 
 	o.address = address

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")
@@ -30,7 +30,7 @@ func TestHTTPOutput(t *testing.T) {
 
 	input := NewTestInput()
 
-	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")
 		}
@@ -106,7 +106,7 @@ func BenchmarkHTTPOutput(b *testing.B) {
 	wg := new(sync.WaitGroup)
 	quit := make(chan int)
 
-	listener := startHTTP(func(req *http.Request) {
+	listener := startHTTP(func(w http.ResponseWriter, req *http.Request) {
 		time.Sleep(50 * time.Millisecond)
 		wg.Done()
 	})

plugins.go

@@ -9,6 +9,8 @@ import (
 type InOutPlugins struct {
 	Inputs  []io.Reader
 	Outputs []io.Writer
+
+	Modifiers []TrafficModifier
 }
 
 type ReaderOrWriter interface{}
@@ -54,6 +56,9 @@ func registerPlugin(constructor interface{}, options ...interface{}) {
 	}
 
 	if _, ok := plugin.(io.Reader); ok {
+		for _, options := range Settings.inputModifier {
+			plugin_wrapper = NewTrafficModifier(plugin_wrapper, options)
+		}
 		Plugins.Inputs = append(Plugins.Inputs, plugin_wrapper.(io.Reader))
 	}
 

settings.go

@@ -41,6 +41,8 @@ type AppSettings struct {
 
 	inputRAW MultiOption
 
+	inputModifier MultiOption
+
 	inputHTTP  MultiOption
 	outputHTTP MultiOption
 
@@ -76,6 +78,8 @@ func init() {
 
 	flag.Var(&Settings.inputRAW, "input-raw", "Capture traffic from given port (use RAW sockets and require *sudo* access):\n\t# Capture traffic from 8080 port\n\tgor --input-raw :8080 --output-http staging.com")
 
+	flag.Var(&Settings.inputModifier, "input-modifier", "Used for modifying input traffic using external command")
+
 	flag.Var(&Settings.inputHTTP, "input-http", "Read requests from HTTP, should be explicitly sent from your application:\n\t# Listen for http on 9000\n\tgor --input-http :9000 --output-http staging.com")
 
 	flag.Var(&Settings.outputHTTP, "output-http", "Forwards incoming requests to given http address.\n\t# Redirect all incoming requests to staging.com address \n\tgor --input-raw :80 --output-http http://staging.com")

traffic_modifier.go

@@ -0,0 +1,98 @@
+package main
+
+import (
+	"bufio"
+	"encoding/hex"
+	"fmt"
+	"io"
+	"log"
+	"os"
+	"os/exec"
+    "strings"
+)
+
+type TrafficModifier struct {
+	plugin  interface{}
+	command string
+
+	data chan []byte
+
+	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)
+
+    commands := strings.Split(command, " ")
+	cmd := exec.Command(commands[0], commands[1:]...)
+
+	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 func() {
+		err := cmd.Start()
+
+		if err != nil {
+			log.Fatal(err)
+		}
+	}()
+
+	defer cmd.Wait()
+
+	return m
+}
+
+func (m *TrafficModifier) copy(to io.Writer, from io.Reader) {
+	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"))
+		}
+	}
+}
+
+func (m *TrafficModifier) read(from io.Reader) {
+	buf := make([]byte, 5*1024*1024)
+
+	scanner := bufio.NewScanner(from)
+
+	for scanner.Scan() {
+		bytes := scanner.Bytes()
+		hex.Decode(buf, bytes)
+
+		Debug("Received:", 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)
+	}
+
+	return
+}
+
+func (m *TrafficModifier) Read(data []byte) (int, error) {
+	Debug("Trying to read channel!")
+	buf := <-m.data
+	copy(data, buf)
+
+	return len(buf), nil
+}
+
+func (m *TrafficModifier) String() string {
+	return fmt.Sprintf("Modifying traffic for %s using '%s' command", m.plugin, m.command)
+}