diff --git a/output_kafka.go b/output_kafka.go index 4ebfc4c..23ba31a 100644 --- a/output_kafka.go +++ b/output_kafka.go @@ -3,6 +3,7 @@ package main import ( "encoding/json" "github.com/Shopify/sarama" + "github.com/Shopify/sarama/mocks" "github.com/buger/gor/proto" "io" "log" @@ -13,8 +14,10 @@ import ( // KafkaConfig should contains required information to // build producers. type KafkaConfig struct { - host string - topic string + host string + topic string + producer sarama.AsyncProducer + useJSON bool } // KafkaOutput should make producer client. @@ -27,6 +30,9 @@ type KafkaOutput struct { // passed as Json to Apache Kafka. type KafkaMessage struct { ReqURL string `json:"Req_URL"` + ReqType string `json:"Req_Type"` + ReqID string `json:"Req_ID"` + ReqTs string `json:"Req_Ts"` ReqMethod string `json:"Req_Method"` ReqBody string `json:"Req_Body,omitempty"` ReqHeaders map[string]string `json:"Req_Headers,omitempty"` @@ -38,15 +44,23 @@ const KafkaOutputFrequency = 500 // NewKafkaOutput creates instance of kafka producer client. func NewKafkaOutput(address string, config *KafkaConfig) io.Writer { c := sarama.NewConfig() - c.Producer.RequiredAcks = sarama.WaitForLocal - c.Producer.Compression = sarama.CompressionSnappy - c.Producer.Flush.Frequency = KafkaOutputFrequency * time.Millisecond - brokerList := strings.Split(config.host, ",") + var producer sarama.AsyncProducer - producer, err := sarama.NewAsyncProducer(brokerList, c) - if err != nil { - log.Fatalln("Failed to start Sarama(Kafka) producer:", err) + if config.producer.(*mocks.AsyncProducer) != nil { + producer = config.producer + } else { + c.Producer.RequiredAcks = sarama.WaitForLocal + c.Producer.Compression = sarama.CompressionSnappy + c.Producer.Flush.Frequency = KafkaOutputFrequency * time.Millisecond + + brokerList := strings.Split(config.host, ",") + + var err error + producer, err = sarama.NewAsyncProducer(brokerList, c) + if err != nil { + log.Fatalln("Failed to start Sarama(Kafka) producer:", err) + } } o := &KafkaOutput{ @@ -70,22 +84,32 @@ func (o *KafkaOutput) ErrorHandler() { } func (o *KafkaOutput) Write(data []byte) (n int, err error) { - headers := make(map[string]string) - proto.ParseHeaders([][]byte{data}, func(header []byte, value []byte) bool { - headers[string(header)] = string(value) - return true - }) + var message sarama.StringEncoder - req := payloadBody(data) + if !o.config.useJSON { + message = sarama.StringEncoder(data) + } else { + headers := make(map[string]string) + proto.ParseHeaders([][]byte{data}, func(header []byte, value []byte) bool { + headers[string(header)] = string(value) + return true + }) - kafkaMessage := KafkaMessage{ - ReqURL: string(proto.Path(req)), - ReqMethod: string(proto.Method(req)), - ReqBody: string(proto.Body(req)), - ReqHeaders: headers, + meta := payloadMeta(data) + req := payloadBody(data) + + kafkaMessage := KafkaMessage{ + ReqURL: string(proto.Path(req)), + ReqType: string(meta[0]), + ReqID: string(meta[1]), + ReqTs: string(meta[2]), + ReqMethod: string(proto.Method(req)), + ReqBody: string(proto.Body(req)), + ReqHeaders: headers, + } + jsonMessage, _ := json.Marshal(&kafkaMessage) + message = sarama.StringEncoder(jsonMessage) } - jsonMessage, _ := json.Marshal(&kafkaMessage) - message := sarama.StringEncoder(jsonMessage) o.producer.Input() <- &sarama.ProducerMessage{ Topic: o.config.topic, diff --git a/output_kafka_test.go b/output_kafka_test.go new file mode 100644 index 0000000..cf9efe0 --- /dev/null +++ b/output_kafka_test.go @@ -0,0 +1,53 @@ +package main + +import ( + "github.com/Shopify/sarama" + "github.com/Shopify/sarama/mocks" + "testing" +) + +func TestOutputKafkaRAW(t *testing.T) { + config := sarama.NewConfig() + config.Producer.Return.Successes = true + producer := mocks.NewAsyncProducer(t, config) + producer.ExpectInputAndSucceed() + + output := NewKafkaOutput("", &KafkaConfig{ + producer: producer, + topic: "test", + useJSON: false, + }) + + output.Write([]byte("1 2 3\nGET / HTTP1.1\r\nHeader: 1\r\n\r\n")) + + resp := <-producer.Successes() + + data, _ := resp.Value.Encode() + + if string(data) != "1 2 3\nGET / HTTP1.1\r\nHeader: 1\r\n\r\n" { + t.Error("Message not properly encoded: ", string(data)) + } +} + +func TestOutputKafkaJSON(t *testing.T) { + config := sarama.NewConfig() + config.Producer.Return.Successes = true + producer := mocks.NewAsyncProducer(t, config) + producer.ExpectInputAndSucceed() + + output := NewKafkaOutput("", &KafkaConfig{ + producer: producer, + topic: "test", + useJSON: true, + }) + + output.Write([]byte("1 2 3\nGET / HTTP1.1\r\nHeader: 1\r\n\r\n")) + + resp := <-producer.Successes() + + data, _ := resp.Value.Encode() + + if string(data) != `{"Req_URL":"/","Req_Type":"1","Req_ID":"2","Req_Ts":"3","Req_Method":"GET","Req_Headers":{"Header":"1"}}` { + t.Error("Message not properly encoded: ", string(data)) + } +} diff --git a/settings.go b/settings.go index f012867..d57deb6 100644 --- a/settings.go +++ b/settings.go @@ -130,6 +130,7 @@ func init() { flag.StringVar(&Settings.outputKafkaConfig.host, "output-kafka-host", "", "Send request and response stats to Kafka:\n\tgor --input-raw :8080 --output-kafka-host '192.168.0.1:9092,192.168.0.2:9092'") flag.StringVar(&Settings.outputKafkaConfig.topic, "output-kafka-topic", "", "Send request and response stats to Kafka:\n\tgor --input-raw :8080 --output-kafka-topic 'kafka-log'") + flag.BoolVar(&Settings.outputKafkaConfig.useJSON, "output-kafka-json-format", false, "If turned on, it will serialize messages from GoReplay text format to JSON.") 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")