feat: -mode encode-marked, and optional MySQL executable-comment lexing

cmd/sqlprocessor/main.go

@@ -15,6 +15,23 @@ import (
 	"github.com/DataDog/go-sqllexer"
 )
 
+// execComments mirrors -executable-comments. It is read-only after flag
+// parsing, so the tokenizers can reach it without threading a parameter
+// through every call site.
+//
+// It defaults off because the encodings this tool emits are what downstream
+// models were trained on, and turning it on changes them: a payload that used
+// to encode as NUMBER MULTILINE_COMMENT now yields the tokens MySQL actually
+// executes. Enable it and retrain together.
+var execComments bool
+
+func newLexer(input string) *sqllexer.Lexer {
+	if execComments {
+		return sqllexer.New(input, sqllexer.WithExecutableComments(true))
+	}
+	return sqllexer.New(input)
+}
+
 type tokenOut struct {
 	Type  string `json:"type"`
 	Value string `json:"value"`
@@ -44,6 +61,8 @@ func main() {
 	outDir := flag.String("outdir", "", "Output directory (default: same as input file)")
 	includeEmpty := flag.Bool("include-empty", false, "Include empty/whitespace-only lines")
 	mode := flag.String("mode", "analyze", "Processing mode: analyze, tokenize or encode")
+	flag.BoolVar(&execComments, "executable-comments", false,
+		"Lex the body of MySQL executable comments (/*! ... */) as SQL instead of emitting one comment token")
 	flag.Parse()
 
 	inputs := make([]string, 0, 1+len(flag.Args()))
@@ -58,7 +77,7 @@ func main() {
 	}
 
 	switch *mode {
-	case "analyze", "tokenize", "encode":
+	case "analyze", "tokenize", "encode", "encode-marked":
 	default:
 		fmt.Fprintf(os.Stderr, "Invalid -mode %q (expected analyze, tokenize or encode)\n", *mode)
 		os.Exit(2)
@@ -187,7 +206,7 @@ func processReader(r io.Reader, format string, out io.Writer, includeEmpty bool,
 		}
 		first := true
 		for {
-			line, err := readLine(reader)
+			line, err := readLine(reader, mode != "encode-marked")
 			if err != nil {
 				if errors.Is(err, io.EOF) {
 					break
@@ -221,7 +240,7 @@ func processReader(r io.Reader, format string, out io.Writer, includeEmpty bool,
 	case "jsonl":
 		writer := bufio.NewWriter(out)
 		for {
-			line, err := readLine(reader)
+			line, err := readLine(reader, mode != "encode-marked")
 			if err != nil {
 				if errors.Is(err, io.EOF) {
 					break
@@ -249,7 +268,7 @@ func processReader(r io.Reader, format string, out io.Writer, includeEmpty bool,
 	case "txt":
 		writer := bufio.NewWriter(out)
 		for {
-			line, err := readLine(reader)
+			line, err := readLine(reader, mode != "encode-marked")
 			if err != nil {
 				if errors.Is(err, io.EOF) {
 					break
@@ -269,6 +288,10 @@ func processReader(r io.Reader, format string, out io.Writer, includeEmpty bool,
 				if _, err := fmt.Fprintf(writer, "%d\t%s\n", lineNum, tokenizeLineTypesOnly(line)); err != nil {
 					return err
 				}
+			} else if mode == "encode-marked" {
+				if _, err := fmt.Fprintf(writer, "%d\t%s\n", lineNum, tokenizeLineTypesOnlyMarked(line)); err != nil {
+					return err
+				}
 			} else {
 				rec := tokenizeLine(line, lineNum)
 				if err := writeTxtRecord(writer, rec); err != nil {
@@ -295,13 +318,15 @@ func encodeValue(mode, line string, lineNum int) any {
 		return tokenizeLineTypesOnly(line)
 	case "encode":
 		return encoded{Line: lineNum, Text: tokenizeLineTypesOnly(line)}
+	case "encode-marked":
+		return encoded{Line: lineNum, Text: tokenizeLineTypesOnlyMarked(line)}
 	default:
 		return tokenizeLine(line, lineNum)
 	}
 }
 
 func tokenizeLineTypesOnly(line string) string {
-	lexer := sqllexer.New(line)
+	lexer := newLexer(line)
 	var types []string
 	for {
 		tok := lexer.Scan()
@@ -313,8 +338,61 @@ func tokenizeLineTypesOnly(line string) string {
 	return strings.Join(types, " ")
 }
 
+// tokenizeLineTypesOnlyMarked is tokenizeLineTypesOnly with quote positions
+// preserved as a QUOTE token.
+//
+// Deleting quotes stops a dangling one swallowing the input, but it also erases
+// the most common injection shape there is. After the strip,
+//
+//	"anything' OR 'x'='x"  and  "anything or x=x"
+//
+// are the same token sequence, so nothing downstream can tell an attack from an
+// ordinary phrase. Here the line is split *at* each quote, each segment lexed
+// separately, and a QUOTE emitted between them:
+//
+//	IDENT QUOTE SPACE KEYWORD SPACE QUOTE IDENT QUOTE OPERATOR QUOTE IDENT
+//
+// A quote still never reaches the lexer, so the swallowing problem stays fixed.
+// This is a different feature encoding, not a refinement of the other one: a
+// quote becomes a lexical boundary, so `12'34` is NUMBER QUOTE NUMBER here and a
+// single NUMBER under the strip. A model trained on one cannot score the other.
+//
+// Callers must read this line with readLine(reader, false); with the quotes
+// already deleted it degrades to tokenizeLineTypesOnly.
+func tokenizeLineTypesOnlyMarked(line string) string {
+	var types []string
+	start := 0
+	for i := 0; i < len(line); i++ {
+		if c := line[i]; c == '\'' || c == '"' {
+			types = appendSegmentTypes(types, line[start:i])
+			types = append(types, quoteTokenName)
+			start = i + 1
+		}
+	}
+	types = appendSegmentTypes(types, line[start:])
+	return strings.Join(types, " ")
+}
+
+// quoteTokenName is not a go-sqllexer type; the encoding inserts it. Upper case
+// and space-free like every other name, so it survives a whitespace split.
+const quoteTokenName = "QUOTE"
+
+func appendSegmentTypes(dst []string, segment string) []string {
+	if segment == "" {
+		return dst
+	}
+	lexer := newLexer(segment)
+	for {
+		tok := lexer.Scan()
+		if tok == nil || tok.Type == sqllexer.EOF {
+			return dst
+		}
+		dst = append(dst, tokenTypeName(tok.Type))
+	}
+}
+
 func tokenizeLine(line string, lineNum int) record {
-	lexer := sqllexer.New(line)
+	lexer := newLexer(line)
 
 	tokens := make([]tokenOut, 0, 32)
 	hasError := false
@@ -416,7 +494,11 @@ func tokenTypeName(t sqllexer.TokenType) string {
 	}
 }
 
-func readLine(reader *bufio.Reader) (string, error) {
+// readLine reads one JSONL record. stripQuotes selects the feature encoding's
+// preprocessing: true deletes quotes before lexing (the "encode" contract),
+// false keeps them for a caller that marks their positions itself
+// ("encode-marked"). See tokenizeLineTypesOnlyMarked for why that matters.
+func readLine(reader *bufio.Reader, stripQuotes bool) (string, error) {
 	var buf []byte
 	for {
 		chunk, err := reader.ReadSlice('\n')
@@ -453,6 +535,9 @@ func readLine(reader *bufio.Reader) (string, error) {
 	// inference must apply the same strip or their features will not match.
 	s = strings.TrimSuffix(s, "\n")
 	s = strings.TrimSuffix(s, "\r")
+	if !stripQuotes {
+		return s, nil
+	}
 	s = strings.ReplaceAll(s, "'", "")
 	s = strings.ReplaceAll(s, "\"", "")
 	return s, nil

cmd/sqlprocessor/main_test.go

@@ -0,0 +1,126 @@
+package main
+
+import (
+	"bufio"
+	"strings"
+	"testing"
+)
+
+// The encodings this binary emits are a contract: a model is fitted to one of
+// them, so a change here silently changes what that model sees. These pin the
+// two that "encode" and "encode-marked" produce.
+
+func TestTokenizeLineTypesOnly(t *testing.T) {
+	// Input reaches this function already stripped of quotes by readLine.
+	for _, tc := range []struct{ in, want string }{
+		{"anything OR x=x", "IDENT SPACE KEYWORD SPACE IDENT OPERATOR IDENT"},
+		{"SELECT * FROM t", "COMMAND SPACE WILDCARD SPACE KEYWORD SPACE IDENT"},
+		{"", ""},
+	} {
+		if got := tokenizeLineTypesOnly(tc.in); got != tc.want {
+			t.Errorf("tokenizeLineTypesOnly(%q)\n   got:  %s\n   want: %s", tc.in, got, tc.want)
+		}
+	}
+}
+
+func TestTokenizeLineTypesOnlyMarked(t *testing.T) {
+	for _, tc := range []struct{ in, want string }{
+		{"anything' OR 'x'='x", "IDENT QUOTE SPACE KEYWORD SPACE QUOTE IDENT QUOTE OPERATOR QUOTE IDENT"},
+		{"O'Brien", "IDENT QUOTE IDENT"},
+		{"'-'", "QUOTE OPERATOR QUOTE"},
+		{"' --", "QUOTE SPACE COMMENT"},
+		{"no quotes here", "IDENT SPACE IDENT SPACE IDENT"},
+		{"'", "QUOTE"},
+		{`"`, "QUOTE"},
+		{"", ""},
+	} {
+		if got := tokenizeLineTypesOnlyMarked(tc.in); got != tc.want {
+			t.Errorf("tokenizeLineTypesOnlyMarked(%q)\n   got:  %s\n   want: %s", tc.in, got, tc.want)
+		}
+	}
+}
+
+// TestMarkedSeparatesWhatTheStripCannot is the reason the mode exists: after the
+// strip a quoted tautology and an ordinary phrase are the same sequence, so
+// nothing downstream can tell them apart.
+func TestMarkedSeparatesWhatTheStripCannot(t *testing.T) {
+	const attack, benign = "anything' OR 'x'='x", "anything or x=x"
+	strip := func(s string) string {
+		return tokenizeLineTypesOnly(strings.NewReplacer("'", "", `"`, "").Replace(s))
+	}
+	if strip(attack) != strip(benign) {
+		t.Fatalf("premise no longer holds: the strip already separates these")
+	}
+	if tokenizeLineTypesOnlyMarked(attack) == tokenizeLineTypesOnlyMarked(benign) {
+		t.Errorf("marked encoding fails to separate them: both %s",
+			tokenizeLineTypesOnlyMarked(attack))
+	}
+}
+
+// TestMarkedNeverLexesAQuote guards the property the strip provides and this
+// mode must not lose: a dangling quote must not swallow the rest of the input.
+func TestMarkedNeverLexesAQuote(t *testing.T) {
+	for _, in := range []string{
+		"'; DROP TABLE users; --", "unbalanced ' quote", `a "b`, "admin'--",
+	} {
+		for _, tok := range strings.Fields(tokenizeLineTypesOnlyMarked(in)) {
+			switch tok {
+			case "STRING", "INCOMPLETE_STRING", "QUOTED_IDENT":
+				t.Errorf("%q produced %s: a quote reached the lexer", in, tok)
+			}
+		}
+	}
+}
+
+func TestReadLineStripQuotesFlag(t *testing.T) {
+	const line = `"1' OR '1'='1"` + "\n" // one JSON-encoded record
+	for _, tc := range []struct {
+		strip bool
+		want  string
+	}{
+		{true, "1 OR 1=1"},
+		{false, "1' OR '1'='1"},
+	} {
+		got, err := readLine(bufio.NewReader(strings.NewReader(line)), tc.strip)
+		if err != nil {
+			t.Fatalf("readLine(strip=%v): %v", tc.strip, err)
+		}
+		if got != tc.want {
+			t.Errorf("readLine(strip=%v) = %q, want %q", tc.strip, got, tc.want)
+		}
+	}
+}
+
+// A MySQL executable comment runs on the server but encodes as a single
+// comment token, so `1/*!50000union select pw from users*/` reaches a model as
+// two tokens with the union hidden inside one of them. -executable-comments
+// unfolds the body; it is off by default so the shipped encodings do not move.
+func TestExecutableCommentsFlag(t *testing.T) {
+	const in = "1/*!50000union select pw from users*/"
+
+	if got, want := tokenizeLineTypesOnly(in), "NUMBER MULTILINE_COMMENT"; got != want {
+		t.Errorf("default\n   got:  %s\n   want: %s", got, want)
+	}
+
+	execComments = true
+	defer func() { execComments = false }()
+
+	want := "NUMBER KEYWORD SPACE COMMAND SPACE IDENT SPACE KEYWORD SPACE IDENT"
+	if got := tokenizeLineTypesOnly(in); got != want {
+		t.Errorf("-executable-comments\n   got:  %s\n   want: %s", got, want)
+	}
+
+	// The marked encoding picks the option up too.
+	if got := tokenizeLineTypesOnlyMarked(in); got != want {
+		t.Errorf("marked -executable-comments\n   got:  %s\n   want: %s", got, want)
+	}
+
+	// A quote ends the segment and the lexer scanning it, so an executable
+	// comment opened before a quote does not stay open across it: the tail is
+	// lexed on its own and the closing */ falls out as WILDCARD OPERATOR.
+	// Quote positions are the point of this encoding, so they win.
+	wantSplit := "NUMBER KEYWORD SPACE COMMAND SPACE QUOTE IDENT WILDCARD OPERATOR"
+	if got := tokenizeLineTypesOnlyMarked("1/*!50000union select 'pw*/"); got != wantSplit {
+		t.Errorf("marked across a quote\n   got:  %s\n   want: %s", got, wantSplit)
+	}
+}

obfuscator_test.go

@@ -321,11 +321,13 @@ func TestObfuscator(t *testing.T) {
 			// postgres #> operator
 			input:    `SELECT * FROM users where '{"a": 1, "b": 2}'::jsonb #> '{a}'`,
 			expected: `SELECT * FROM users where ?::jsonb #> ?`,
+			dbms:     DBMSPostgres,
 		},
 		{
 			// postgres #>> operator
 			input:    `SELECT * FROM users where '{"a": 1, "b": 2}'::jsonb #>> '{a}'`,
 			expected: `SELECT * FROM users where ?::jsonb #>> ?`,
+			dbms:     DBMSPostgres,
 		},
 		{
 			// postgres ? operator