Merge branch 'main' into standarize_internal_metrics

pkg/internal/ebpf/common/sql_detect_transform.go

@@ -1,6 +1,9 @@
 package ebpfcommon
 
 import (
+	"encoding/binary"
+	"errors"
+	"fmt"
 	"strings"
 	"unsafe"
 
@@ -29,6 +32,52 @@ func asciiToUpper(input string) string {
 	return string(out)
 }
 
+func isASCII(s string) bool {
+	for i := 0; i < len(s); i++ {
+		c := s[i]
+		if (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') ||
+			c == '.' || c == '_' || c == ' ' || c == '-' {
+			continue
+		}
+		return false
+	}
+
+	return true
+}
+
+func detectSQLBytes(b []byte) (string, string, string) {
+	op, table, sql := detectSQL(string(b))
+	if !validSQL(op, table) {
+		if isPostgresBindCommand(b) {
+			statement, portal, args, err := parsePostgresBindCommand(b)
+			if err == nil {
+				op = "BIND"
+				table = fmt.Sprintf("%s.%s", statement, portal)
+				for _, arg := range args {
+					if isASCII(arg) {
+						sql += arg + " "
+					}
+				}
+			}
+		} else if isPostgresQueryCommand(b) {
+			text, err := parsePosgresQueryCommand(b)
+			if err == nil {
+				query := asciiToUpper(text)
+				if strings.HasPrefix(query, "EXECUTE ") {
+					parts := strings.Split(text, " ")
+					op = parts[0]
+					if len(parts) > 1 {
+						table = parts[1]
+					}
+					sql = text
+				}
+			}
+		}
+	}
+
+	return op, table, sql
+}
+
 func detectSQL(buf string) (string, string, string) {
 	b := asciiToUpper(buf)
 	for _, q := range []string{"SELECT", "UPDATE", "DELETE", "INSERT", "ALTER", "CREATE", "DROP"} {
@@ -44,6 +93,120 @@ func detectSQL(buf string) (string, string, string) {
 	return "", "", ""
 }
 
+func isPostgresBindCommand(b []byte) bool {
+	return isPostgresCommand('B', b)
+}
+
+func isPostgresQueryCommand(b []byte) bool {
+	return isPostgresCommand('Q', b)
+}
+
+func isPostgresCommand(lookup byte, b []byte) bool {
+	if len(b) < 5 {
+		return false
+	}
+
+	if b[0] == lookup {
+		size := int32(binary.BigEndian.Uint32(b[1:5]))
+		if size < 0 || size > 1000 {
+			return false
+		}
+		return true
+	}
+
+	return false
+}
+
+// nolint:cyclop
+func parsePostgresBindCommand(buf []byte) (string, string, []string, error) {
+	statement := []byte{}
+	portal := []byte{}
+	args := []string{}
+
+	size := int(binary.BigEndian.Uint32(buf[1:5]))
+	if size > len(buf) {
+		size = len(buf)
+	}
+	ptr := 5
+
+	// parse statement, zero terminated string
+	for {
+		if ptr >= size {
+			return string(statement), string(portal), args, errors.New("too short, while parsing statement")
+		}
+		b := buf[ptr]
+		ptr++
+
+		if b == 0 {
+			break
+		}
+		statement = append(statement, b)
+	}
+
+	// parse portal, zero terminated string
+	for {
+		if ptr >= size {
+			return string(statement), string(portal), args, errors.New("too short, while parsing portal")
+		}
+		b := buf[ptr]
+		ptr++
+
+		if b == 0 {
+			break
+		}
+		portal = append(portal, b)
+	}
+
+	if ptr+2 >= size {
+		return string(statement), string(portal), args, errors.New("too short, while parsing format codes")
+	}
+
+	formats := int16(binary.BigEndian.Uint16(buf[ptr : ptr+2]))
+	ptr += 2
+	for i := 0; i < int(formats); i++ {
+		// ignore format codes
+		if ptr+2 >= size {
+			return string(statement), string(portal), args, errors.New("too short, while parsing format codes")
+		}
+		ptr += 2
+	}
+
+	params := int16(binary.BigEndian.Uint16(buf[ptr : ptr+2]))
+	ptr += 2
+	for i := 0; i < int(params); i++ {
+		if ptr+4 >= size {
+			return string(statement), string(portal), args, errors.New("too short, while parsing params")
+		}
+		argLen := int(binary.BigEndian.Uint32(buf[ptr : ptr+4]))
+		ptr += 4
+		arg := []byte{}
+		for j := 0; j < int(argLen); j++ {
+			if ptr >= size {
+				break
+			}
+			arg = append(arg, buf[ptr])
+			ptr++
+		}
+		args = append(args, string(arg))
+	}
+
+	return string(statement), string(portal), args, nil
+}
+
+func parsePosgresQueryCommand(buf []byte) (string, error) {
+	size := int(binary.BigEndian.Uint32(buf[1:5]))
+	if size > len(buf) {
+		size = len(buf)
+	}
+	ptr := 5
+
+	if ptr > size {
+		return "", errors.New("too short")
+	}
+
+	return string(buf[ptr:size]), nil
+}
+
 func TCPToSQLToSpan(trace *TCPRequestInfo, op, table, sql string) request.Span {
 	peer := ""
 	hostname := ""

pkg/internal/ebpf/common/sql_detect_transform_test.go

@@ -0,0 +1,200 @@
+package ebpfcommon
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+type bindParseResult struct {
+	statement    string
+	portal       string
+	args         []string
+	hasErr       bool
+	hasASCIIArgs bool
+}
+
+type bindTest struct {
+	name   string
+	bytes  []byte
+	isBind bool
+	result bindParseResult
+}
+
+func TestPostgresBindParsing(t *testing.T) {
+	for _, ts := range []bindTest{
+		{
+			name:   "Valid bind",
+			bytes:  []byte{66, 0, 0, 0, 52, 0, 101, 99, 116, 111, 95, 49, 49, 53, 56, 0, 0, 1, 0, 1, 0, 1, 0, 0, 0, 19, 114, 101, 99, 111, 109, 109, 101, 110, 100, 97, 116, 105, 111, 110, 67, 97, 99, 104, 101, 0, 3, 0, 1, 0, 1, 0, 1, 69, 0, 0, 0, 9, 0, 0, 0, 0, 0, 83, 0, 0, 0, 4, 0, 4, 34, 101, 110, 97, 98, 108, 101, 100, 34, 32, 70, 82, 79, 77, 32, 34, 102, 101, 97, 116, 117, 114, 101, 102, 108, 97, 103, 115, 34, 32, 65, 83, 32, 102, 48, 32, 87, 72, 69, 82, 69, 32, 40, 102, 48, 46, 34, 110, 97, 109, 101, 34, 32, 61, 32, 36, 49, 41, 0, 0, 1, 0, 0, 0, 25, 68, 0, 0, 0, 15, 83, 101, 99, 116, 111, 95, 49, 49, 53, 56, 0, 72, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
+			isBind: true,
+			result: bindParseResult{
+				statement:    "",
+				portal:       "ecto_1158",
+				args:         []string{"recommendationCache"},
+				hasErr:       false,
+				hasASCIIArgs: true,
+			},
+		},
+		{
+			name:   "Less length than needed",
+			bytes:  []byte{66, 0, 0, 0, 12, 0, 101, 99, 116, 111, 95, 49, 49, 53, 56, 0, 0, 1, 0, 1, 0, 1, 0, 0, 0, 19, 114, 101, 99, 111, 109, 109, 101, 110, 100, 97, 116, 105, 111, 110, 67, 97, 99, 104, 101, 0, 3, 0, 1, 0, 1, 0, 1, 69, 0, 0, 0, 9, 0, 0, 0, 0, 0, 83, 0, 0, 0, 4, 0, 4, 34, 101, 110, 97, 98, 108, 101, 100, 34, 32, 70, 82, 79, 77, 32, 34, 102, 101, 97, 116, 117, 114, 101, 102, 108, 97, 103, 115, 34, 32, 65, 83, 32, 102, 48, 32, 87, 72, 69, 82, 69, 32, 40, 102, 48, 46, 34, 110, 97, 109, 101, 34, 32, 61, 32, 36, 49, 41, 0, 0, 1, 0, 0, 0, 25, 68, 0, 0, 0, 15, 83, 101, 99, 116, 111, 95, 49, 49, 53, 56, 0, 72, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
+			isBind: true,
+			result: bindParseResult{
+				statement:    "",
+				portal:       "ecto_1",
+				args:         []string{},
+				hasErr:       true,
+				hasASCIIArgs: true,
+			},
+		},
+		{
+			name:   "Not a bind",
+			bytes:  []byte{67, 0, 0, 0, 52, 0, 101, 99, 116, 111, 95, 49, 49, 53, 56, 0, 0, 1, 0, 1, 0, 1, 0, 0, 0, 19, 114, 101, 99, 111, 109, 109, 101, 110, 100, 97, 116, 105, 111, 110, 67, 97, 99, 104, 101, 0, 3, 0, 1, 0, 1, 0, 1, 69, 0, 0, 0, 9, 0, 0, 0, 0, 0, 83, 0, 0, 0, 4, 0, 4, 34, 101, 110, 97, 98, 108, 101, 100, 34, 32, 70, 82, 79, 77, 32, 34, 102, 101, 97, 116, 117, 114, 101, 102, 108, 97, 103, 115, 34, 32, 65, 83, 32, 102, 48, 32, 87, 72, 69, 82, 69, 32, 40, 102, 48, 46, 34, 110, 97, 109, 101, 34, 32, 61, 32, 36, 49, 41, 0, 0, 1, 0, 0, 0, 25, 68, 0, 0, 0, 15, 83, 101, 99, 116, 111, 95, 49, 49, 53, 56, 0, 72, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
+			isBind: false,
+			result: bindParseResult{},
+		},
+		{
+			name:   "Too long",
+			bytes:  []byte{66, 100, 0, 0, 52, 0, 101, 99, 116, 111, 95, 49, 49, 53, 56, 0, 0, 1, 0, 1, 0, 1, 0, 0, 0, 19, 114, 101, 99, 111, 109, 109, 101, 110, 100, 97, 116, 105, 111, 110, 67, 97, 99, 104, 101, 0, 3, 0, 1, 0, 1, 0, 1, 69, 0, 0, 0, 9, 0, 0, 0, 0, 0, 83, 0, 0, 0, 4, 0, 4, 34, 101, 110, 97, 98, 108, 101, 100, 34, 32, 70, 82, 79, 77, 32, 34, 102, 101, 97, 116, 117, 114, 101, 102, 108, 97, 103, 115, 34, 32, 65, 83, 32, 102, 48, 32, 87, 72, 69, 82, 69, 32, 40, 102, 48, 46, 34, 110, 97, 109, 101, 34, 32, 61, 32, 36, 49, 41, 0, 0, 1, 0, 0, 0, 25, 68, 0, 0, 0, 15, 83, 101, 99, 116, 111, 95, 49, 49, 53, 56, 0, 72, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
+			isBind: false,
+			result: bindParseResult{},
+		},
+		{
+			name:   "Too short",
+			bytes:  []byte{67, 100},
+			isBind: false,
+			result: bindParseResult{},
+		},
+		{
+			name:   "Empty",
+			bytes:  []byte{},
+			isBind: false,
+			result: bindParseResult{},
+		},
+		{
+			name:   "A bind, but without anything reasonable",
+			bytes:  []byte{66, 0, 0, 0, 12},
+			isBind: true,
+			result: bindParseResult{
+				statement:    "",
+				portal:       "",
+				args:         []string{},
+				hasErr:       true,
+				hasASCIIArgs: true,
+			},
+		},
+		{
+			name:   "Crazy long argument length",
+			bytes:  []byte{66, 0, 0, 0, 52, 0, 101, 99, 116, 111, 95, 49, 49, 53, 56, 0, 0, 1, 0, 1, 0, 1, 0, 0, 100, 19, 114, 101, 99, 111, 109, 109, 101, 110, 100, 97, 116, 105, 111, 110, 67, 97, 99, 104, 101, 0, 3, 0, 1, 0, 1, 0, 1, 69, 0, 0, 0, 9, 0, 0, 0, 0, 0, 83, 0, 0, 0, 4, 0, 4, 34, 101, 110, 97, 98, 108, 101, 100, 34, 32, 70, 82, 79, 77, 32, 34, 102, 101, 97, 116, 117, 114, 101, 102, 108, 97, 103, 115, 34, 32, 65, 83, 32, 102, 48, 32, 87, 72, 69, 82, 69, 32, 40, 102, 48, 46, 34, 110, 97, 109, 101, 34, 32, 61, 32, 36, 49, 41, 0, 0, 1, 0, 0, 0, 25, 68, 0, 0, 0, 15, 83, 101, 99, 116, 111, 95, 49, 49, 53, 56, 0, 72, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0},
+			isBind: true,
+			result: bindParseResult{
+				statement:    "",
+				portal:       "ecto_1158",
+				args:         []string{"recommendationCache"},
+				hasErr:       false,
+				hasASCIIArgs: false,
+			},
+		},
+	} {
+		t.Run(ts.name, func(t *testing.T) {
+			ok := isPostgresBindCommand(ts.bytes)
+			assert.Equal(t, ts.isBind, ok)
+			if ok {
+				statement, portal, args, err := parsePostgresBindCommand(ts.bytes)
+				if ts.result.hasErr {
+					assert.True(t, err != nil)
+				} else {
+					assert.True(t, err == nil)
+				}
+				assert.Equal(t, ts.result.statement, statement)
+				assert.Equal(t, ts.result.portal, portal)
+				if ts.result.hasASCIIArgs {
+					assert.Equal(t, ts.result.args, args)
+				} else {
+					for _, arg := range args {
+						assert.False(t, isASCII(arg))
+					}
+				}
+			}
+		})
+	}
+}
+
+type qSQLTest struct {
+	name  string
+	bytes []byte
+	op    string
+	table string
+	sql   string
+}
+
+func TestPostgresQueryParsing(t *testing.T) {
+	for _, ts := range []qSQLTest{
+		{
+			name:  "Query prepared statement",
+			bytes: []byte{81, 0, 0, 0, 28, 101, 120, 101, 99, 117, 116, 101, 32, 109, 121, 95, 99, 111, 110, 116, 97, 99, 116, 115, 32, 40, 49, 41, 0, 69, 76, 69, 67, 84, 32, 42, 32, 102, 114, 111, 109, 32, 97, 99, 99, 111, 117, 110, 116, 105, 110, 103, 46, 99, 111, 110, 116, 97, 99, 116, 115, 32, 87, 72, 69, 82, 69, 32, 105, 100, 32, 61, 32, 36, 49, 0, 53, 90, 51, 106, 119, 55, 54, 111, 100, 85, 115, 57, 78, 75, 72, 73, 76, 119, 120, 104, 108, 81, 118, 50, 98, 122, 70, 72, 111, 73, 70, 48, 61},
+			op:    "execute",
+			table: "my_contacts",
+			sql:   "execute my_contacts (1)",
+		},
+		{
+			name:  "Query prepared statement bad len",
+			bytes: []byte{81, 0, 0, 0, 7, 101, 120, 101, 99, 117, 116, 101, 32, 109, 121, 95, 99, 111, 110, 116, 97, 99, 116, 115, 32, 40, 49, 41, 0, 69, 76, 69, 67, 84, 32, 42, 32, 102, 114, 111, 109, 32, 97, 99, 99, 111, 117, 110, 116, 105, 110, 103, 46, 99, 111, 110, 116, 97, 99, 116, 115, 32, 87, 72, 69, 82, 69, 32, 105, 100, 32, 61, 32, 36, 49, 0, 53, 90, 51, 106, 119, 55, 54, 111, 100, 85, 115, 57, 78, 75, 72, 73, 76, 119, 120, 104, 108, 81, 118, 50, 98, 122, 70, 72, 111, 73, 70, 48, 61},
+			op:    "",
+			table: "",
+			sql:   "",
+		},
+		{
+			name:  "small len",
+			bytes: []byte{81, 0, 0},
+			op:    "",
+			table: "",
+			sql:   "",
+		},
+		{
+			name:  "empty",
+			bytes: []byte{},
+			op:    "",
+			table: "",
+			sql:   "",
+		},
+	} {
+		t.Run(ts.name, func(t *testing.T) {
+			op, table, sql := detectSQLBytes(ts.bytes)
+			assert.Equal(t, ts.op, op)
+			assert.Equal(t, ts.table, table)
+			assert.Equal(t, ts.sql, sql)
+		})
+	}
+}
+
+type asciiSQLTest struct {
+	name string
+	s    string
+	ok   bool
+}
+
+func TestIsASCII(t *testing.T) {
+	for _, ts := range []asciiSQLTest{
+		{
+			name: "Positive test",
+			s:    "This is a test_.-1234",
+			ok:   true,
+		},
+		{
+			name: "Bad char",
+			s:    "This is\x00 a test_.-1234",
+			ok:   false,
+		},
+		{
+			name: "Empty",
+			s:    "",
+			ok:   true,
+		},
+	} {
+		t.Run(ts.name, func(t *testing.T) {
+			res := isASCII(ts.s)
+			assert.Equal(t, ts.ok, res)
+		})
+	}
+}