diff --git a/cmd/sqlprocessor/main.go b/cmd/sqlprocessor/main.go index 3b3e32a..1104693 100644 --- a/cmd/sqlprocessor/main.go +++ b/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 diff --git a/cmd/sqlprocessor/main_test.go b/cmd/sqlprocessor/main_test.go new file mode 100644 index 0000000..b907b7b --- /dev/null +++ b/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) + } +} diff --git a/obfuscator_test.go b/obfuscator_test.go index 7dd2d2a..787f4af 100644 --- a/obfuscator_test.go +++ b/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 diff --git a/sqllexer.go b/sqllexer.go index 55dad0a..ecdc5b6 100644 --- a/sqllexer.go +++ b/sqllexer.go @@ -66,6 +66,11 @@ func (t *Token) getLastValueToken() *LastValueToken { type LexerConfig struct { DBMS DBMSType `json:"dbms,omitempty"` + + // ExecutableComments makes the lexer scan the body of a MySQL executable + // comment (/*! ... */ and /*!NNNNN ... */) as SQL rather than emitting the + // whole construct as a single MULTILINE_COMMENT token. + ExecutableComments bool `json:"executable_comments,omitempty"` } type lexerOption func(*LexerConfig) @@ -77,6 +82,24 @@ func WithDBMS(dbms DBMSType) lexerOption { } } +// WithExecutableComments controls how MySQL executable comments are tokenized. +// +// MySQL executes the body of /*! ... */ and /*!NNNNN ... */ when the server +// version is at least NNNNN, so `id=1/*!50000union select pw from users*/` +// runs the union on any MySQL >= 5.0. Lexing the construct as one comment +// token hides that statement from anything reading the token stream. +// +// With this enabled the opening delimiter and its version digits are consumed +// and the body is lexed as ordinary SQL, so the stream matches what the server +// executes. The default is disabled: query normalization wants a MySQL +// optimizer hint to stay a comment, and changing that would alter obfuscated +// output for existing callers. +func WithExecutableComments(enabled bool) lexerOption { + return func(c *LexerConfig) { + c.ExecutableComments = enabled + } +} + // SQL Lexer inspired from Rob Pike's talk on Lexical Scanning in Go type Lexer struct { src string // the input src string @@ -88,6 +111,8 @@ type Lexer struct { hasDigits bool // true if the token has digits isTableIndicator bool // true if the token is a table indicator isSimpleIdentifier bool // true if current quoted ident started with a letter and only used alphanumerics afterwards + execComments bool // mirrors config.ExecutableComments, read once per token + inExecComment bool // true while scanning the body of a MySQL executable comment } func New(input string, opts ...lexerOption) *Lexer { @@ -99,12 +124,16 @@ func New(input string, opts ...lexerOption) *Lexer { for _, opt := range opts { opt(lexer.config) } + lexer.execComments = lexer.config.ExecutableComments return lexer } // Scan scans the next token and returns it. func (s *Lexer) Scan() *Token { ch := s.peek() + if s.execComments { + ch = s.skipExecutableCommentDelimiters(ch) + } switch { case isSpace(ch): return s.scanWhitespace() @@ -154,13 +183,27 @@ func (s *Lexer) Scan() *Token { } return s.scanUnknown() // backtick is only valid in mysql case ch == '#': + // SQL Server names temporary tables #temp, so there a # opens an + // identifier. It is the only dialect that does, and saying so is the + // caller's job. if s.config.DBMS == DBMSSQLServer { return s.scanIdentifier(ch) - } else if s.config.DBMS == DBMSMySQL { - // MySQL treats # as a comment - return s.scanSingleLineComment(ch) } - return s.scanOperator(ch) + // PostgreSQL spells three JSON path operators #> #>> #-, but only + // PostgreSQL does, so recognising them requires being told. Guessing + // from the two characters alone is what a bypass is made of: MySQL + // reads `admin'#-` as a comment to end of line and logs you in. + if s.config.DBMS == DBMSPostgres && (s.lookAhead(1) == '>' || s.lookAhead(1) == '-') { + return s.scanOperator(ch) + } + // Everything else: a comment to end of line. MySQL says so outright, + // and a caller that has not named its dialect gets the same reading, + // because the alternative is worse for the one that cannot afford to + // guess. A WAF scanning untrusted input has no idea what is behind it; + // scanning `admin'#` as an identifier and an operator hides a login + // bypass that MySQL would execute, while scanning a stray # in some + // other dialect as a comment costs a token. + return s.scanSingleLineComment(ch) case ch == '@': if s.lookAhead(1) == '@' { if isAlphaNumeric(s.lookAhead(2)) { @@ -399,7 +442,8 @@ func (s *Lexer) scanIdentifier(ch rune) *Token { } // If we found a complete keyword and next char is whitespace - if node.isEnd && (isPunctuation(ch) || isSpace(ch) || isMultiLineComment(ch, s.lookAhead(1)) || isEOF(ch)) { + if node.isEnd && (isPunctuation(ch) || isSpace(ch) || isMultiLineComment(ch, s.lookAhead(1)) || + s.atExecCommentClose(ch) || isEOF(ch)) { s.cursor = pos + 1 // Include the last matched character s.isTableIndicator = node.isTableIndicator return s.emit(node.tokenType) @@ -572,6 +616,62 @@ func (s *Lexer) scanMultiLineComment() *Token { return s.emit(MULTILINE_COMMENT) } +// skipExecutableCommentDelimiters consumes any run of MySQL executable comment +// delimiters at the cursor and returns the rune Scan should dispatch on. The +// delimiters emit no token of their own, so scanning has to carry on past them +// to reach one; looping rather than recursing keeps a run of empty comments +// (/*!*//*!*/...) from growing the stack once per delimiter. +func (s *Lexer) skipExecutableCommentDelimiters(ch rune) rune { + for { + // Closing delimiter of a body we are lexing. Reached only outside + // strings and comments, so a */ inside a quoted literal does not end + // the body -- which is what MySQL does too. + if s.atExecCommentClose(ch) { + s.nextBy(2) + s.inExecComment = false + ch = s.peek() + continue + } + if isMultiLineComment(ch, s.lookAhead(1)) && s.lookAhead(2) == '!' { + s.consumeExecutableCommentOpening() + ch = s.peek() + continue + } + return ch + } +} + +// atExecCommentClose reports whether the cursor sits on the `*/` that ends the +// executable comment body currently being lexed. +// +// It is a token boundary like a space is: `/*!50000or*/` executes as the +// operator, so `or` has to reach the keyword trie as a complete word rather +// than decay to an IDENT because the rune after it is a `*`. +func (s *Lexer) atExecCommentClose(ch rune) bool { + return s.inExecComment && ch == '*' && s.lookAhead(1) == '/' +} + +// consumeExecutableCommentOpening consumes the opening `/*!` of a MySQL +// executable comment plus the version gate when one is present, leaving the +// cursor on the body so it is lexed as ordinary SQL until Scan reaches the +// closing `*/`. +// +// A version gate is exactly five digits (50000 is 5.0.0). Anything shorter is +// not a gate, so those digits stay part of the body and lex as a number. +func (s *Lexer) consumeExecutableCommentOpening() { + s.nextBy(3) // consume the opening slash, asterisk and bang + + digits := 0 + for digits < 5 && isDigit(s.lookAhead(digits)) { + digits++ + } + if digits == 5 { + s.nextBy(5) + } + + s.inExecComment = true +} + func (s *Lexer) scanPunctuation() *Token { s.start = s.cursor s.next() diff --git a/sqllexer_test.go b/sqllexer_test.go index 63b94d7..100cd3f 100644 --- a/sqllexer_test.go +++ b/sqllexer_test.go @@ -2,6 +2,7 @@ package sqllexer import ( "fmt" + "strings" "testing" ) @@ -893,6 +894,7 @@ here */`, {SPACE, " "}, {IDENT, "users"}, }, + lexerOpts: []lexerOption{WithDBMS(DBMSPostgres)}, }, { name: "extracts JSON sub-object at the specified path as text", @@ -912,6 +914,7 @@ here */`, {SPACE, " "}, {IDENT, "users"}, }, + lexerOpts: []lexerOption{WithDBMS(DBMSPostgres)}, }, { name: "JSON path return any item for the specified JSON value", @@ -1148,6 +1151,178 @@ here */`, {IDENT, "my_table"}, }, }, + { + name: "WAITFOR DELAY is a keyword pair, not two identifiers", + input: "1 WAITFOR DELAY '0:0:5'", + expected: []TokenSpec{ + {NUMBER, "1"}, + {SPACE, " "}, + {KEYWORD, "WAITFOR"}, + {SPACE, " "}, + {KEYWORD, "DELAY"}, + {SPACE, " "}, + {STRING, "'0:0:5'"}, + }, + }, + { + name: "# opens a comment without being told the dialect", + input: "1 or 1=1#", + expected: []TokenSpec{ + {NUMBER, "1"}, + {SPACE, " "}, + {KEYWORD, "or"}, + {SPACE, " "}, + {NUMBER, "1"}, + {OPERATOR, "="}, + {NUMBER, "1"}, + {COMMENT, "#"}, + }, + }, + { + // #> is a JSON path operator only in PostgreSQL. Undeclared, the + // MySQL reading wins: everything after the # is commented out. + // Guessing from the two characters would leave `admin'#-` scanning + // as an operator while MySQL executes it as a login bypass. + name: "an undeclared #> is a comment, not a JSON operator", + input: "data #> '{a}'", + expected: []TokenSpec{ + {IDENT, "data"}, + {SPACE, " "}, + {COMMENT, "#> '{a}'"}, + }, + }, + { + name: "a declared PostgreSQL #- stays a JSON operator", + input: "data #- '{a}'", + lexerOpts: []lexerOption{WithDBMS(DBMSPostgres)}, + expected: []TokenSpec{ + {IDENT, "data"}, + {SPACE, " "}, + {JSON_OP, "#-"}, + {SPACE, " "}, + {STRING, "'{a}'"}, + }, + }, + { + name: "mysql executable comment stays one comment token by default", + input: "1/*!50000union select pw from users*/", + expected: []TokenSpec{ + {NUMBER, "1"}, + {MULTILINE_COMMENT, "/*!50000union select pw from users*/"}, + }, + }, + { + name: "mysql executable comment body is lexed when enabled", + input: "1/*!50000union select pw from users*/", + lexerOpts: []lexerOption{WithExecutableComments(true)}, + expected: []TokenSpec{ + {NUMBER, "1"}, + {KEYWORD, "union"}, + {SPACE, " "}, + {COMMAND, "select"}, + {SPACE, " "}, + {IDENT, "pw"}, + {SPACE, " "}, + {KEYWORD, "from"}, + {SPACE, " "}, + {IDENT, "users"}, + }, + }, + { + name: "mysql executable comment without a version gate", + input: "1/*!union select 1*/", + lexerOpts: []lexerOption{WithExecutableComments(true)}, + expected: []TokenSpec{ + {NUMBER, "1"}, + {KEYWORD, "union"}, + {SPACE, " "}, + {COMMAND, "select"}, + {SPACE, " "}, + {NUMBER, "1"}, + }, + }, + { + name: "a version gate is exactly five digits, shorter runs stay in the body", + input: "1/*!500or 1=1*/", + lexerOpts: []lexerOption{WithExecutableComments(true)}, + expected: []TokenSpec{ + {NUMBER, "1"}, + {NUMBER, "500"}, + {KEYWORD, "or"}, + {SPACE, " "}, + {NUMBER, "1"}, + {OPERATOR, "="}, + {NUMBER, "1"}, + }, + }, + { + name: "plain multiline comments are untouched when enabled", + input: "SELECT /* plain */ 1", + lexerOpts: []lexerOption{WithExecutableComments(true)}, + expected: []TokenSpec{ + {COMMAND, "SELECT"}, + {SPACE, " "}, + {MULTILINE_COMMENT, "/* plain */"}, + {SPACE, " "}, + {NUMBER, "1"}, + }, + }, + { + name: "a closing delimiter inside a string does not end the body", + input: "1/*!50000select '*/' */", + lexerOpts: []lexerOption{WithExecutableComments(true)}, + expected: []TokenSpec{ + {NUMBER, "1"}, + {COMMAND, "select"}, + {SPACE, " "}, + {STRING, "'*/'"}, + {SPACE, " "}, + }, + }, + { + name: "an empty executable comment yields no tokens", + input: "/*!*/", + lexerOpts: []lexerOption{WithExecutableComments(true)}, + expected: []TokenSpec{}, + }, + { + name: "a keyword against the closing delimiter is still a keyword", + input: "1 /*!50000or*/ sleep(1)", + lexerOpts: []lexerOption{WithExecutableComments(true)}, + expected: []TokenSpec{ + {NUMBER, "1"}, + {SPACE, " "}, + {KEYWORD, "or"}, + {SPACE, " "}, + {FUNCTION, "sleep"}, + {PUNCTUATION, "("}, + {NUMBER, "1"}, + {PUNCTUATION, ")"}, + }, + }, + { + name: "a command against the closing delimiter is still a command", + input: "1/*!50000union*/select 1", + lexerOpts: []lexerOption{WithExecutableComments(true)}, + expected: []TokenSpec{ + {NUMBER, "1"}, + {KEYWORD, "union"}, + {COMMAND, "select"}, + {SPACE, " "}, + {NUMBER, "1"}, + }, + }, + { + name: "an unterminated executable comment ends at EOF", + input: "1/*!50000union select", + lexerOpts: []lexerOption{WithExecutableComments(true)}, + expected: []TokenSpec{ + {NUMBER, "1"}, + {KEYWORD, "union"}, + {SPACE, " "}, + {COMMAND, "select"}, + }, + }, } for _, tt := range tests { @@ -1508,3 +1683,23 @@ func ExampleLexer() { fmt.Println(token) } } + +// Executable comment delimiters emit no token, so Scan has to keep scanning to +// find one. It does that in a loop; this pins that, since recursing once per +// delimiter would grow the stack in proportion to the input. +func TestExecutableCommentsDoNotGrowTheStack(t *testing.T) { + input := strings.Repeat("/*!*/", 200000) + "1" + + lexer := New(input, WithExecutableComments(true)) + got := 0 + for { + if tok := lexer.Scan(); tok.Type == EOF { + break + } + got++ + } + + if got != 1 { + t.Errorf("got %d tokens, want 1", got) + } +} diff --git a/sqllexer_utils.go b/sqllexer_utils.go index 97add6f..41cabf9 100644 --- a/sqllexer_utils.go +++ b/sqllexer_utils.go @@ -154,6 +154,12 @@ var keywords = []string{ "SKIP", "IF", "ONLY", + + // T-SQL statement delay. WAITFOR DELAY '00:00:05' is the SQL Server + // equivalent of SLEEP(5), and the oracle a blind injection reads one bit + // at a time; without these both words scan as ordinary identifiers. + "WAITFOR", + "DELAY", } var (