Rust: Add Database Sources + tokio-postgres support

rust/ql/lib/codeql/rust/Concepts.qll

@@ -100,6 +100,32 @@ class ModeledEnvironmentSource extends EnvironmentSource::Range {
   ModeledEnvironmentSource() { sourceNode(this, "environment-source") }
 }
 
+/**
+ * A data flow source corresponding to the program's database reads.
+ */
+final class DatabaseSource = DatabaseSource::Range;
+
+/**
+ * Provides a class for modeling new sources for the program's database reads.
+ */
+module DatabaseSource {
+  /**
+   * A data flow source corresponding to the program's database reads.
+   */
+  abstract class Range extends ThreatModelSource::Range {
+    override string getThreatModel() { result = "database" }
+
+    override string getSourceType() { result = "DatabaseSource" }
+  }
+}
+
+/**
+ * An externally modeled source for data from the program's database.
+ */
+class ModeledDatabaseSource extends DatabaseSource::Range {
+  ModeledDatabaseSource() { sourceNode(this, "database") }
+}
+
 /**
  * A data flow source for remote (network) data.
  */

rust/ql/lib/codeql/rust/frameworks/tokio-postgres.model.yml

@@ -0,0 +1,24 @@
+extensions:
+  - addsTo:
+      pack: codeql/rust-all
+      extensible: sinkModel
+    data:
+      - ["repo:https://github.com/sfackler/rust-postgres:tokio-postgres", "<crate::client::Client>::execute", "Argument[0]", "sql-injection", "manual"]
+      - ["repo:https://github.com/sfackler/rust-postgres:tokio-postgres", "<crate::client::Client>::batch_execute", "Argument[0]", "sql-injection", "manual"]
+      - ["repo:https://github.com/sfackler/rust-postgres:tokio-postgres", "<crate::client::Client>::execute_raw", "Argument[0]", "sql-injection", "manual"]
+      - ["repo:https://github.com/sfackler/rust-postgres:tokio-postgres", "<crate::client::Client>::prepare", "Argument[0]", "sql-injection", "manual"]
+      - ["repo:https://github.com/sfackler/rust-postgres:tokio-postgres", "<crate::client::Client>::prepare_typed", "Argument[0]", "sql-injection", "manual"]
+      - ["repo:https://github.com/sfackler/rust-postgres:tokio-postgres", "<crate::client::Client>::query", "Argument[0]", "sql-injection", "manual"]
+      - ["repo:https://github.com/sfackler/rust-postgres:tokio-postgres", "<crate::client::Client>::query_opt", "Argument[0]", "sql-injection", "manual"]
+      - ["repo:https://github.com/sfackler/rust-postgres:tokio-postgres", "<crate::client::Client>::query_raw", "Argument[0]", "sql-injection", "manual"]
+      - ["repo:https://github.com/sfackler/rust-postgres:tokio-postgres", "<crate::client::Client>::query_typed", "Argument[0]", "sql-injection", "manual"]
+      - ["repo:https://github.com/sfackler/rust-postgres:tokio-postgres", "<crate::client::Client>::query_typed_raw", "Argument[0]", "sql-injection", "manual"]
+      - ["repo:https://github.com/sfackler/rust-postgres:tokio-postgres", "<crate::client::Client>::simple_query", "Argument[0]", "sql-injection", "manual"]
+      - ["repo:https://github.com/sfackler/rust-postgres:tokio-postgres", "<crate::client::Client>::simple_query_raw", "Argument[0]", "sql-injection", "manual"]
+
+  - addsTo:
+      pack: codeql/rust-all
+      extensible: sourceModel
+    data:
+      - ["repo:https://github.com/sfackler/rust-postgres:tokio-postgres", "<crate::row::Row>::get", "ReturnValue", "database", "manual"]
+      - ["repo:https://github.com/sfackler/rust-postgres:tokio-postgres", "<crate::row::Row>::try_get", "ReturnValue.Variant[crate::result::Result::Ok(0)]", "database", "manual"]

rust/ql/test/library-tests/frameworks/postgres/Postgres.ql

@@ -1,9 +1,10 @@
 import rust
+import codeql.rust.Concepts
 import codeql.rust.security.SqlInjectionExtensions
 import utils.test.InlineExpectationsTest
 
 module PostgresTest implements TestSig {
-  string getARelevantTag() { result = "sql-sink" }
+  string getARelevantTag() { result = ["sql-sink", "database-read"] }
 
   predicate hasActualResult(Location location, string element, string tag, string value) {
     exists(SqlInjection::Sink sink |
@@ -13,6 +14,14 @@ module PostgresTest implements TestSig {
       tag = "sql-sink" and
       value = ""
     )
+    or
+    exists(ModeledDatabaseSource source |
+      location = source.getLocation() and
+      location.getFile().getBaseName() != "" and
+      element = source.toString() and
+      tag = "database-read" and
+      value = ""
+    )
   }
 }
 

rust/ql/test/library-tests/frameworks/postgres/main.rs

@@ -33,9 +33,9 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     // conn.query_typed_raw(query.as_str(), &[])?;
 
     for row in &conn.query("SELECT id, name, age FROM person", &[])? {  // $ sql-sink
-        let id: i32 = row.get("id");
-        let name: &str = row.get("name");
-        let age: i32 = row.get("age");
+        let id: i32 = row.get("id");        // $ database-read
+        let name: &str = row.try_get("name")?;   // $ database-read
+        let age: i32 = row.try_get("age").unwrap();      // $ database-read
         println!("found person: {} {} {}", id, name, age);
     }