diff --git a/Dockerfile b/Dockerfile index 98fb7a4..91435a0 100644 --- a/Dockerfile +++ b/Dockerfile @@ -15,8 +15,8 @@ RUN wget http://www.tcpdump.org/release/libpcap-1.7.4.tar.gz && tar xzf libpcap- RUN go get github.com/google/gopacket RUN go get -u github.com/golang/lint/golint -WORKDIR /go/src/github.com/buger/gor/ -ADD . /go/src/github.com/buger/gor/ +WORKDIR /go/src/github.com/buger/gor-pro/ +ADD . /go/src/github.com/buger/gor-pro/ RUN javac -cp /tmp/commons-io-2.4/commons-io-2.4.jar ./examples/middleware/echo.java RUN go get \ No newline at end of file diff --git a/Makefile b/Makefile index c0c1e45..f7807d8 100644 --- a/Makefile +++ b/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_null.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 middleware.go protocol.go output_file_settings.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_null.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 middleware.go protocol.go output_file_settings.go tcp_client.go output_binary.go SOURCE_PATH = /go/src/github.com/buger/gor/ PORT = 8000 FADDR = :8000 diff --git a/input_raw.go b/input_raw.go index 0236a2c..f733670 100644 --- a/input_raw.go +++ b/input_raw.go @@ -1,8 +1,8 @@ package main import ( - "github.com/buger/gor/proto" - raw "github.com/buger/gor/raw_socket_listener" + "github.com/buger/gor-pro/proto" + raw "github.com/buger/gor-pro/raw_socket_listener" "log" "net" "time" @@ -18,6 +18,7 @@ type RAWInput struct { realIPHeader []byte trackResponse bool listener *raw.Listener + protocol raw.TCPProtocol } // Available engines for intercepting traffic @@ -28,7 +29,7 @@ const ( ) // NewRAWInput constructor for RAWInput. Accepts address with port as argument. -func NewRAWInput(address string, engine int, trackResponse bool, expire time.Duration, realIPHeader string) (i *RAWInput) { +func NewRAWInput(address string, engine int, trackResponse bool, expire time.Duration, realIPHeader string, protocol string) (i *RAWInput) { i = new(RAWInput) i.data = make(chan *raw.TCPMessage) i.address = address @@ -38,6 +39,15 @@ func NewRAWInput(address string, engine int, trackResponse bool, expire time.Dur i.quit = make(chan bool) i.trackResponse = trackResponse + switch protocol { + case "http": + i.protocol = raw.ProtocolHTTP + case "binary": + i.protocol = raw.ProtocolBinary + default: + log.Fatal("Unsupported protocol:", protocol) + } + i.listen(address) i.listener.IsReady() @@ -80,7 +90,7 @@ func (i *RAWInput) listen(address string) { log.Fatal("input-raw: error while parsing address", err) } - i.listener = raw.NewListener(host, port, i.engine, i.trackResponse, i.expire) + i.listener = raw.NewListener(host, port, i.engine, i.trackResponse, i.expire, i.protocol) ch := i.listener.Receiver() diff --git a/input_raw_test.go b/input_raw_test.go index b8fe1b1..b44e37e 100644 --- a/input_raw_test.go +++ b/input_raw_test.go @@ -43,7 +43,7 @@ func TestRAWInputIPv4(t *testing.T) { var respCounter, reqCounter int64 - input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "X-Real-IP") + input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "X-Real-IP", "http") defer input.Close() output := NewTestOutput(func(data []byte) { @@ -102,7 +102,7 @@ func TestRAWInputIPv6(t *testing.T) { var respCounter, reqCounter int64 - input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "") + input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "", "http") defer input.Close() output := NewTestOutput(func(data []byte) { @@ -153,7 +153,7 @@ func TestInputRAW100Expect(t *testing.T) { originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1) - input := NewRAWInput(originAddr, EnginePcap, true, time.Second, "") + input := NewRAWInput(originAddr, EnginePcap, true, time.Second, "", "http") defer input.Close() // We will use it to get content of raw HTTP request @@ -216,7 +216,7 @@ func TestInputRAWChunkedEncoding(t *testing.T) { })) originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1) - input := NewRAWInput(originAddr, EnginePcap, true, time.Second, "") + input := NewRAWInput(originAddr, EnginePcap, true, time.Second, "", "http") defer input.Close() replay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -280,7 +280,7 @@ func TestInputRAWLargePayload(t *testing.T) { })) originAddr := strings.Replace(origin.Listener.Addr().String(), "[::]", "127.0.0.1", -1) - input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "") + input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "", "http") defer input.Close() replay := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { @@ -325,7 +325,7 @@ func BenchmarkRAWInput(b *testing.B) { var respCounter, reqCounter int64 - input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "") + input := NewRAWInput(originAddr, EnginePcap, true, testRawExpire, "", "http") defer input.Close() output := NewTestOutput(func(data []byte) { diff --git a/middleware_test.go b/middleware_test.go index beaed15..46f9b28 100644 --- a/middleware_test.go +++ b/middleware_test.go @@ -117,7 +117,7 @@ func TestEchoMiddleware(t *testing.T) { // Catch traffic from one service fromAddr := strings.Replace(from.Listener.Addr().String(), "[::]", "127.0.0.1", -1) - input := NewRAWInput(fromAddr, EnginePcap, true, testRawExpire, "") + input := NewRAWInput(fromAddr, EnginePcap, true, testRawExpire, "", "http") defer input.Close() // And redirect to another @@ -179,7 +179,7 @@ func TestTokenMiddleware(t *testing.T) { fromAddr := strings.Replace(from.Listener.Addr().String(), "[::]", "127.0.0.1", -1) // Catch traffic from one service - input := NewRAWInput(fromAddr, EnginePcap, true, testRawExpire, "") + input := NewRAWInput(fromAddr, EnginePcap, true, testRawExpire, "", "http") defer input.Close() // And redirect to another diff --git a/output_binary.go b/output_binary.go new file mode 100644 index 0000000..6134e65 --- /dev/null +++ b/output_binary.go @@ -0,0 +1,175 @@ +package main + +import ( + "io" + "sync/atomic" + "time" +) + +// BinaryOutputConfig struct for holding binary output configuration +type BinaryOutputConfig struct { + workers int + Timeout time.Duration + BufferSize int + Debug bool + TrackResponses bool +} + +// BinaryOutput plugin manage pool of workers which send request to replayed server +// By default workers pool is dynamic and starts with 10 workers +// You can specify fixed number of workers using `--output-tcp-workers` +type BinaryOutput struct { + // Keep this as first element of struct because it guarantees 64bit + // alignment. atomic.* functions crash on 32bit machines if operand is not + // aligned at 64bit. See https://github.com/golang/go/issues/599 + activeWorkers int64 + + address string + queue chan []byte + + responses chan response + + needWorker chan int + + config *BinaryOutputConfig + + queueStats *GorStat +} + +// NewBinaryOutput constructor for BinaryOutput +// Initialize workers +func NewBinaryOutput(address string, config *BinaryOutputConfig) io.Writer { + o := new(BinaryOutput) + + o.address = address + o.config = config + + o.queue = make(chan []byte, 1000) + o.responses = make(chan response, 1000) + o.needWorker = make(chan int, 1) + + // Initial workers count + if o.config.workers == 0 { + o.needWorker <- initialDynamicWorkers + } else { + o.needWorker <- o.config.workers + } + + if len(Settings.middleware) > 0 { + o.config.TrackResponses = true + } + + go o.workerMaster() + + return o +} + +func (o *BinaryOutput) workerMaster() { + for { + newWorkers := <-o.needWorker + for i := 0; i < newWorkers; i++ { + go o.startWorker() + } + + // Disable dynamic scaling if workers poll fixed size + if o.config.workers != 0 { + return + } + } +} + +func (o *BinaryOutput) startWorker() { + client := NewTCPClient(o.address, &TCPClientConfig{ + Debug: o.config.Debug, + Timeout: o.config.Timeout, + ResponseBufferSize: o.config.BufferSize, + }) + + deathCount := 0 + + atomic.AddInt64(&o.activeWorkers, 1) + + for { + select { + case data := <-o.queue: + o.sendRequest(client, data) + deathCount = 0 + case <-time.After(time.Millisecond * 100): + // When dynamic scaling enabled workers die after 2s of inactivity + if o.config.workers == 0 { + deathCount++ + } else { + continue + } + + if deathCount > 20 { + workersCount := atomic.LoadInt64(&o.activeWorkers) + + // At least 1 startWorker should be alive + if workersCount != 1 { + atomic.AddInt64(&o.activeWorkers, -1) + return + } + } + } + } +} + +func (o *BinaryOutput) Write(data []byte) (n int, err error) { + if !isRequestPayload(data) { + return len(data), nil + } + + buf := make([]byte, len(data)) + copy(buf, data) + + o.queue <- buf + + if o.config.workers == 0 { + workersCount := atomic.LoadInt64(&o.activeWorkers) + + if len(o.queue) > int(workersCount) { + o.needWorker <- len(o.queue) + } + } + + return len(data), nil +} + +func (o *BinaryOutput) Read(data []byte) (int, error) { + resp := <-o.responses + + Debug("[OUTPUT-TCP] Received response:", string(resp.payload)) + + header := payloadHeader(ReplayedResponsePayload, resp.uuid, resp.startedAt, resp.roundTripTime) + copy(data[0:len(header)], header) + copy(data[len(header):], resp.payload) + + return len(resp.payload) + len(header), nil +} + +func (o *BinaryOutput) sendRequest(client *TCPClient, request []byte) { + meta := payloadMeta(request) + if len(meta) < 2 { + return + } + uuid := meta[1] + + body := payloadBody(request) + + start := time.Now() + resp, err := client.Send(body) + stop := time.Now() + + if err != nil { + Debug("Request error:", err) + } + + if o.config.TrackResponses { + o.responses <- response{resp, uuid, start.UnixNano(), stop.UnixNano() - start.UnixNano()} + } +} + +func (o *BinaryOutput) String() string { + return "TCP output: " + o.address +} diff --git a/output_http.go b/output_http.go index d442b83..397c96c 100644 --- a/output_http.go +++ b/output_http.go @@ -5,7 +5,7 @@ import ( "sync/atomic" "time" - "github.com/buger/gor/proto" + "github.com/buger/gor-pro/proto" ) const initialDynamicWorkers = 10 diff --git a/plugins.go b/plugins.go index 1a0f6db..42e5699 100644 --- a/plugins.go +++ b/plugins.go @@ -107,7 +107,7 @@ func InitPlugins() { } for _, options := range Settings.inputRAW { - registerPlugin(NewRAWInput, options, engine, Settings.inputRAWTrackResponse, time.Duration(0), Settings.inputRAWRealIPHeader) + registerPlugin(NewRAWInput, options, engine, Settings.inputRAWTrackResponse, time.Duration(0), Settings.inputRAWRealIPHeader, Settings.inputRAWProtocol) } for _, options := range Settings.inputTCP { @@ -142,4 +142,8 @@ func InitPlugins() { for _, options := range Settings.outputHTTP { registerPlugin(NewHTTPOutput, options, &Settings.outputHTTPConfig) } + + for _, options := range Settings.outputBinary { + registerPlugin(NewBinaryOutput, options, &Settings.outputBinaryConfig) + } } diff --git a/protocol.go b/protocol.go index 50d6943..63c4560 100644 --- a/protocol.go +++ b/protocol.go @@ -69,6 +69,7 @@ func payloadHeader(payloadType byte, uuid []byte, timing int64, latency int64) ( copy(header[3+len(uuid):], sTime) if latency != -1 { + header[3+len(uuid)+len(sTime)] = ' ' copy(header[4+len(uuid)+len(sTime):], sLatency) } diff --git a/raw_socket_listener/listener.go b/raw_socket_listener/listener.go index a4c88fd..4d21cf9 100644 --- a/raw_socket_listener/listener.go +++ b/raw_socket_listener/listener.go @@ -16,7 +16,7 @@ import ( "bytes" "encoding/binary" "fmt" - "github.com/buger/gor/proto" + "github.com/buger/gor-pro/proto" "github.com/google/gopacket" "github.com/google/gopacket/layers" "github.com/google/gopacket/pcap" @@ -69,6 +69,8 @@ type Listener struct { quit chan bool readyCh chan bool + + protocol TCPProtocol } type request struct { @@ -85,7 +87,7 @@ const ( ) // NewListener creates and initializes new Listener object -func NewListener(addr string, port string, engine int, trackResponse bool, expire time.Duration) (l *Listener) { +func NewListener(addr string, port string, engine int, trackResponse bool, expire time.Duration, protocol TCPProtocol) (l *Listener) { l = &Listener{} l.packetsChan = make(chan []byte, 10000) @@ -99,6 +101,7 @@ func NewListener(addr string, port string, engine int, trackResponse bool, expir l.respAliases = make(map[uint32]*TCPMessage) l.respWithoutReq = make(map[uint32]tcpID) l.trackResponse = trackResponse + l.protocol = protocol l.addr = addr _port, _ := strconv.Atoi(port) @@ -176,7 +179,7 @@ func (t *Listener) dispatchMessage(message *TCPMessage) { t.deleteMessage(message) - if !message.complete { + if t.protocol == ProtocolHTTP && !message.complete { if !message.IsIncoming { delete(t.respAliases, message.Ack) delete(t.respWithoutReq, message.Ack) @@ -667,28 +670,30 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) { isIncoming := packet.DestPort == t.port - // Seek for 100-expect chunks - if parentAck, ok := t.seqWithData[packet.Seq]; ok { - // In case if non-first data chunks comes first - for _, m := range t.messages { - if m.Ack == packet.Ack && bytes.Equal(m.packets[0].Addr, packet.Addr) { - t.deleteMessage(m) + if t.protocol == ProtocolHTTP { + // Seek for 100-expect chunks + if parentAck, ok := t.seqWithData[packet.Seq]; ok { + // In case if non-first data chunks comes first + for _, m := range t.messages { + if m.Ack == packet.Ack && bytes.Equal(m.packets[0].Addr, packet.Addr) { + t.deleteMessage(m) - if m.AssocMessage != nil { - m.setAssocMessage(nil) - } + if m.AssocMessage != nil { + m.setAssocMessage(nil) + } - for _, pkt := range m.packets { - // log.Println("Updating ack", parentAck, pkt.Ack) - pkt.UpdateAck(parentAck) - // Re-queue this packets - t.processTCPPacket(pkt) + for _, pkt := range m.packets { + // log.Println("Updating ack", parentAck, pkt.Ack) + pkt.UpdateAck(parentAck) + // Re-queue this packets + t.processTCPPacket(pkt) + } } } - } - t.ackAliases[packet.Ack] = parentAck - packet.UpdateAck(parentAck) + t.ackAliases[packet.Ack] = parentAck + packet.UpdateAck(parentAck) + } } if alias, ok := t.ackAliases[packet.Ack]; ok { @@ -704,7 +709,7 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) { message, ok := t.messages[packet.ID] if !ok { - message = NewTCPMessage(packet.Seq, packet.Ack, isIncoming) + message = NewTCPMessage(packet.Seq, packet.Ack, isIncoming, t.protocol) t.messages[packet.ID] = message if !isIncoming { @@ -721,7 +726,7 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) { message.AddPacket(packet) // Handling Expect: 100-continue requests - if message.expectType == httpExpect100Continue && len(message.packets) == message.headerPacket+1 { + if t.protocol == ProtocolHTTP && message.expectType == httpExpect100Continue && len(message.packets) == message.headerPacket+1 { seq := packet.Seq + uint32(message.Size()) t.seqWithData[seq] = packet.Ack message.DataSeq = seq diff --git a/raw_socket_listener/listener_test.go b/raw_socket_listener/listener_test.go index b66b23a..eee85dd 100644 --- a/raw_socket_listener/listener_test.go +++ b/raw_socket_listener/listener_test.go @@ -12,7 +12,7 @@ import ( func TestRawListenerInput(t *testing.T) { var req, resp *TCPMessage - listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond) + listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP) defer listener.Close() reqPacket := buildPacket(true, 1, 1, []byte("GET / HTTP/1.1\r\n\r\n")) @@ -49,7 +49,7 @@ func TestRawListenerInput(t *testing.T) { func TestRawListenerInputWithoutResponse(t *testing.T) { var req *TCPMessage - listener := NewListener("", "0", EnginePcap, false, 10*time.Millisecond) + listener := NewListener("", "0", EnginePcap, false, 10*time.Millisecond, ProtocolHTTP) defer listener.Close() reqPacket := buildPacket(true, 1, 1, []byte("GET / HTTP/1.1\r\n\r\n")) @@ -71,7 +71,7 @@ func TestRawListenerInputWithoutResponse(t *testing.T) { func TestRawListenerResponse(t *testing.T) { var req, resp *TCPMessage - listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond) + listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP) defer listener.Close() reqPacket := buildPacket(true, 1, 1, []byte("GET / HTTP/1.1\r\n\r\n")) @@ -109,7 +109,7 @@ func TestRawListenerResponse(t *testing.T) { } func TestShort100Continue(t *testing.T) { - listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond) + listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP) defer listener.Close() reqPacket1 := buildPacket(true, 1, 1, []byte("POST / HTTP/1.1\r\nContent-Length: 2\r\nExpect: 100-continue\r\n\r\n")) @@ -129,7 +129,7 @@ func TestShort100Continue(t *testing.T) { // Response comes before Request func Test100ContinueWrongOrder(t *testing.T) { - listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond) + listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP) defer listener.Close() reqPacket1 := buildPacket(true, 1, 1, []byte("POST / HTTP/1.1\r\nContent-Length: 2\r\nExpect: 100-continue\r\n\r\n")) @@ -148,7 +148,7 @@ func Test100ContinueWrongOrder(t *testing.T) { } func TestAlt100ContinueHeaderOrder(t *testing.T) { - listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond) + listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP) defer listener.Close() reqPacket1 := buildPacket(true, 1, 1, []byte("POST / HTTP/1.1\r\nExpect: 100-continue\r\nContent-Length: 2\r\n\r\n")) @@ -309,7 +309,7 @@ func permutation(n int, list []*TCPPacket) []*TCPPacket { // Response comes before Request func TestRawListenerChunkedWrongOrder(t *testing.T) { - listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond) + listener := NewListener("", "0", EnginePcap, true, 10*time.Millisecond, ProtocolHTTP) defer listener.Close() reqPacket1 := buildPacket(true, 1, 1, []byte("POST / HTTP/1.1\r\nTransfer-Encoding: chunked\r\nExpect: 100-continue\r\n\r\n")) @@ -387,7 +387,7 @@ func getMessage() []*TCPPacket { // Response comes before Request func TestRawListenerBench(t *testing.T) { - l := NewListener("", "0", EnginePcap, true, 200*time.Millisecond) + l := NewListener("", "0", EnginePcap, true, 200*time.Millisecond, ProtocolHTTP) defer l.Close() // Should re-construct message from all possible combinations diff --git a/raw_socket_listener/tcp_message.go b/raw_socket_listener/tcp_message.go index 15837bd..ba9388d 100644 --- a/raw_socket_listener/tcp_message.go +++ b/raw_socket_listener/tcp_message.go @@ -14,6 +14,13 @@ import ( var _ = log.Println +type TCPProtocol uint8 + +const ( + ProtocolHTTP TCPProtocol = 0 + ProtocolBinary TCPProtocol = 1 +) + // TCPMessage ensure that all TCP packets for given request is received, and processed in right sequence // Its needed because all TCP message can be fragmented or re-transmitted // @@ -37,6 +44,8 @@ type TCPMessage struct { delChan chan *TCPMessage + protocol TCPProtocol + /* HTTP specific variables */ methodType httpMethodType bodyType httpBodyType @@ -48,8 +57,8 @@ type TCPMessage struct { } // NewTCPMessage pointer created from a Acknowledgment number and a channel of messages readuy to be deleted -func NewTCPMessage(Seq, Ack uint32, IsIncoming bool) (msg *TCPMessage) { - msg = &TCPMessage{Seq: Seq, Ack: Ack, IsIncoming: IsIncoming} +func NewTCPMessage(Seq, Ack uint32, IsIncoming bool, protocol TCPProtocol) (msg *TCPMessage) { + msg = &TCPMessage{Seq: Seq, Ack: Ack, IsIncoming: IsIncoming, protocol: protocol} msg.Start = time.Now() return @@ -95,48 +104,46 @@ func (t *TCPMessage) Size() (size int) { // AddPacket to the message and ensure packet uniqueness // TCP allows that packet can be re-send multiple times func (t *TCPMessage) AddPacket(packet *TCPPacket) { - packetFound := false - for _, pkt := range t.packets { if packet.Seq == pkt.Seq { - packetFound = true - break + return } } - if !packetFound { - // Packets not always captured in same Seq order, and sometimes we need to prepend - if len(t.packets) == 0 || packet.Seq > t.packets[len(t.packets)-1].Seq { - t.packets = append(t.packets, packet) - } else if packet.Seq < t.packets[0].Seq { - t.packets = append([]*TCPPacket{packet}, t.packets...) - t.Seq = packet.Seq // Message Seq should indicated starting seq - } else { // insert somewhere in the middle... - for i, p := range t.packets { - if packet.Seq < p.Seq { - t.packets = append(t.packets[:i], append([]*TCPPacket{packet}, t.packets[i:]...)...) - break - } + // Packets not always captured in same Seq order, and sometimes we need to prepend + if len(t.packets) == 0 || packet.Seq > t.packets[len(t.packets)-1].Seq { + t.packets = append(t.packets, packet) + } else if packet.Seq < t.packets[0].Seq { + t.packets = append([]*TCPPacket{packet}, t.packets...) + t.Seq = packet.Seq // Message Seq should indicated starting seq + } else { // insert somewhere in the middle... + for i, p := range t.packets { + if packet.Seq < p.Seq { + t.packets = append(t.packets[:i], append([]*TCPPacket{packet}, t.packets[i:]...)...) + break } } + } - if t.IsIncoming { - t.End = time.Now() - } else { - t.End = time.Now().Add(time.Millisecond) - } + if t.IsIncoming { + t.End = time.Now() + } else { + t.End = time.Now().Add(time.Millisecond) + } - if packet.OrigAck != 0 { - t.DataAck = packet.OrigAck - } + if packet.OrigAck != 0 { + t.DataAck = packet.OrigAck } t.checkSeqIntegrity() - t.updateHeadersPacket() - t.updateMethodType() - t.updateBodyType() - t.checkIfComplete() - t.check100Continue() + + if t.protocol == ProtocolHTTP { + t.updateHeadersPacket() + t.updateMethodType() + t.updateBodyType() + t.check100Continue() + t.checkIfComplete() + } } // Check if there is missing packet @@ -156,7 +163,7 @@ func (t *TCPMessage) checkSeqIntegrity() { nextSeq := p.Seq + uint32(len(p.Data)) if np.Seq != nextSeq { - if t.expectType == httpExpect100Continue { + if t.protocol == ProtocolHTTP && t.expectType == httpExpect100Continue { if np.Seq != nextSeq+22 { t.seqMissing = true return diff --git a/raw_socket_listener/tcp_message_test.go b/raw_socket_listener/tcp_message_test.go index 0af3bd9..618ea3f 100644 --- a/raw_socket_listener/tcp_message_test.go +++ b/raw_socket_listener/tcp_message_test.go @@ -36,7 +36,7 @@ func buildMessage(p *TCPPacket) *TCPMessage { isIncoming = true } - m := NewTCPMessage(p.Seq, p.Ack, isIncoming) + m := NewTCPMessage(p.Seq, p.Ack, isIncoming, ProtocolHTTP) m.AddPacket(p) return m diff --git a/settings.go b/settings.go index dd0675c..425c13d 100644 --- a/settings.go +++ b/settings.go @@ -49,13 +49,18 @@ type AppSettings struct { inputRAWEngine string inputRAWTrackResponse bool inputRAWRealIPHeader string + inputRAWProtocol string middleware string inputHTTP MultiOption - outputHTTP MultiOption + outputHTTP MultiOption outputHTTPConfig HTTPOutputConfig + + outputBinary MultiOption + outputBinaryConfig BinaryOutputConfig + modifierConfig HTTPModifierConfig } @@ -106,6 +111,8 @@ func init() { flag.StringVar(&Settings.inputRAWEngine, "input-raw-engine", "libpcap", "Intercept traffic using `libpcap` (default), and `raw_socket`") + flag.StringVar(&Settings.inputRAWProtocol, "input-raw-protocol", "http", "Specify application protocol of intercepted traffic. Possible values: http, binary") + flag.StringVar(&Settings.inputRAWRealIPHeader, "input-raw-realip-header", "", "If not blank, injects header with given name and real IP value to the request payload. Usually this header should be named: X-Real-IP") flag.StringVar(&Settings.middleware, "middleware", "", "Used for modifying traffic using external command") @@ -113,16 +120,27 @@ func init() { 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") + + /* outputHTTPConfig */ flag.IntVar(&Settings.outputHTTPConfig.BufferSize, "output-http-response-buffer", 0, "HTTP response buffer size, all data after this size will be discarded.") flag.IntVar(&Settings.outputHTTPConfig.workers, "output-http-workers", 0, "Gor uses dynamic worker scaling by default. Enter a number to run a set number of workers.") flag.IntVar(&Settings.outputHTTPConfig.redirectLimit, "output-http-redirects", 0, "Enable how often redirects should be followed.") flag.DurationVar(&Settings.outputHTTPConfig.Timeout, "output-http-timeout", 0, "Specify HTTP request/response timeout. By default 5s. Example: --output-http-timeout 30s") - flag.BoolVar(&Settings.outputHTTPConfig.stats, "output-http-stats", false, "Report http output queue stats to console every 5 seconds.") flag.BoolVar(&Settings.outputHTTPConfig.OriginalHost, "http-original-host", false, "Normally gor replaces the Host http header with the host supplied with --output-http. This option disables that behavior, preserving the original Host header.") flag.BoolVar(&Settings.outputHTTPConfig.Debug, "output-http-debug", false, "Enables http debug output.") - flag.StringVar(&Settings.outputHTTPConfig.elasticSearch, "output-http-elasticsearch", "", "Send request and response stats to ElasticSearch:\n\tgor --input-raw :8080 --output-http staging.com --output-http-elasticsearch 'es_host:api_port/index_name'") + /* outputHTTPConfig */ + + + flag.Var(&Settings.outputBinary, "output-binary", "Forwards incoming binary payloads to given address.\n\t# Redirect all incoming requests to staging.com address \n\tgor --input-raw :80 --input-raw-protocol binary --output-binary staging.com:80") + /* outputBinaryConfig */ + flag.IntVar(&Settings.outputBinaryConfig.BufferSize, "output-tcp-response-buffer", 0, "TCP response buffer size, all data after this size will be discarded.") + flag.IntVar(&Settings.outputBinaryConfig.workers, "output-binary-workers", 0, "Gor uses dynamic worker scaling by default. Enter a number to run a set number of workers.") + flag.DurationVar(&Settings.outputBinaryConfig.Timeout, "output-binary-timeout", 0, "Specify HTTP request/response timeout. By default 5s. Example: --output-binary-timeout 30s") + + flag.BoolVar(&Settings.outputBinaryConfig.Debug, "output-binary-debug", false, "Enables binary debug output.") + /* outputBinaryConfig */ flag.Var(&Settings.modifierConfig.headers, "http-set-header", "Inject additional headers to http reqest:\n\tgor --input-raw :8080 --output-http staging.com --http-set-header 'User-Agent: Gor'") flag.Var(&Settings.modifierConfig.headers, "output-http-header", "WARNING: `--output-http-header` DEPRECATED, use `--http-set-header` instead") diff --git a/tcp_client.go b/tcp_client.go new file mode 100644 index 0000000..677c2d4 --- /dev/null +++ b/tcp_client.go @@ -0,0 +1,190 @@ +package main + +import ( + "crypto/tls" + "io" + "log" + "net" + "runtime/debug" + "syscall" + "time" +) + +type TCPClientConfig struct { + Debug bool + ConnectionTimeout time.Duration + Timeout time.Duration + ResponseBufferSize int + Secure bool +} + +type TCPClient struct { + baseURL string + addr string + conn net.Conn + respBuf []byte + config *TCPClientConfig + redirectsCount int +} + +func NewTCPClient(addr string, config *TCPClientConfig) *TCPClient { + if config.Timeout.Nanoseconds() == 0 { + config.Timeout = 5 * time.Second + } + + config.ConnectionTimeout = config.Timeout + + if config.ResponseBufferSize == 0 { + config.ResponseBufferSize = 100 * 1024 // 100kb + } + + client := &TCPClient{config: config, addr: addr} + client.respBuf = make([]byte, config.ResponseBufferSize) + + return client +} + +func (c *TCPClient) Connect() (err error) { + c.Disconnect() + + c.conn, err = net.DialTimeout("tcp", c.addr, c.config.ConnectionTimeout) + + if c.config.Secure { + tlsConn := tls.Client(c.conn, &tls.Config{InsecureSkipVerify: true}) + + if err = tlsConn.Handshake(); err != nil { + return + } + + c.conn = tlsConn + } + + return +} + +func (c *TCPClient) Disconnect() { + if c.conn != nil { + c.conn.Close() + c.conn = nil + Debug("[TCPClient] Disconnected: ", c.baseURL) + } +} + +func (c *TCPClient) isAlive() bool { + one := make([]byte, 1) + + // Ready 1 byte from socket without timeout to check if it not closed + c.conn.SetReadDeadline(time.Now().Add(time.Millisecond)) + _, err := c.conn.Read(one) + + if err == nil { + return true + } else if err == io.EOF { + if c.config.Debug { + Debug("[TCPClient] connection closed, reconnecting") + } + return false + } else if err == syscall.EPIPE { + Debug("Detected broken pipe.", err) + return false + } + + return true +} + +func (c *TCPClient) Send(data []byte) (response []byte, err error) { + // Don't exit on panic + defer func() { + if r := recover(); r != nil { + Debug("[TCPClient]", r, string(data)) + + if _, ok := r.(error); !ok { + log.Println("[TCPClient] Failed to send request: ", string(data)) + log.Println("PANIC: pkg:", r, debug.Stack()) + } + } + }() + + if c.conn == nil || !c.isAlive() { + Debug("[TCPClient] Connecting:", c.baseURL) + if err = c.Connect(); err != nil { + log.Println("[TCPClient] Connection error:", err) + return + } + } + + timeout := time.Now().Add(c.config.Timeout) + + c.conn.SetWriteDeadline(timeout) + + if c.config.Debug { + Debug("[TCPClient] Sending:", string(data)) + } + + if _, err = c.conn.Write(data); err != nil { + Debug("[TCPClient] Write error:", err, c.baseURL) + return + } + + var readBytes, n int + var currentChunk []byte + timeout = time.Now().Add(c.config.Timeout) + + for { + c.conn.SetReadDeadline(timeout) + + if readBytes < len(c.respBuf) { + n, err = c.conn.Read(c.respBuf[readBytes:]) + readBytes += n + + if err != nil { + if err == io.EOF { + err = nil + } + break + } + } else { + if currentChunk == nil { + currentChunk = make([]byte, readChunkSize) + } + + n, err = c.conn.Read(currentChunk) + + if err == io.EOF { + break + } else if err != nil { + Debug("[TCPClient] Read the whole body error:", err, c.baseURL) + break + } + + readBytes += int(n) + } + + if readBytes >= maxResponseSize { + Debug("[TCPClient] Body is more than the max size", maxResponseSize, + c.baseURL) + break + } + + // For following chunks expect less timeout + timeout = time.Now().Add(c.config.Timeout / 5) + } + + if err != nil { + Debug("[TCPClient] Response read error", err, c.conn, readBytes) + return + } + + if readBytes > len(c.respBuf) { + readBytes = len(c.respBuf) + } + + payload := make([]byte, readBytes) + copy(payload, c.respBuf[:readBytes]) + + if c.config.Debug { + Debug("[TCPClient] Received:", string(payload)) + } + + return payload, err +}