Correctly escape strings

Cargo.lock

@@ -725,6 +725,7 @@ dependencies = [
  "fallible-iterator 0.1.5 (registry+https://github.com/rust-lang/crates.io-index)",
  "indicatif 0.9.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "postgres 0.15.2 (registry+https://github.com/rust-lang/crates.io-index)",
+ "rand 0.5.5 (registry+https://github.com/rust-lang/crates.io-index)",
  "rayon 1.0.2 (registry+https://github.com/rust-lang/crates.io-index)",
  "state-map 0.1.0 (git+https://github.com/matrix-org/rust-matrix-state-map)",
  "string_cache 0.7.3 (registry+https://github.com/rust-lang/crates.io-index)",

Cargo.toml

@@ -9,6 +9,7 @@ clap = "2.32.0"
 fallible-iterator = "0.1.5"
 indicatif = "0.9.0"
 postgres = "0.15.2"
+rand = "0.5.5"
 rayon = "1.0.2"
 string_cache = "0.7.3"
 

src/compressor.rs

@@ -264,7 +264,8 @@ fn test_new_map() {
         (10, 9),
         (11, 10),
         (13, 12),
-    ].into_iter()
+    ]
+    .into_iter()
     .collect();
 
     for sg in 0i64..=13i64 {

src/database.rs

@@ -15,8 +15,12 @@
 use fallible_iterator::FallibleIterator;
 use indicatif::{ProgressBar, ProgressStyle};
 use postgres::{Connection, TlsMode};
+use rand::distributions::Alphanumeric;
+use rand::{thread_rng, Rng};
 
+use std::borrow::Cow;
 use std::collections::BTreeMap;
+use std::fmt;
 
 use StateGroupEntry;
 
@@ -47,7 +51,8 @@ pub fn get_data_from_db(db_url: &str, room_id: &str) -> BTreeMap<i64, StateGroup
                 } else {
                     None
                 }
-            }).collect();
+            })
+            .collect();
 
         if missing_sgs.is_empty() {
             break;
@@ -74,7 +79,8 @@ fn get_initial_data_from_db(conn: &Connection, room_id: &str) -> BTreeMap<i64, S
                 LEFT JOIN state_group_edges AS e ON (m.id = e.state_group)
                 WHERE m.room_id = $1
             "#,
-        ).unwrap();
+        )
+        .unwrap();
 
     let trans = conn.transaction().unwrap();
     let mut rows = stmt.lazy_query(&trans, &[&room_id], 1000).unwrap();
@@ -123,7 +129,8 @@ fn get_missing_from_db(conn: &Connection, missing_sgs: &[i64]) -> BTreeMap<i64,
                 FROM state_group_edges
                 WHERE state_group = ANY($1)
             "#,
-        ).unwrap();
+        )
+        .unwrap();
     let trans = conn.transaction().unwrap();
 
     let mut rows = stmt.lazy_query(&trans, &[&missing_sgs], 100).unwrap();
@@ -140,3 +147,38 @@ fn get_missing_from_db(conn: &Connection, missing_sgs: &[i64]) -> BTreeMap<i64,
 
     state_group_map
 }
+
+/// Helper function that escapes the wrapped text when writing SQL
+pub struct PGEscapse<'a>(pub &'a str);
+
+impl<'a> fmt::Display for PGEscapse<'a> {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        let mut delim = Cow::from("$$");
+        while self.0.contains(&delim as &str) {
+            let s: String = thread_rng().sample_iter(&Alphanumeric).take(10).collect();
+
+            delim = format!("${}$", s).into();
+        }
+
+        write!(f, "{}{}{}", delim, self.0, delim)
+    }
+}
+
+#[test]
+fn test_pg_escape() {
+    let s = format!("{}", PGEscapse("test"));
+    assert_eq!(s, "$$test$$");
+
+    let dodgy_string = "test$$ing";
+
+    let s = format!("{}", PGEscapse(dodgy_string));
+
+    // prefix and suffixes should match
+    let start_pos = s.find(dodgy_string).expect("expected to find dodgy string");
+    let end_pos = start_pos + dodgy_string.len();
+    assert_eq!(s[..start_pos], s[end_pos..]);
+
+    // .. and they should start and end with '$'
+    assert_eq!(&s[0..1], "$");
+    assert_eq!(&s[start_pos - 1..start_pos], "$");
+}

src/main.rs

@@ -12,7 +12,6 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-
 //! This is a tool that attempts to further compress state maps within a
 //! Synapse instance's database. Specifically, it aims to reduce the number of
 //! rows that a given room takes up in the `state_groups_state` table.
@@ -22,6 +21,7 @@ extern crate clap;
 extern crate fallible_iterator;
 extern crate indicatif;
 extern crate postgres;
+extern crate rand;
 extern crate rayon;
 extern crate state_map;
 extern crate string_cache;
@@ -30,6 +30,7 @@ mod compressor;
 mod database;
 
 use compressor::Compressor;
+use database::PGEscapse;
 
 use clap::{App, Arg};
 use indicatif::{ProgressBar, ProgressStyle};
@@ -259,7 +260,15 @@ fn main() {
                         } else {
                             write!(output, "    ,");
                         }
-                        writeln!(output, "({}, '{}', '{}', '{}', '{}')", sg, room_id, t, s, e);
+                        writeln!(
+                            output,
+                            "({}, {}, {}, {}, {})",
+                            sg,
+                            PGEscapse(room_id),
+                            PGEscapse(t),
+                            PGEscapse(s),
+                            PGEscapse(e)
+                        );
                     }
                     writeln!(output, ";");
                 }
@@ -268,7 +277,6 @@ fn main() {
                     writeln!(output, "COMMIT;");
                 }
                 writeln!(output);
-
             }
 
             pb.inc(1);
@@ -304,7 +312,8 @@ fn main() {
             } else {
                 Ok(())
             }
-        }).expect("expected state to match");
+        })
+        .expect("expected state to match");
 
     pb.finish();