Merge pull request #46 from BloodHoundAD/gzip-compression

feat: add gzip compression to BHE ingest
SpecterOps · Sep 14, 2023 · d2ff1b6 · d2ff1b6
2 parents c8ea1a5 + e556b5d
commit d2ff1b6
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 9 deletions.
diff --git a/cmd/start.go b/cmd/start.go
@@ -18,6 +18,8 @@
 package cmd
 
 import (
+	"bytes"
+	"compress/gzip"
 	"context"
 	"encoding/json"
 	"errors"
@@ -182,20 +184,32 @@ func ingest(ctx context.Context, bheUrl url.URL, bheClient *http.Client, in <-ch
 	)
 
 	for data := range pipeline.OrDone(ctx.Done(), in) {
-		body := models.IngestRequest{
+		var (
+			body bytes.Buffer
+			gw   = gzip.NewWriter(&body)
+		)
+
+		ingestData := models.IngestRequest{
 			Meta: models.Meta{
 				Type: "azure",
 			},
 			Data: data,
 		}
 
-		headers := make(map[string]string)
-		headers["Prefer"] = "wait=60"
+		err := json.NewEncoder(gw).Encode(ingestData)
+		if err != nil {
+			log.Error(err, unrecoverableErrMsg)
+		}
+		gw.Close()
 
-		if req, err := rest.NewRequest(ctx, "POST", endpoint, body, nil, headers); err != nil {
+		if req, err := http.NewRequestWithContext(ctx, "POST", endpoint.String(), &body); err != nil {
 			log.Error(err, unrecoverableErrMsg)
 			return true
 		} else {
+			req.Header.Set("User-Agent", constants.UserAgent())
+			req.Header.Set("Accept", "application/json")
+			req.Header.Set("Prefer", "wait=60")
+			req.Header.Set("Content-Encoding", "gzip")
 			for retry := 0; retry < maxRetries; retry++ {
 				//No retries on regular err cases, only on HTTP 504 Gateway Timeout and HTTP 503 Service Unavailable
 				if response, err := bheClient.Do(req); err != nil {

diff --git a/cmd/utils.go b/cmd/utils.go
@@ -26,8 +26,8 @@ import (
 	"crypto/tls"
 	"encoding/base64"
 	"fmt"
+	"io"
 	"io/fs"
-	"io/ioutil"
 	"net"
 	"net/http"
 	"net/url"
@@ -219,7 +219,7 @@ func newAzureClient() (client.AzureClient, error) {
 	)
 
 	if file, ok := certFile.(string); ok && file != "" {
-		if content, err := ioutil.ReadFile(certFile.(string)); err != nil {
+		if content, err := os.ReadFile(certFile.(string)); err != nil {
 			return nil, fmt.Errorf("unable to read provided certificate: %w", err)
 		} else {
 			clientCert = string(content)
@@ -228,7 +228,7 @@ func newAzureClient() (client.AzureClient, error) {
 	}
 
 	if file, ok := keyFile.(string); ok && file != "" {
-		if content, err := ioutil.ReadFile(keyFile.(string)); err != nil {
+		if content, err := os.ReadFile(keyFile.(string)); err != nil {
 			return nil, fmt.Errorf("unable to read provided key file: %w", err)
 		} else {
 			clientKey = string(content)
@@ -304,8 +304,8 @@ func (s signingTransport) RoundTrip(req *http.Request) (*http.Response, error) {
 		if contentLength, err := body.ReadFrom(req.Body); err != nil {
 			return nil, err
 		} else if contentLength != 0 {
-			req.Body = ioutil.NopCloser(bytes.NewReader(body.Bytes()))
-			clone.Body = ioutil.NopCloser(bytes.NewReader(body.Bytes()))
+			req.Body = io.NopCloser(bytes.NewReader(body.Bytes()))
+			clone.Body = io.NopCloser(bytes.NewReader(body.Bytes()))
 		}
 	}
 	if _, err := digester.Write(body.Bytes()); err != nil {