fix(rust): SQL injections

Iron-E · Jun 22, 2023 · 3ff82eb · 3ff82eb
1 parent 5c832c2
commit 3ff82eb
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 14 deletions.
diff --git a/after/queries/rust/highlights.scm b/after/queries/rust/highlights.scm
@@ -84,43 +84,43 @@
 		(field_identifier) @punctuation.special
 		(identifier) @punctuation.special
 	]
-	(#lua-match? @punctuation.special "^r#")
+	(#lua-match? @punctuation.special "^r#\"")
 	(#offset-from! "start" @punctuation.special 0 0 0 2)
 	(#set! priority 128)
 )
 
 ((raw_string_literal) @punctuation.special
-	(#lua-match? @punctuation.special "^r")
+	(#lua-match? @punctuation.special "^r\"")
 	(#offset-from! "start" @punctuation.special 0 0 0 1)
 )
 
 ((raw_string_literal) @punctuation.special
-	(#lua-match? @punctuation.special "^r#")
+	(#lua-match? @punctuation.special "^r#\"")
 	(#offset-from! "start" @punctuation.special 0 0 0 2)
 )
 
 ((raw_string_literal) @punctuation.special
-	(#lua-match? @punctuation.special "^r#")
+	(#lua-match? @punctuation.special "^r#\"")
 	(#offset-from! "end" @punctuation.special 0 -1 0 0)
 )
 
 ((raw_string_literal) @punctuation.special
-	(#lua-match? @punctuation.special "^r##")
+	(#lua-match? @punctuation.special "^r##\"")
 	(#offset-from! "start" @punctuation.special 0 0 0 3)
 )
 
 ((raw_string_literal) @punctuation.special
-	(#lua-match? @punctuation.special "^r##")
+	(#lua-match? @punctuation.special "^r##\"")
 	(#offset-from! "end" @punctuation.special 0 -2 0 0)
 )
 
 ((raw_string_literal) @punctuation.special
-	(#lua-match? @punctuation.special "^r###")
+	(#lua-match? @punctuation.special "^r###\"")
 	(#offset-from! "start" @punctuation.special 0 0 0 4)
 )
 
 ((raw_string_literal) @punctuation.special
-	(#lua-match? @punctuation.special "^r###")
+	(#lua-match? @punctuation.special "^r###\"")
 	(#offset-from! "end" @punctuation.special 0 -3 0 0)
 )
 

diff --git a/after/queries/rust/injections.scm b/after/queries/rust/injections.scm
@@ -5,10 +5,56 @@
 		path: (identifier) @_path (#eq? @_path "sqlx")
 		name: (identifier) @_name (#any-of? @_name "query" "query_as")
 	)
-	(token_tree
-		[
-			(string_literal)
-			(raw_string_literal)
-		] @sql
+	(token_tree (string_literal) @sql
+		(#offset! @sql 0 1 0 -1)
+		(#set! "priority" 102)
 	)
-) @foo
+)
+
+(macro_invocation
+	macro: (scoped_identifier
+		path: (identifier) @_path (#eq? @_path "sqlx")
+		name: (identifier) @_name (#any-of? @_name "query" "query_as")
+	)
+	(token_tree (raw_string_literal) @sql
+		(#lua-match? @sql "^r\"")
+		(#offset! @sql 0 2 0 -1)
+		(#set! "priority" 102)
+	)
+)
+
+(macro_invocation
+	macro: (scoped_identifier
+		path: (identifier) @_path (#eq? @_path "sqlx")
+		name: (identifier) @_name (#any-of? @_name "query" "query_as")
+	)
+	(token_tree (raw_string_literal) @sql
+		(#lua-match? @sql "^r#\"")
+		(#offset! @sql 0 3 0 -2)
+		(#set! "priority" 102)
+	)
+)
+
+(macro_invocation
+	macro: (scoped_identifier
+		path: (identifier) @_path (#eq? @_path "sqlx")
+		name: (identifier) @_name (#any-of? @_name "query" "query_as")
+	)
+	(token_tree (raw_string_literal) @sql
+		(#lua-match? @sql "^r##\"")
+		(#offset! @sql 0 4 0 -3)
+		(#set! "priority" 102)
+	)
+)
+
+(macro_invocation
+	macro: (scoped_identifier
+		path: (identifier) @_path (#eq? @_path "sqlx")
+		name: (identifier) @_name (#any-of? @_name "query" "query_as")
+	)
+	(token_tree (raw_string_literal) @sql
+		(#lua-match? @sql "^r###\"")
+		(#offset! @sql 0 5 0 -4)
+		(#set! "priority" 102)
+	)
+)