diff --git a/emitter.go b/emitter.go index ec5065c..0769295 100644 --- a/emitter.go +++ b/emitter.go @@ -56,8 +56,13 @@ func CopyMulty(src io.Reader, writers ...io.Writer) (err error) { if nr > 0 && len(buf) > nr { payload := buf[0:nr] + _maxN := nr + if nr > 500 { + _maxN = 500 + } + if Settings.debug { - Debug("[EMITTER] input:", string(payload[0:500])) + Debug("[EMITTER] input:", string(payload[0:_maxN])) } if modifier != nil { @@ -69,7 +74,7 @@ func CopyMulty(src io.Reader, writers ...io.Writer) (err error) { } if Settings.debug { - Debug("[EMITTER] Rewrittern input:", len(payload), "First 500 bytes:", string(payload[0:500])) + Debug("[EMITTER] Rewrittern input:", len(payload), "First 500 bytes:", string(payload[0:_maxN])) } } diff --git a/examples/echo_modifier.sh b/examples/echo_modifier.sh index e3f7dd9..105b4bb 100755 --- a/examples/echo_modifier.sh +++ b/examples/echo_modifier.sh @@ -11,6 +11,9 @@ while read line; do >&2 echo "[DEBUG][ECHO] ===================================" case ${header:0:1} in + "3") + >&2 echo "[DEBUG][ECHO] Request type: Original Response" + ;; "2") >&2 echo "[DEBUG][ECHO] Request type: Replayed Response" ;; diff --git a/input_raw.go b/input_raw.go index 6fdb9ee..762655c 100644 --- a/input_raw.go +++ b/input_raw.go @@ -10,17 +10,21 @@ import ( // RAWInput used for intercepting traffic for given address type RAWInput struct { - data chan []byte + requests chan []byte + responses chan []byte address string expire time.Duration + captureResponse bool } // NewRAWInput constructor for RAWInput. Accepts address with port as argument. -func NewRAWInput(address string, expire time.Duration) (i *RAWInput) { +func NewRAWInput(address string, expire time.Duration, captureResponse bool) (i *RAWInput) { i = new(RAWInput) - i.data = make(chan []byte) + i.requests = make(chan []byte) + i.responses = make(chan []byte) i.address = address i.expire = expire + i.captureResponse = captureResponse go i.listen(address) @@ -28,18 +32,25 @@ func NewRAWInput(address string, expire time.Duration) (i *RAWInput) { } func (i *RAWInput) Read(data []byte) (int, error) { - buf := <-i.data + select { + case buf := <- i.requests: + if i.captureResponse { + header := []byte("1\n") + copy(data[0:len(header)], header) + copy(data[len(header):], buf) - if len(Settings.middleware) > 0 { - header := []byte("1\n") + return len(buf) + len(header), nil + } else { + copy(data, buf) + + return len(buf), nil + } + case buf := <- i.responses: + header := []byte("3\n") copy(data[0:len(header)], header) copy(data[len(header):], buf) return len(buf) + len(header), nil - } else { - copy(data, buf) - - return len(buf), nil } } @@ -54,13 +65,17 @@ func (i *RAWInput) listen(address string) { log.Fatal("input-raw: error while parsing address", err) } - listener := raw.NewListener(host, port, i.expire) + listener := raw.NewListener(host, port, i.expire, i.captureResponse) for { // Receiving TCPMessage object m := listener.Receive() - i.data <- m.Bytes() + i.requests <- m.RequestBytes() + + if i.captureResponse { + i.responses <- m.ResponseBytes() + } } } diff --git a/input_raw_test.go b/input_raw_test.go index a145845..c822029 100644 --- a/input_raw_test.go +++ b/input_raw_test.go @@ -24,7 +24,7 @@ func TestRAWInput(t *testing.T) { listener := startHTTP(func(w http.ResponseWriter, req *http.Request) {}) - input := NewRAWInput(listener.Addr().String(), testRawExpire) + input := NewRAWInput(listener.Addr().String(), testRawExpire, false) output := NewTestOutput(func(data []byte) { wg.Done() }) @@ -66,7 +66,7 @@ func TestInputRAW100Expect(t *testing.T) { originAddr := strings.Replace(origin.Addr().String(), "[::]", "127.0.0.1", -1) - input := NewRAWInput(originAddr, testRawExpire) + input := NewRAWInput(originAddr, testRawExpire, false) // We will use it to get content of raw HTTP request testOutput := NewTestOutput(func(data []byte) { @@ -123,7 +123,7 @@ func TestInputRAWChunkedEncoding(t *testing.T) { originAddr := strings.Replace(origin.Addr().String(), "[::]", "127.0.0.1", -1) - input := NewRAWInput(originAddr, testRawExpire) + input := NewRAWInput(originAddr, testRawExpire, false) listener := startHTTP(func(w http.ResponseWriter, req *http.Request) { defer req.Body.Close() @@ -181,7 +181,7 @@ func TestInputRAWLargePayload(t *testing.T) { })) originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1) - input := NewRAWInput(originAddr, testRawExpire) + input := NewRAWInput(originAddr, testRawExpire, false) replay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { req.Body = http.MaxBytesReader(w, req.Body, 1*1024*1024) @@ -214,3 +214,36 @@ func TestInputRAWLargePayload(t *testing.T) { wg.Wait() close(quit) } + +func TestInputRAWResponse(t *testing.T) { + wg := new(sync.WaitGroup) + quit := make(chan int) + + listener := startHTTP(func(w http.ResponseWriter, req *http.Request) {}) + + input := NewRAWInput(listener.Addr().String(), testRawExpire, true) + output := NewTestOutput(func(data []byte) { + wg.Done() + }) + + Plugins.Inputs = []io.Reader{input} + Plugins.Outputs = []io.Writer{output} + + address := strings.Replace(listener.Addr().String(), "[::]", "127.0.0.1", -1) + + client := NewHTTPClient(address, &HTTPClientConfig{}) + + time.Sleep(time.Millisecond) + go Start(quit) + + for i := 0; i < 100; i++ { + // 2 because we track both request and response + wg.Add(2) + client.Get("/") + } + + wg.Wait() + close(quit) + + time.Sleep(100*time.Millisecond) +} \ No newline at end of file diff --git a/middleware_test.go b/middleware_test.go index 71748a8..e323fd9 100644 --- a/middleware_test.go +++ b/middleware_test.go @@ -105,7 +105,7 @@ func TestEchoMiddleware(t *testing.T) { quit := make(chan int) // Catch traffic from one service - input := NewRAWInput(from.Listener.Addr().String(), testRawExpire) + input := NewRAWInput(from.Listener.Addr().String(), testRawExpire, true) // And redirect to another output := NewHTTPOutput(to.URL, &HTTPOutputConfig{Debug: true}) @@ -158,14 +158,14 @@ func TestTokenMiddleware(t *testing.T) { quit := make(chan int) // Catch traffic from one service - input := NewRAWInput(from, testRawExpire) + input := NewRAWInput(from, testRawExpire, true) // And redirect to another output := NewHTTPOutput(to, &HTTPOutputConfig{}) Plugins.Inputs = []io.Reader{input} Plugins.Outputs = []io.Writer{output} - // Settings.middleware = "./examples/echo_modifier.sh" + Settings.middleware = "./examples/echo_modifier.sh" // Start Gor go Start(quit) diff --git a/plugins.go b/plugins.go index d0ba71a..c92435c 100644 --- a/plugins.go +++ b/plugins.go @@ -58,6 +58,7 @@ func registerPlugin(constructor interface{}, options ...interface{}) { _, isR := plugin.(io.Reader) _, isW := plugin.(io.Writer) + // Some of the output can be Readers as well because return responses if isR && !isW { Plugins.Inputs = append(Plugins.Inputs, pluginWrapper.(io.Reader)) } diff --git a/raw_socket_listener/listener.go b/raw_socket_listener/listener.go index 53857ac..ef8b283 100644 --- a/raw_socket_listener/listener.go +++ b/raw_socket_listener/listener.go @@ -32,6 +32,8 @@ type Listener struct { // To get ACK of second message we need to compute its Seq and wait for them message seqWithData map[uint32]uint32 + respAliases map[uint32]uint32 + // Messages ready to be send to client packetsChan chan *TCPPacket @@ -42,14 +44,16 @@ type Listener struct { messageDelChan chan *TCPMessage addr string // IP to listen - port int // Port to listen + port uint16 // Port to listen messageExpire time.Duration + + captureResponse bool } // NewListener creates and initializes new Listener object -func NewListener(addr string, port string, expire time.Duration) (rawListener *Listener) { - rawListener = &Listener{} +func NewListener(addr string, port string, expire time.Duration, captureResponse bool) (rawListener *Listener) { + rawListener = &Listener{captureResponse: captureResponse} rawListener.packetsChan = make(chan *TCPPacket, 10000) rawListener.messagesChan = make(chan *TCPMessage, 10000) @@ -58,9 +62,11 @@ func NewListener(addr string, port string, expire time.Duration) (rawListener *L rawListener.messages = make(map[string]*TCPMessage) rawListener.ackAliases = make(map[uint32]uint32) rawListener.seqWithData = make(map[uint32]uint32) + rawListener.respAliases = make(map[uint32]uint32) rawListener.addr = addr - rawListener.port, _ = strconv.Atoi(port) + _port, _ := strconv.Atoi(port) + rawListener.port = uint16(_port) if expire.Nanoseconds() == 0 { expire = 2000 * time.Millisecond @@ -115,18 +121,19 @@ func (t *Listener) readRAWSocket() { } func (t *Listener) parsePacket(addr net.Addr, buf []byte) { - if t.isIncomingDataPacket(buf) { + if t.isValidPacket(buf) { t.packetsChan <- ParseTCPPacket(addr, buf) } } -func (t *Listener) isIncomingDataPacket(buf []byte) bool { +func (t *Listener) isValidPacket(buf []byte) bool { // To avoid full packet parsing every time, we manually parsing values needed for packet filtering // http://en.wikipedia.org/wiki/Transmission_Control_Protocol destPort := binary.BigEndian.Uint16(buf[2:4]) + srcPort := binary.BigEndian.Uint16(buf[0:2]) // Because RAW_SOCKET can't be bound to port, we have to control it by ourself - if int(destPort) == t.port { + if destPort == t.port || (t.captureResponse && srcPort == t.port) { // Get the 'data offset' (size of the TCP header in 32-bit words) dataOffset := (buf[12] & 0xF0) >> 4 @@ -161,12 +168,20 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) { packet.Ack = alias } - mID := packet.Addr.String() + strconv.Itoa(int(packet.SrcPort)) + strconv.Itoa(int(packet.Ack)) + // if response + if t.captureResponse && packet.DestPort != t.port { + if alias, ok := t.respAliases[packet.Ack]; ok { + packet.Ack = alias + } + } + + + mID := packet.Addr.String() + strconv.Itoa(int(packet.Ack)) message, ok := t.messages[mID] if !ok { // We sending messageDelChan channel, so message object can communicate with Listener and notify it if message completed - message = NewTCPMessage(mID, t.messageDelChan, packet.Ack, &t.messageExpire) + message = NewTCPMessage(mID, t.messageDelChan, packet.Ack, &t.messageExpire, t.port) t.messages[mID] = message } @@ -179,6 +194,11 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) { } } + if t.captureResponse { + // Response tracking + t.respAliases[packet.Seq+uint32(len(packet.Data))] = packet.Ack + } + // Adding packet to message message.packetsChan <- packet } diff --git a/raw_socket_listener/tcp_message.go b/raw_socket_listener/tcp_message.go index e6e31bd..5e0f8a6 100644 --- a/raw_socket_listener/tcp_message.go +++ b/raw_socket_listener/tcp_message.go @@ -15,6 +15,7 @@ import ( type TCPMessage struct { ID string // Message ID Ack uint32 + port uint16 packets []*TCPPacket timer *time.Timer // Used for expire check @@ -27,8 +28,8 @@ type TCPMessage struct { } // NewTCPMessage pointer created from a Acknowledgment number and a channel of messages readuy to be deleted -func NewTCPMessage(ID string, delChan chan *TCPMessage, Ack uint32, expire *time.Duration) (msg *TCPMessage) { - msg = &TCPMessage{ID: ID, Ack: Ack, expire: expire} +func NewTCPMessage(ID string, delChan chan *TCPMessage, Ack uint32, expire *time.Duration, port uint16) (msg *TCPMessage) { + msg = &TCPMessage{ID: ID, Ack: Ack, expire: expire, port: port} msg.packetsChan = make(chan *TCPPacket) msg.delChan = delChan // used for notifying that message completed or expired @@ -73,11 +74,26 @@ func (t *TCPMessage) Timeout() { } // Bytes sorts packets in right orders and return message content -func (t *TCPMessage) Bytes() (output []byte) { +func (t *TCPMessage) RequestBytes() (output []byte) { sort.Sort(sortBySeq(t.packets)) for _, v := range t.packets { - output = append(output, v.Data...) + if v.DestPort == t.port { + output = append(output, v.Data...) + } + } + + return output +} + +// Bytes sorts packets in right orders and return message content +func (t *TCPMessage) ResponseBytes() (output []byte) { + sort.Sort(sortBySeq(t.packets)) + + for _, v := range t.packets { + if v.DestPort != t.port { + output = append(output, v.Data...) + } } return output diff --git a/raw_socket_listener/tcp_packet.go b/raw_socket_listener/tcp_packet.go index 3ebcf88..998eeb4 100644 --- a/raw_socket_listener/tcp_packet.go +++ b/raw_socket_listener/tcp_packet.go @@ -50,7 +50,6 @@ func ParseTCPPacket(addr net.Addr, b []byte) (p *TCPPacket) { // Parse TCP Packet, inspired by: https://github.com/miekg/pcap/blob/master/packet.go func (t *TCPPacket) Parse() { t.ParseBasic() - t.DestPort = binary.BigEndian.Uint16(t.Data[2:4]) t.Flags = binary.BigEndian.Uint16(t.Data[12:14]) & 0x1FF t.Window = binary.BigEndian.Uint16(t.Data[14:16]) t.Checksum = binary.BigEndian.Uint16(t.Data[16:18]) @@ -59,6 +58,7 @@ func (t *TCPPacket) Parse() { // ParseBasic set of fields func (t *TCPPacket) ParseBasic() { + t.DestPort = binary.BigEndian.Uint16(t.Data[2:4]) t.SrcPort = binary.BigEndian.Uint16(t.Data[0:2]) t.Seq = binary.BigEndian.Uint32(t.Data[4:8]) t.Ack = binary.BigEndian.Uint32(t.Data[8:12])