Add gzip compression support

Add go-modules support
eventbrite · Sep 8, 2021 · e352744 · e352744
1 parent 10df423
commit e352744
Show file tree

Hide file tree

Showing 7 changed files with 219 additions and 121 deletions.
diff --git a/client.go b/client.go
@@ -2,6 +2,7 @@ package hec
 
 import (
 	"bytes"
+	"compress/gzip"
 	"context"
 	"encoding/json"
 	"fmt"
@@ -52,6 +53,9 @@ type Client struct {
 
 	// Mutex to allow threadsafe acknowledgement checking
 	ackMux sync.Mutex
+
+	// Compression type, "" and "gzip" are supported
+	compression string
 }
 
 func NewClient(serverURL string, token string) HEC {
@@ -88,6 +92,10 @@ func (hec *Client) SetMaxContentLength(size int) {
 	hec.maxLength = size
 }
 
+func (hec *Client) SetCompression(compression string) {
+	hec.compression = compression
+}
+
 func (hec *Client) WriteEventWithContext(ctx context.Context, event *Event) error {
 	if event.empty() {
 		return nil // skip empty events
@@ -314,7 +322,21 @@ func (res *Response) String() string {
 func (hec *Client) makeRequest(ctx context.Context, endpoint string, data []byte) (*Response, error) {
 	retries := 0
 RETRY:
-	req, err := http.NewRequest(http.MethodPost, hec.serverURL+endpoint, bytes.NewReader(data))
+	var reader io.Reader
+	if hec.compression == "gzip" {
+		var buffer bytes.Buffer
+		gzipWriter := gzip.NewWriter(&buffer)
+		_, err := gzipWriter.Write(data)
+		gzipWriter.Close()
+		if err != nil {
+			return nil, err
+		}
+		reader = &buffer
+	} else {
+		reader = bytes.NewReader(data)
+	}
+
+	req, err := http.NewRequest(http.MethodPost, hec.serverURL+endpoint, reader)
 	if err != nil {
 		return nil, err
 	}
@@ -323,6 +345,9 @@ RETRY:
 		req.Header.Set("Connection", "keep-alive")
 	}
 	req.Header.Set("Authorization", "Splunk "+hec.token)
+	if hec.compression == "gzip" {
+		req.Header.Set("Content-Encoding", "gzip")
+	}
 	res, err := hec.httpClient.Do(req)
 	if err != nil {
 		return nil, err

diff --git a/client_test.go b/client_test.go
@@ -1,6 +1,7 @@
 package hec
 
 import (
+	"compress/gzip"
 	"crypto/tls"
 	"encoding/json"
 	"net/http"
@@ -24,11 +25,23 @@ var testHttpClient *http.Client = &http.Client{
 	Timeout: 100 * time.Millisecond,
 }
 
-func jsonEndpoint(t *testing.T) http.Handler {
+func jsonEndpoint(t *testing.T, compression string) http.Handler {
 	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		failed := false
 		input := make(map[string]interface{})
-		j := json.NewDecoder(r.Body)
+		content := r.Body
+		if compression == "gzip" {
+			var err error
+			content, err = gzip.NewReader(r.Body)
+			if err != nil {
+				t.Errorf("Unexpected error in gzip: %v", err)
+			}
+			header := r.Header.Get("Content-Encoding")
+			if header != "gzip" {
+				t.Errorf("Content-Encoding header wasn't sent for gzip")
+			}
+		}
+		j := json.NewDecoder(content)
 		err := j.Decode(&input)
 		if err != nil {
 			t.Errorf("Decoding JSON: %v", err)
@@ -66,20 +79,23 @@ func jsonEndpoint(t *testing.T) http.Handler {
 }
 
 func TestHEC_WriteEvent(t *testing.T) {
-	event := &Event{
-		Index:      String("main"),
-		Source:     String("test-hec-raw"),
-		SourceType: String("manual"),
-		Host:       String("localhost"),
-		Time:       String("1485237827.123"),
-		Event:      "hello, world",
-	}
+	for _, compression := range []string{"", "gzip"} {
+		event := &Event{
+			Index:      String("main"),
+			Source:     String("test-hec-raw"),
+			SourceType: String("manual"),
+			Host:       String("localhost"),
+			Time:       String("1485237827.123"),
+			Event:      "hello, world",
+		}
 
-	ts := httptest.NewServer(jsonEndpoint(t))
-	c := NewClient(ts.URL, testSplunkToken)
-	c.SetHTTPClient(testHttpClient)
-	err := c.WriteEvent(event)
-	assert.NoError(t, err)
+		ts := httptest.NewServer(jsonEndpoint(t, compression))
+		c := NewClient(ts.URL, testSplunkToken)
+		c.SetHTTPClient(testHttpClient)
+		c.SetCompression(compression)
+		err := c.WriteEvent(event)
+		assert.NoError(t, err)
+	}
 }
 
 func TestHEC_WriteEventServerFailure(t *testing.T) {
@@ -103,23 +119,26 @@ func TestHEC_WriteEventServerFailure(t *testing.T) {
 }
 
 func TestHEC_WriteObjectEvent(t *testing.T) {
-	event := &Event{
-		Index:      String("main"),
-		Source:     String("test-hec-raw"),
-		SourceType: String("manual"),
-		Host:       String("localhost"),
-		Time:       String("1485237827.123"),
-		Event: map[string]interface{}{
-			"str":  "hello",
-			"time": time.Now(),
-		},
-	}
+	for _, compression := range []string{"", "gzip"} {
+		event := &Event{
+			Index:      String("main"),
+			Source:     String("test-hec-raw"),
+			SourceType: String("manual"),
+			Host:       String("localhost"),
+			Time:       String("1485237827.123"),
+			Event: map[string]interface{}{
+				"str":  "hello",
+				"time": time.Now(),
+			},
+		}
 
-	ts := httptest.NewServer(jsonEndpoint(t))
-	c := NewClient(ts.URL, testSplunkToken)
-	c.SetHTTPClient(testHttpClient)
-	err := c.WriteEvent(event)
-	assert.NoError(t, err)
+		ts := httptest.NewServer(jsonEndpoint(t, compression))
+		c := NewClient(ts.URL, testSplunkToken)
+		c.SetHTTPClient(testHttpClient)
+		c.SetCompression(compression)
+		err := c.WriteEvent(event)
+		assert.NoError(t, err)
+	}
 }
 
 func TestHEC_WriteLongEvent(t *testing.T) {
@@ -132,7 +151,7 @@ func TestHEC_WriteLongEvent(t *testing.T) {
 		Event:      "hello, world",
 	}
 
-	ts := httptest.NewServer(jsonEndpoint(t))
+	ts := httptest.NewServer(jsonEndpoint(t, ""))
 	c := NewClient(ts.URL, testSplunkToken)
 
 	c.SetHTTPClient(testHttpClient)
@@ -143,62 +162,74 @@ func TestHEC_WriteLongEvent(t *testing.T) {
 }
 
 func TestHEC_WriteEventBatch(t *testing.T) {
-	events := []*Event{
-		{Event: "event one"},
-		{Event: "event two"},
-	}
+	for _, compression := range []string{"", "gzip"} {
+		events := []*Event{
+			{Event: "event one"},
+			{Event: "event two"},
+		}
 
-	ts := httptest.NewServer(jsonEndpoint(t))
-	c := NewClient(ts.URL, testSplunkToken)
+		ts := httptest.NewServer(jsonEndpoint(t, compression))
+		c := NewClient(ts.URL, testSplunkToken)
 
-	c.SetHTTPClient(testHttpClient)
-	err := c.WriteBatch(events)
-	assert.NoError(t, err)
+		c.SetHTTPClient(testHttpClient)
+		c.SetCompression(compression)
+		err := c.WriteBatch(events)
+		assert.NoError(t, err)
+	}
 }
 
 func TestHEC_WriteLongEventBatch(t *testing.T) {
-	events := []*Event{
-		{Event: "event one"},
-		{Event: "event two"},
-	}
+	for _, compression := range []string{"", "gzip"} {
+		events := []*Event{
+			{Event: "event one"},
+			{Event: "event two"},
+		}
 
-	ts := httptest.NewServer(jsonEndpoint(t))
-	c := NewClient(ts.URL, testSplunkToken)
-	c.SetHTTPClient(testHttpClient)
-	c.SetMaxContentLength(25)
-	err := c.WriteBatch(events)
-	assert.NoError(t, err)
+		ts := httptest.NewServer(jsonEndpoint(t, compression))
+		c := NewClient(ts.URL, testSplunkToken)
+		c.SetHTTPClient(testHttpClient)
+		c.SetMaxContentLength(25)
+		c.SetCompression(compression)
+		err := c.WriteBatch(events)
+		assert.NoError(t, err)
+	}
 }
 
 func TestHEC_WriteEventRaw(t *testing.T) {
-	events := `2017-01-24T06:07:10.488Z Raw event one
+	for _, compression := range []string{"", "gzip"} {
+		events := `2017-01-24T06:07:10.488Z Raw event one
 2017-01-24T06:07:12.434Z Raw event two`
-	metadata := EventMetadata{
-		Source: String("test-hec-raw"),
+		metadata := EventMetadata{
+			Source: String("test-hec-raw"),
+		}
+		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.Write([]byte(`{"text":"Success","code":0}`))
+		}))
+		c := NewClient(ts.URL, testSplunkToken)
+		c.SetHTTPClient(testHttpClient)
+		c.SetCompression(compression)
+		err := c.WriteRaw(strings.NewReader(events), &metadata)
+		assert.NoError(t, err)
 	}
-	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		w.Write([]byte(`{"text":"Success","code":0}`))
-	}))
-	c := NewClient(ts.URL, testSplunkToken)
-	c.SetHTTPClient(testHttpClient)
-	err := c.WriteRaw(strings.NewReader(events), &metadata)
-	assert.NoError(t, err)
 }
 
 func TestHEC_WriteLongEventRaw(t *testing.T) {
-	events := `2017-01-24T06:07:10.488Z Raw event one
+	for _, compression := range []string{"", "gzip"} {
+		events := `2017-01-24T06:07:10.488Z Raw event one
 2017-01-24T06:07:12.434Z Raw event two`
-	metadata := EventMetadata{
-		Source: String("test-hec-raw"),
+		metadata := EventMetadata{
+			Source: String("test-hec-raw"),
+		}
+		ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.Write([]byte(`{"text":"Success","code":0}`))
+		}))
+		c := NewClient(ts.URL, testSplunkToken)
+		c.SetMaxContentLength(40)
+		c.SetHTTPClient(testHttpClient)
+		c.SetCompression(compression)
+		err := c.WriteRaw(strings.NewReader(events), &metadata)
+		assert.NoError(t, err)
 	}
-	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		w.Write([]byte(`{"text":"Success","code":0}`))
-	}))
-	c := NewClient(ts.URL, testSplunkToken)
-	c.SetMaxContentLength(40)
-	c.SetHTTPClient(testHttpClient)
-	err := c.WriteRaw(strings.NewReader(events), &metadata)
-	assert.NoError(t, err)
 }
 
 func TestHEC_WriteRawFailure(t *testing.T) {

diff --git a/cluster.go b/cluster.go
@@ -78,6 +78,14 @@ func (c *Cluster) SetMaxContentLength(size int) {
 	c.mtx.Unlock()
 }
 
+func (c *Cluster) SetCompression(compression string) {
+	c.mtx.Lock()
+	for _, client := range c.clients {
+		client.SetCompression(compression)
+	}
+	c.mtx.Unlock()
+}
+
 func (c *Cluster) WriteEvent(event *Event) error {
 	return c.retry(func(client *Client) error {
 		return client.WriteEvent(event)