go fmt

elasticsearch.go

@@ -2,8 +2,8 @@ package main
 
 import (
 	"encoding/json"
-	"github.com/mattbaird/elastigo/lib"
 	"github.com/buger/gor/proto"
+	"github.com/mattbaird/elastigo/lib"
 	"log"
 	"regexp"
 	"time"

gor.go

@@ -132,4 +132,4 @@ func profileMEM(memprofile string) {
 			f.Close()
 		})
 	}
-}
+}

http_client.go

@@ -3,6 +3,7 @@ package main
 import (
 	"bytes"
 	"crypto/tls"
+	"encoding/base64"
 	"io"
 	"log"
 	"net"
@@ -44,6 +45,7 @@ type HTTPClient struct {
 	baseURL        string
 	scheme         string
 	host           string
+	auth           string
 	conn           net.Conn
 	respBuf        []byte
 	config         *HTTPClientConfig
@@ -79,6 +81,10 @@ func NewHTTPClient(baseURL string, config *HTTPClientConfig) *HTTPClient {
 	client.respBuf = make([]byte, config.ResponseBufferSize)
 	client.config = config
 
+	if u.User != nil {
+		client.auth = "Basic " + base64.StdEncoding.EncodeToString([]byte(u.User.String()))
+	}
+
 	return client
 }
 
@@ -164,6 +170,10 @@ func (c *HTTPClient) Send(data []byte) (response []byte, err error) {
 		data = proto.SetHost(data, []byte(c.baseURL), []byte(c.host))
 	}
 
+	if c.auth != "" {
+		data = proto.SetHeader(data, []byte("Authorization"), []byte(c.auth))
+	}
+
 	if c.config.Debug {
 		Debug("[HTTPClient] Sending:", string(data))
 	}

http_client_test.go

@@ -11,6 +11,7 @@ import (
 	"net/http/httptest"
 	"net/http/httputil"
 	_ "reflect"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -92,7 +93,7 @@ func TestHTTPClientResonseByClose(t *testing.T) {
 
 	payload := []byte("GET / HTTP/1.1\r\n\r\n")
 	ln, _ := net.Listen("tcp", ":0")
-	go func(){
+	go func() {
 		for {
 			conn, _ := ln.Accept()
 			buf := make([]byte, 4096)
@@ -356,6 +357,45 @@ func TestHTTPClientRedirectLimit(t *testing.T) {
 	wg.Wait()
 }
 
+func TestHTTPClientBasicAuth(t *testing.T) {
+	wg := new(sync.WaitGroup)
+	wg.Add(2)
+
+	GETPayload := []byte("GET / HTTP/1.1\r\n\r\n")
+
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		user, pass, _ := r.BasicAuth()
+
+		if user != "user" || pass != "pass" {
+			http.Error(w, "Unauthorized.", 401)
+			wg.Done()
+			return
+		}
+
+		wg.Done()
+	}))
+	defer server.Close()
+
+	client := NewHTTPClient(server.URL, &HTTPClientConfig{Debug: false})
+	resp, _ := client.Send(GETPayload)
+	client.Disconnect()
+
+	if !bytes.Equal(proto.Status(resp), []byte("401")) {
+		t.Error("Should return unauthorized error", string(resp))
+	}
+
+	authUrl := strings.Replace(server.URL, "http://", "http://user:pass@", -1)
+	client = NewHTTPClient(authUrl, &HTTPClientConfig{Debug: false})
+	resp, _ = client.Send(GETPayload)
+	client.Disconnect()
+
+	if !bytes.Equal(proto.Status(resp), []byte("200")) {
+		t.Error("Should return proper response", string(resp))
+	}
+
+	wg.Wait()
+}
+
 func TestHTTPClientHandleHTTP10(t *testing.T) {
 	wg := new(sync.WaitGroup)
 

input_raw_test.go

@@ -92,10 +92,10 @@ func TestRAWInputNoKeepAlive(t *testing.T) {
 		t.Fatal(err)
 	}
 	origin := &http.Server{
-		Handler:      http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-				w.Write([]byte("a"))
-				w.Write([]byte("b"))
-			}),
+		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.Write([]byte("a"))
+			w.Write([]byte("b"))
+		}),
 		ReadTimeout:  10 * time.Second,
 		WriteTimeout: 10 * time.Second,
 	}

output_http.go

@@ -179,7 +179,7 @@ func (o *HTTPOutput) Read(data []byte) (int, error) {
 
 	Debug("[OUTPUT-HTTP] Received response:", string(resp.payload))
 
-	header := payloadHeader(ReplayedResponsePayload, resp.uuid, resp.startedAt, resp.roundTripTime)
+	header := payloadHeader(ReplayedResponsePayload, resp.uuid, resp.roundTripTime, resp.startedAt)
 	copy(data[0:len(header)], header)
 	copy(data[len(header):], resp.payload)
 

output_null.go

@@ -6,13 +6,13 @@ type NullOutput struct {
 
 // NullOutput constructor for NullOutput
 func NewNullOutput() (o *NullOutput) {
-    return new(NullOutput)
+	return new(NullOutput)
 }
 
 func (o *NullOutput) Write(data []byte) (int, error) {
-    return len(data), nil
+	return len(data), nil
 }
 
 func (o *NullOutput) String() string {
-    return "Null Output"
+	return "Null Output"
 }

raw_socket_listener/listener.go

@@ -115,8 +115,6 @@ func NewListener(addr string, port string, engine int, trackResponse bool, expir
 	// Special case for testing
 	if l.port != 0 {
 		switch engine {
-		case EngineRawSocket:
-			go l.readRAWSocket()
 		case EnginePcap:
 			go l.readPcap()
 		case EnginePcapFile:
@@ -559,7 +557,14 @@ func (t *Listener) readPcapFile() {
 			if tcpLayer := packet.Layer(layers.LayerTypeTCP); tcpLayer != nil {
 				tcp, _ := tcpLayer.(*layers.TCP)
 				data = append(tcp.LayerContents(), tcp.LayerPayload()...)
-				copy(data[2:4], []byte{0, 1})
+
+				if tcp.SrcPort >= 32768 && tcp.SrcPort <= 61000 {
+					copy(data[0:2], []byte{0, 0})
+					copy(data[2:4], []byte{0, 1})
+				} else {
+					copy(data[0:2], []byte{0, 1})
+					copy(data[2:4], []byte{0, 0})
+				}
 			} else {
 				continue
 			}
@@ -576,10 +581,11 @@ func (t *Listener) readPcapFile() {
 			}
 
 			dataOffset := (data[12] & 0xF0) >> 4
+			isFIN := data[13]&0x01 != 0
 
 			// We need only packets with data inside
 			// Check that the buffer is larger than the size of the TCP header
-			if len(data) <= int(dataOffset*4) {
+			if len(data) <= int(dataOffset*4) && !isFIN {
 				continue
 			}
 
@@ -663,8 +669,6 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) {
 		}
 	}()
 
-	// log.Println("Processing packet:", packet.Ack, packet.Seq, packet.ID)
-
 	var message *TCPMessage
 
 	isIncoming := packet.DestPort == t.port
@@ -693,6 +697,14 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) {
 		packet.UpdateAck(parentAck)
 	}
 
+	if isIncoming && packet.IsFIN {
+		if ma, ok := t.respAliases[packet.Seq]; ok {
+			if ma.packets[0].SrcPort == packet.SrcPort {
+				packet.UpdateAck(ma.Ack)
+			}
+		}
+	}
+
 	if alias, ok := t.ackAliases[packet.Ack]; ok {
 		packet.UpdateAck(alias)
 	}
@@ -764,8 +776,11 @@ func (t *Listener) processTCPPacket(packet *TCPPacket) {
 
 	// If message contains only single packet immediately dispatch it
 	if message.complete {
+		// log.Println("COMPLETE!", isIncoming, message)
 		if isIncoming {
 			if t.trackResponse {
+				// log.Println("Found response!", message.ResponseID, t.messages)
+
 				if resp, ok := t.messages[message.ResponseID]; ok {
 					if resp.complete {
 						t.dispatchMessage(resp)

raw_socket_listener/tcp_message.go

@@ -352,6 +352,26 @@ func (t *TCPMessage) updateBodyType() {
 		return
 	}
 
+	var lengthB, encB, connB []byte
+
+	proto.ParseHeaders(t.packetsData(), func(header, value []byte)bool{
+		if proto.HeadersEqual(header, []byte("Content-Length")) {
+			lengthB = value
+			return false
+		}
+
+		if proto.HeadersEqual(header, []byte("Transfer-Encoding")) {
+			encB = value
+			return false
+		}
+
+		if proto.HeadersEqual(header, []byte("Connection")) {
+			connB = value
+		}
+
+		return true
+	})
+
 	switch t.methodType {
 	case httpMethodNotFound:
 		return
@@ -359,27 +379,6 @@ func (t *TCPMessage) updateBodyType() {
 		t.bodyType = httpBodyEmpty
 		return
 	case httpMethodWithBody:
-		var lengthB, encB, connB []byte
-
-		proto.ParseHeaders(t.packetsData(), func(header, value []byte)bool{
-			if proto.HeadersEqual(header, []byte("Content-Length")) {
-				lengthB = value
-				return false
-			}
-
-			if proto.HeadersEqual(header, []byte("Transfer-Encoding")) {
-				encB = value
-				return false
-			}
-
-			if proto.HeadersEqual(header, []byte("Connection")) {
-				connB = value
-				return false
-			}
-
-			return true
-		})
-
 		if len(lengthB) > 0 {
 			t.contentLength, _ = strconv.Atoi(string(lengthB))
 
@@ -461,6 +460,10 @@ func (t *TCPMessage) setAssocMessage(m *TCPMessage) {
 // UpdateResponseAck should be called after packet is added
 func (t *TCPMessage) UpdateResponseAck() uint32 {
 	lastPacket := t.packets[len(t.packets)-1]
+	if lastPacket.IsFIN && len(t.packets) > 1 {
+		lastPacket = t.packets[len(t.packets)-2]
+	}
+
 	respAck := lastPacket.Seq + uint32(len(lastPacket.Data))
 
 	if t.ResponseAck != respAck {