From d2d5f31599186ee885aa0b1ad38b62663fd7d3f3 Mon Sep 17 00:00:00 2001
From: Nick Rolfe <nickrolfe@github.com>
Date: Tue, 4 May 2021 16:52:35 +0100
Subject: [PATCH 1/3] Escape keys for files and folders

---
 extractor/src/extractor.rs | 35 +++++++++++++++++++++++++++++++++--
 1 file changed, 33 insertions(+), 2 deletions(-)

diff --git a/extractor/src/extractor.rs b/extractor/src/extractor.rs
index 1f282d8211c..0c586e223e2 100644
--- a/extractor/src/extractor.rs
+++ b/extractor/src/extractor.rs
@@ -190,6 +190,24 @@ pub fn extract(
     Ok(Program(visitor.trap_writer.trap_output))
 }
 
+/// Escapes a string for use in a TRAP key, by replacing special characters with
+/// HTML entities.
+fn escape_key(s: &str) -> String {
+    let mut escaped = String::new();
+    for c in s.chars() {
+        match c {
+            '&' => escaped.push_str("&amp;"),
+            '{' => escaped.push_str("&lbrace;"),
+            '}' => escaped.push_str("&rbrace;"),
+            '"' => escaped.push_str("&quot;"),
+            '@' => escaped.push_str("&commat;"),
+            '#' => escaped.push_str("&num;"),
+            _ => escaped.push(c),
+        }
+    }
+    escaped
+}
+
 /// Normalizes the path according the common CodeQL specification. Assumes that
 /// `path` has already been canonicalized using `std::fs::canonicalize`.
 fn normalize_path(path: &Path) -> String {
@@ -230,11 +248,11 @@ fn normalize_path(path: &Path) -> String {
 }
 
 fn full_id_for_file(path: &Path) -> String {
-    format!("{};sourcefile", normalize_path(path))
+    format!("{};sourcefile", escape_key(&normalize_path(path)))
 }
 
 fn full_id_for_folder(path: &Path) -> String {
-    format!("{};folder", normalize_path(path))
+    format!("{};folder", escape_key(&normalize_path(path)))
 }
 
 struct ChildNode {
@@ -731,3 +749,16 @@ fn limit_string_test() {
     assert_eq!("hi ☹", limit_string(&"hi ☹☹".to_owned(), 6));
     assert_eq!("hi ", limit_string(&"hi ☹☹".to_owned(), 5));
 }
+
+#[test]
+fn escape_key_test() {
+    assert_eq!("foo!", escape_key("foo!"));
+    assert_eq!("foo&lbrace;&rbrace;", escape_key("foo{}"));
+    assert_eq!("&lbrace;&rbrace;", escape_key("{}"));
+    assert_eq!("", escape_key(""));
+    assert_eq!("/path/to/foo.rb", escape_key("/path/to/foo.rb"));
+    assert_eq!(
+        "/path/to/foo&amp;&lbrace;&rbrace;&quot;&commat;&num;.rb",
+        escape_key("/path/to/foo&{}\"@#.rb")
+    );
+}

From 99ae17de03c193cb36ee542fc158ca8662fe9c7b Mon Sep 17 00:00:00 2001
From: Nick Rolfe <nickrolfe@github.com>
Date: Wed, 5 May 2021 12:54:23 +0100
Subject: [PATCH 2/3] Avoid copying key when it doesn't need escaping

---
 extractor/src/extractor.rs | 40 +++++++++++++++++++++++++++-----------
 1 file changed, 29 insertions(+), 11 deletions(-)

diff --git a/extractor/src/extractor.rs b/extractor/src/extractor.rs
index 0c586e223e2..afb058098a6 100644
--- a/extractor/src/extractor.rs
+++ b/extractor/src/extractor.rs
@@ -1,4 +1,5 @@
 use node_types::{EntryKind, Field, NodeTypeMap, Storage, TypeName};
+use std::borrow::Cow;
 use std::collections::BTreeMap as Map;
 use std::collections::BTreeSet as Set;
 use std::fmt;
@@ -192,20 +193,37 @@ pub fn extract(
 
 /// Escapes a string for use in a TRAP key, by replacing special characters with
 /// HTML entities.
-fn escape_key(s: &str) -> String {
-    let mut escaped = String::new();
-    for c in s.chars() {
+fn escape_key<'a, S: Into<Cow<'a, str>>>(key: S) -> Cow<'a, str> {
+    fn needs_escaping(c: char) -> bool {
         match c {
-            '&' => escaped.push_str("&amp;"),
-            '{' => escaped.push_str("&lbrace;"),
-            '}' => escaped.push_str("&rbrace;"),
-            '"' => escaped.push_str("&quot;"),
-            '@' => escaped.push_str("&commat;"),
-            '#' => escaped.push_str("&num;"),
-            _ => escaped.push(c),
+            '&' => true,
+            '{' => true,
+            '}' => true,
+            '"' => true,
+            '@' => true,
+            '#' => true,
+            _ => false,
         }
     }
-    escaped
+
+    let key = key.into();
+    if key.contains(needs_escaping) {
+        let mut escaped = String::with_capacity(key.len());
+        for c in key.chars() {
+            match c {
+                '&' => escaped.push_str("&amp;"),
+                '{' => escaped.push_str("&lbrace;"),
+                '}' => escaped.push_str("&rbrace;"),
+                '"' => escaped.push_str("&quot;"),
+                '@' => escaped.push_str("&commat;"),
+                '#' => escaped.push_str("&num;"),
+                _ => escaped.push(c),
+            }
+        }
+        Cow::Owned(escaped)
+    } else {
+        key
+    }
 }
 
 /// Normalizes the path according the common CodeQL specification. Assumes that

From c37f390efc7404ebd6be50edb4c4bb6d8b6cc24d Mon Sep 17 00:00:00 2001
From: Nick Rolfe <nickrolfe@github.com>
Date: Wed, 5 May 2021 13:21:16 +0100
Subject: [PATCH 3/3] Reserve more capacity for escaped key

---
 extractor/src/extractor.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/extractor/src/extractor.rs b/extractor/src/extractor.rs
index afb058098a6..0c28c4a20ca 100644
--- a/extractor/src/extractor.rs
+++ b/extractor/src/extractor.rs
@@ -208,7 +208,7 @@ fn escape_key<'a, S: Into<Cow<'a, str>>>(key: S) -> Cow<'a, str> {
 
     let key = key.into();
     if key.contains(needs_escaping) {
-        let mut escaped = String::with_capacity(key.len());
+        let mut escaped = String::with_capacity(2 * key.len());
         for c in key.chars() {
             match c {
                 '&' => escaped.push_str("&amp;"),