Merge pull request #6 from armosec/chunks

fix SplitSlice2Chunks size overflow

httputils/httphelper_test.go

@@ -103,7 +103,7 @@ func TestSplitSlice2Chunks(t *testing.T) {
 			Name: "CVE-2006-1611",
 		},
 		{
-			Name: "CVE-2017-18589",
+			Name: "This is a very long CVE name to test sizes and make sure we don't exceed max size",
 		},
 		{
 			Name: "CVE-2019-10743",
@@ -132,7 +132,8 @@ func TestSplitSlice2Chunks(t *testing.T) {
 	}
 
 	//split splice to chunks
-	chunksChan, totalTestTypes := SplitSlice2Chunks(testTypeSlice, 100, 10)
+	maxSize := 100
+	chunksChan, totalTestTypes := SplitSlice2Chunks(testTypeSlice, maxSize, 10)
 	testWg := sync.WaitGroup{}
 	var totalReceived, numOfChunks, maxChunkSize, minChunkSize, maxChunkLength, minChunkLength int
 	testWg.Add(1)
@@ -144,6 +145,9 @@ func TestSplitSlice2Chunks(t *testing.T) {
 			vSize := JSONSize(v)
 			vLen := len(v)
 			totalReceived += vLen
+			if vSize > maxSize {
+				t.Errorf("chunk size %d is more than expected max size %d", vSize, maxSize)
+			}
 			if maxChunkSize < vSize {
 				maxChunkSize = vSize
 			}
@@ -162,11 +166,11 @@ func TestSplitSlice2Chunks(t *testing.T) {
 	testWg.Wait()
 	//compare with expected
 	assert.Equal(t, totalTestTypes, totalReceived, "total elements received is not equal to number of element sent")
-	assert.Equal(t, 3, minChunkLength, "minChunkLength must be same as expected minChunkLength")
+	assert.Equal(t, 1, minChunkLength, "minChunkLength must be same as expected minChunkLength")
 	assert.Equal(t, 3, maxChunkLength, "maxChunkLength must be same as expected maxChunkLength")
-	assert.Equal(t, 77, minChunkSize, "minChunkSize must be same as expected minChunkSize")
-	assert.Equal(t, 89, maxChunkSize, "maxChunkSize must be same as expected maxChunkSize")
-	assert.Equal(t, 11, numOfChunks, "numOfChunks must be same as expected numOfChunks")
+	assert.Equal(t, 64, minChunkSize, "minChunkSize must be same as expected minChunkSize")
+	assert.Equal(t, 95, maxChunkSize, "maxChunkSize must be same as expected maxChunkSize")
+	assert.Equal(t, 12, numOfChunks, "numOfChunks must be same as expected numOfChunks")
 
 }
 

httputils/httphelpers.go

@@ -6,7 +6,6 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
-	"math"
 	"net/http"
 	"strings"
 	"sync"
@@ -87,13 +86,15 @@ func setHeaders(req *http.Request, headers map[string]string) {
 	}
 }
 
-// HTTPRespToString parses the body as string and checks the HTTP status code, it closes the body reader at the end
+// HttpRespToString parses the body as string and checks the HTTP status code, it closes the body reader at the end
 func HttpRespToString(resp *http.Response) (string, error) {
 	if resp == nil || resp.Body == nil {
 		return "", nil
 	}
 	strBuilder := strings.Builder{}
-	defer resp.Body.Close()
+	defer func(Body io.ReadCloser) {
+		_ = Body.Close()
+	}(resp.Body)
 	if resp.ContentLength > 0 {
 		strBuilder.Grow(int(resp.ContentLength))
 	}
@@ -176,36 +177,32 @@ func splitSlice2Chunks[T any](slice []T, maxSize int, chunks chan<- []T, wg *syn
 			return
 		}
 		//slice is bigger than max size
-		//calculate the average size + 5% of a single element T
-		avgTSize := int(math.Round(float64(jsonSize) * 1.05 / float64(len(slice))))
-		//calculate the average number of elements that will not exceed max size
-		avgSliceSize := maxSize / avgTSize
-		last := len(slice)
-		if avgSliceSize >= last {
-			avgSliceSize = last / 2
-		} else if avgSliceSize < 1 {
-			avgSliceSize = 1
-		}
-
-		//split the slice to slices of avgSliceSize size
-		startIndex := 0
-		for i := avgSliceSize; i < last; i += avgSliceSize {
-			splitSlice2Chunks(slice[startIndex:i], maxSize, chunks, wg)
-			startIndex = i
+		//split the slice to slices smaller than max size
+		index := 0
+		for i := 1; i < len(slice); i++ {
+			jsonSize = JSONSize(slice[index:i])
+			if jsonSize > maxSize {
+				//send the part of the slice that is smaller than max size
+				splitSlice2Chunks(slice[index:i-1], maxSize, chunks, wg)
+				index = i - 1
+			}
 		}
 		//send the last part of the slice
-		splitSlice2Chunks(slice[startIndex:last], maxSize, chunks, wg)
+		splitSlice2Chunks(slice[index:], maxSize, chunks, wg)
 	}(slice, maxSize, chunks, wg)
 }
 
-// jsonSize returns the size in bytes of the json encoding of i
+// JSONSize returns the size in bytes of the json encoding of i
 func JSONSize(i interface{}) int {
 	if i == nil {
 		return 0
 	}
 	counter := bytesCounter{}
 	enc := json.NewEncoder(&counter)
-	enc.Encode(i)
+	err := enc.Encode(i)
+	if err != nil {
+		return 0
+	}
 	return counter.count
 }