From f0b96af420a636aeae36b9afc918d9645a417793 Mon Sep 17 00:00:00 2001
From: Leonid Bugaev <leonsbox@gmail.com>
Date: Sun, 23 Aug 2015 13:07:29 +0300
Subject: [PATCH] Add protection for malformed requests

---
 http_modifier.go |  4 ++++
 output_http.go   |  8 +++++++-
 proto/proto.go   | 15 +++++++++++++++
 3 files changed, 26 insertions(+), 1 deletion(-)

diff --git a/http_modifier.go b/http_modifier.go
index 14cdf85..fbca0c3 100644
--- a/http_modifier.go
+++ b/http_modifier.go
@@ -30,6 +30,10 @@ func NewHTTPModifier(config *HTTPModifierConfig) *HTTPModifier {
 }
 
 func (m *HTTPModifier) Rewrite(payload []byte) (response []byte) {
+	if !proto.IsHTTPPayload(payload) {
+		return payload
+	}
+
 	if len(m.config.methods) > 0 {
 		method := proto.Method(payload)
 
diff --git a/output_http.go b/output_http.go
index 0528e38..cafe637 100644
--- a/output_http.go
+++ b/output_http.go
@@ -5,6 +5,7 @@ import (
 	"log"
 	"sync/atomic"
 	"time"
+	"github.com/buger/gor/proto"
 )
 
 const initialDynamicWorkers = 10
@@ -186,8 +187,13 @@ func (o *HTTPOutput) sendRequest(client *HTTPClient, request []byte) {
 	meta := payloadMeta(request)
 	uuid := meta[1]
 
+	body := payloadBody(request)
+	if !proto.IsHTTPPayload(body) {
+		return
+	}
+
 	start := time.Now()
-	resp, err := client.Send(payloadBody(request))
+	resp, err := client.Send(body)
 	stop := time.Now()
 
 	if err != nil {
diff --git a/proto/proto.go b/proto/proto.go
index da53cc0..549cec3 100644
--- a/proto/proto.go
+++ b/proto/proto.go
@@ -208,3 +208,18 @@ func Method(payload []byte) []byte {
 func Status(payload []byte) []byte {
 	return Path(payload)
 }
+
+var httpMethods []string = []string{
+	"GET ", "OPTI", "HEAD", "POST", "PUT ", "DELE", "TRAC", "CONN",
+}
+
+func IsHTTPPayload(payload []byte) bool {
+	method := string(payload[0:4])
+
+	for _, m := range httpMethods {
+		if method == m {
+			return true
+		}
+	}
+	return false
+}