Merge pull request #93 from msaf1980/fix/tags_buffer_overflow

Fix uploader tags buffer overflow

Author: Felixoid

helper/RowBinary/write_buffer.go

@@ -7,6 +7,12 @@ import (
 	"sync"
 )
 
+const (
+	SIZE_INT64 = 8
+	SIZE_INT16 = 2
+	SIZE_INT32 = 4
+)
+
 var writeBufferPool = sync.Pool{
 	New: func() interface{} {
 		return &WriteBuffer{}
@@ -59,6 +65,14 @@ func (wb *WriteBuffer) Reset() *WriteBuffer {
 	return wb
 }
 
+func (wb *WriteBuffer) Len() int {
+	return wb.Used
+}
+
+func (wb *WriteBuffer) FreeSize() int {
+	return len(wb.Body) - wb.Used
+}
+
 func (wb *WriteBuffer) Empty() bool {
 	return wb.Used == 0
 }

uploader/tagged.go

@@ -10,6 +10,7 @@ import (
 	"time"
 
 	"github.com/lomik/carbon-clickhouse/helper/RowBinary"
+	"go.uber.org/zap"
 )
 
 type Tagged struct {
@@ -20,6 +21,8 @@ type Tagged struct {
 var _ Uploader = &Tagged{}
 var _ UploaderWithReset = &Tagged{}
 
+var errBufOverflow = fmt.Errorf("output buffer overflow")
+
 func NewTagged(base *Base) *Tagged {
 	u := &Tagged{}
 	u.cached = newCached(base)
@@ -49,6 +52,74 @@ func urlParse(rawurl string) (*url.URL, error) {
 	return m, err
 }
 
+func (u *Tagged) parseName(name string, days uint16,
+	// reusable buffers
+	tag1 []string, wb *RowBinary.WriteBuffer, tagsBuf *RowBinary.WriteBuffer) error {
+
+	m, err := urlParse(name)
+	if err != nil {
+		return err
+	}
+
+	version := uint32(time.Now().Unix())
+
+	wb.Reset()
+	tagsBuf.Reset()
+	tag1 = tag1[:0]
+
+	t := fmt.Sprintf("__name__=%s", m.Path)
+	tag1 = append(tag1, t)
+	tagsBuf.WriteString(t)
+
+	// calc size for prevent buffer overflow
+	sizeTags := RowBinary.SIZE_INT16 /* days */ +
+		RowBinary.SIZE_INT64 + len(m.Path) +
+		RowBinary.SIZE_INT64 + len(name) +
+		RowBinary.SIZE_INT64 + //  tagsBuf.Len() not known at this step
+		RowBinary.SIZE_INT16 //version
+
+	// don't upload any other tag but __name__
+	// if either main metric (m.Path) or each metric (*) is ignored
+	ignoreAllButName := u.ignoredMetrics[m.Path] || u.ignoredMetrics["*"]
+	tagsWritten := 1
+	for k, v := range m.Query() {
+		t := fmt.Sprintf("%s=%s", k, v[0])
+
+		sizeTags += RowBinary.SIZE_INT16 /* days */ +
+			RowBinary.SIZE_INT64 + len(t) +
+			RowBinary.SIZE_INT64 + len(name) +
+			RowBinary.SIZE_INT64 + // tagsBuf.Len() not known at this step
+			RowBinary.SIZE_INT16 //version
+
+		if sizeTags >= wb.FreeSize() {
+			return errBufOverflow
+		}
+
+		tagsBuf.WriteString(t)
+		tagsWritten++
+
+		if !ignoreAllButName {
+			tag1 = append(tag1, t)
+		}
+	}
+
+	sizeTags += len(tag1) * tagsBuf.Len()
+	if sizeTags >= wb.FreeSize() {
+		return errBufOverflow
+	}
+
+	for i := 0; i < len(tag1); i++ {
+		wb.WriteUint16(days)
+		wb.WriteString(tag1[i])
+		wb.WriteString(name)
+		wb.WriteUVarint(uint64(tagsWritten))
+		wb.Write(tagsBuf.Bytes())
+		wb.WriteUint32(version)
+	}
+
+	return nil
+}
+
 func (u *Tagged) parseFile(filename string, out io.Writer) (uint64, map[string]bool, error) {
 	var reader *RowBinary.Reader
 	var err error
@@ -60,8 +131,6 @@ func (u *Tagged) parseFile(filename string, out io.Writer) (uint64, map[string]b
 	}
 	defer reader.Close()
 
-	version := uint32(time.Now().Unix())
-
 	newTagged := make(map[string]bool)
 
 	wb := RowBinary.GetWriteBuffer()
@@ -83,59 +152,30 @@ LineLoop:
 			continue
 		}
 
-		key := strconv.Itoa(int(reader.Days())) + ":" + unsafeString(name)
+		nameStr := unsafeString(name)
 
+		days := reader.Days()
+		key := strconv.Itoa(int(days)) + ":" + nameStr
 		if u.existsCache.Exists(key) {
 			continue LineLoop
 		}
 
 		if newTagged[key] {
+			// already processed
 			continue LineLoop
 		}
-		n++
-
-		m, err := urlParse(unsafeString(name))
-		if err != nil {
-			continue
-		}
-
-		newTagged[key] = true
-
-		wb.Reset()
-		tagsBuf.Reset()
-		tag1 = tag1[:0]
 
-		t := fmt.Sprintf("__name__=%s", m.Path)
-		tag1 = append(tag1, t)
-		tagsBuf.WriteString(t)
-
-		// don't upload any other tag but __name__
-		// if either main metric (m.Path) or each metric (*) is ignored
-		ignoreAllButName := u.ignoredMetrics[m.Path] || u.ignoredMetrics["*"]
-		tagsWritten := 1
-		for k, v := range m.Query() {
-			t := fmt.Sprintf("%s=%s", k, v[0])
-			tagsBuf.WriteString(t)
-			tagsWritten++
-
-			if !ignoreAllButName {
-				tag1 = append(tag1, t)
-			}
-		}
-
-		for i := 0; i < len(tag1); i++ {
-			wb.WriteUint16(reader.Days())
-			wb.WriteString(tag1[i])
-			wb.WriteBytes(name)
-			wb.WriteUVarint(uint64(tagsWritten))
-			wb.Write(tagsBuf.Bytes())
-			wb.WriteUint32(version)
-		}
+		n++
 
-		_, err = out.Write(wb.Bytes())
-		if err != nil {
+		if err = u.parseName(nameStr, days, tag1, wb, tagsBuf); err != nil {
+			u.logger.Warn("parse",
+				zap.String("metric", nameStr), zap.String("type", "tagged"), zap.String("name", filename), zap.Error(err),
+			)
+			continue LineLoop
+		} else if _, err = out.Write(wb.Bytes()); err != nil {
 			return n, nil, err
 		}
+		newTagged[key] = true
 	}
 
 	return n, newTagged, nil

uploader/tagged_test.go

@@ -4,9 +4,12 @@ import (
 	"fmt"
 	"net/url"
 	"strconv"
+	"strings"
 	"testing"
 
+	"github.com/lomik/carbon-clickhouse/helper/RowBinary"
 	"github.com/lomik/carbon-clickhouse/helper/escape"
+	"github.com/lomik/zapwriter"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -57,3 +60,30 @@ func BenchmarkKeyConcat(b *testing.B) {
 		_ = strconv.Itoa(int(unum)) + ":" + path
 	}
 }
+
+func TestTagged_parseName_Overflow(t *testing.T) {
+	var tag1 []string
+	wb := RowBinary.GetWriteBuffer()
+	tagsBuf := RowBinary.GetWriteBuffer()
+	defer wb.Release()
+	defer tagsBuf.Release()
+
+	logger := zapwriter.Logger("upload")
+	base := &Base{
+		queue:   make(chan string, 1024),
+		inQueue: make(map[string]bool),
+		logger:  logger,
+		config:  &Config{TableName: "test"},
+	}
+	var sb strings.Builder
+	sb.WriteString("very_long_name_field1.very_long_name_field2.very_long_name_field3.very_long_name_field4?")
+	for i := 0; i < 100; i++ {
+		if i > 0 {
+			sb.WriteString("&")
+		}
+		sb.WriteString(fmt.Sprintf("very_long_tag%d=very_long_value%d", i, i))
+	}
+	u := NewTagged(base)
+	err := u.parseName(sb.String(), 10, tag1, wb, tagsBuf)
+	assert.Equal(t, errBufOverflow, err)
+}