datasecurity: add postgres backend engine and report sub task scan errors in payload

Author: aimenebelfodil

Cargo.lock

@@ -3,11 +3,18 @@ name = "datasecurity"
 version = "0.1.0"
 edition = "2024"
 
+[features]
+default = ["engine-postgres"]
+# Each backend engine is gated behind its own feature so the set of compiled
+# engines is explicit. Add new engines as `engine-<name>` here.
+engine-postgres = ["dep:postgres"]
+
 [dependencies]
 anyhow = "1.0.100"
 core = { path = "../../core" }
 dd_sds = { package = "dd-sensitive-data-scanner", version = "=0.1.0-20260715-fcd8f9716b61", default-features = false, features = ["dd-sds"] }
 libc = "0.2.182"
+postgres = { version = "0.19", optional = true }
 serde = { version = "1", features = ["derive"] }
 serde_json = "1"
 

pkg/collector/sharedlibrary/rustchecks/Cargo.lock

@@ -0,0 +1,57 @@
+//! Backend scan engines: run a sub task's query and return its column data.
+
+use anyhow::Result;
+use serde_json::Value;
+
+use crate::config::SubTask;
+
+#[cfg(feature = "engine-postgres")]
+mod postgres;
+
+/// A data-source engine that runs a sub task's query and returns the result as
+/// a `{ column: [values] }` map ready for the scanner.
+pub trait ScanEngine: Sync {
+    /// Engine name, matched against the sub task platform.
+    fn name(&self) -> &'static str;
+    /// Runs the sub task's query and returns its columns.
+    fn run_scan(&self, sub_task: &SubTask) -> Result<Value>;
+}
+
+/// Compiled engines. Add a new engine here behind its `engine-*` feature.
+fn engines() -> &'static [&'static dyn ScanEngine] {
+    &[
+        #[cfg(feature = "engine-postgres")]
+        &postgres::ENGINE,
+    ]
+}
+
+fn engine_for(platform: &str) -> Result<&'static dyn ScanEngine> {
+    engines()
+        .iter()
+        .copied()
+        .find(|engine| engine.name() == platform)
+        .ok_or_else(|| {
+            let available = engines()
+                .iter()
+                .map(|engine| engine.name())
+                .collect::<Vec<_>>()
+                .join(", ");
+            anyhow::anyhow!("unsupported platform {platform:?} (compiled engines: [{available}])")
+        })
+}
+
+/// Runs the sub task on the engine selected by its platform.
+pub fn execute_scan(sub_task: &SubTask) -> Result<Value> {
+    engine_for(&sub_task.platform)?.run_scan(sub_task)
+}
+
+#[cfg(all(test, feature = "engine-postgres"))]
+mod tests {
+    use super::engine_for;
+
+    #[test]
+    fn resolves_postgres_engine() {
+        assert_eq!(engine_for("postgres").unwrap().name(), "postgres");
+        assert!(engine_for("mysql").is_err());
+    }
+}

pkg/collector/sharedlibrary/rustchecks/checks/datasecurity/Cargo.toml

@@ -0,0 +1,102 @@
+//! Postgres scan engine.
+
+use std::time::Duration;
+
+use anyhow::{Context, Result};
+use postgres::config::SslMode;
+use postgres::{Config, NoTls, Row};
+use serde_json::{Map, Value};
+
+use crate::backend::ScanEngine;
+use crate::config::SubTask;
+
+pub struct PostgresEngine;
+pub const ENGINE: PostgresEngine = PostgresEngine;
+
+impl ScanEngine for PostgresEngine {
+    fn name(&self) -> &'static str {
+        "postgres"
+    }
+
+    fn run_scan(&self, sub_task: &SubTask) -> Result<Value> {
+        let conn = &sub_task.connection;
+        let timeout = sub_task.timeout_seconds;
+        println!(
+            "datasecurity: connecting to postgres host={} port={} dbname={} user={} timeout={}s",
+            conn.host, conn.port, conn.dbname, conn.username, timeout
+        );
+
+        let mut config = Config::new();
+        config
+            .host(&conn.host)
+            .port(conn.port)
+            .dbname(&conn.dbname)
+            .user(&conn.username)
+            .password(&conn.password)
+            .ssl_mode(SslMode::Disable)
+            // `0` means no statement timeout in postgres.
+            .options(&format!("-c statement_timeout={}", timeout * 1000));
+        if timeout > 0 {
+            config.connect_timeout(Duration::from_secs(timeout));
+        }
+
+        let mut client = config.connect(NoTls).context("connecting to postgres")?;
+
+        let rows = client
+            .query(sub_task.query.as_str(), &[])
+            .context("running postgres query")?;
+
+        Ok(rows_to_columns(&rows))
+    }
+}
+
+/// Builds a `{ column: [string values] }` map from the query rows.
+fn rows_to_columns(rows: &[Row]) -> Value {
+    let Some(first) = rows.first() else {
+        return Value::Object(Map::new());
+    };
+
+    let names: Vec<String> = first.columns().iter().map(|c| c.name().to_string()).collect();
+    let mut columns: Map<String, Value> = names
+        .iter()
+        .map(|name| (name.clone(), Value::Array(Vec::new())))
+        .collect();
+
+    for row in rows {
+        for (i, name) in names.iter().enumerate() {
+            if let Some(Value::Array(values)) = columns.get_mut(name) {
+                values.push(cell_to_value(row, i));
+            }
+        }
+    }
+
+    Value::Object(columns)
+}
+
+/// Renders a postgres cell as a string value (or null), so the scanner sees a
+/// uniform `{ column: [string values] }` shape.
+fn cell_to_value(row: &Row, index: usize) -> Value {
+    fn string(v: Option<impl ToString>) -> Value {
+        v.map(|n| Value::String(n.to_string())).unwrap_or(Value::Null)
+    }
+
+    if let Ok(v) = row.try_get::<_, Option<String>>(index) {
+        return string(v);
+    }
+    if let Ok(v) = row.try_get::<_, Option<i64>>(index) {
+        return string(v);
+    }
+    if let Ok(v) = row.try_get::<_, Option<i32>>(index) {
+        return string(v);
+    }
+    if let Ok(v) = row.try_get::<_, Option<i16>>(index) {
+        return string(v);
+    }
+    if let Ok(v) = row.try_get::<_, Option<f64>>(index) {
+        return string(v);
+    }
+    if let Ok(v) = row.try_get::<_, Option<bool>>(index) {
+        return string(v);
+    }
+    Value::Null
+}

pkg/collector/sharedlibrary/rustchecks/checks/datasecurity/src/backend/mod.rs

@@ -1,9 +1,9 @@
 use anyhow::{Context, Result};
 use core::*;
-use serde_json::Value;
 
+use crate::backend;
 use crate::config::{CheckConfig, SubTask};
-use crate::payload::ScanEventPayload;
+use crate::payload::{Match, ScanEventPayload, ScanStatus};
 use crate::scanning::Scanner;
 
 /// Check implementation (scaffolding).
@@ -33,25 +33,35 @@ fn run_sub_task(
     sub_task: &SubTask,
 ) -> Result<()> {
     println!(
-        "datasecurity: running sub task (sub_task_id={})",
-        sub_task.sub_task_id
+        "datasecurity: running sub task (sub_task_id={}, platform={})",
+        sub_task.sub_task_id, sub_task.platform
     );
 
-    // TODO(DSEC-139): fetch the rows from postgres.
-    let data = fetch_data(sub_task);
-    let matches = scanner.scan(&data).context("scanning sub task data")?;
+    // A sub task failure is reported inside the payload (status=ERROR) rather
+    // than aborting the check, so every sub task produces exactly one event.
+    let (status, failure_reason, matches) = match run_scan(scanner, sub_task) {
+        Ok(matches) => {
+            println!("datasecurity: sub task succeeded ({} match(es))", matches.len());
+            (ScanStatus::Success, String::new(), matches)
+        }
+        Err(err) => {
+            let reason = format!("{err:#}");
+            eprintln!(
+                "datasecurity: sub task {} failed: {reason}",
+                sub_task.sub_task_id
+            );
+            (ScanStatus::Error, reason, Vec::new())
+        }
+    };
 
     let payload = ScanEventPayload {
         task_id: config.task_id.clone(),
         sub_task_id: sub_task.sub_task_id.clone(),
+        status,
+        failure_reason,
         matches,
     };
 
-    println!(
-        "datasecurity: built scaffold event payload ({} match(es))",
-        payload.matches.len()
-    );
-
     // TODO(DSEC-140): send sdsresult rather than an event
     let payload_json =
         serde_json::to_string(&payload).context("serializing scan event payload")?;
@@ -71,8 +81,8 @@ fn run_sub_task(
     Ok(())
 }
 
-/// Mimics a postgres fetch by returning the sub task's dummy response.
-// TODO(DSEC-139): replace with a real postgres query.
-fn fetch_data(sub_task: &SubTask) -> Value {
-    sub_task.dummy_response.clone()
+/// Fetches the sub task's data and scans it, returning the matches.
+fn run_scan(scanner: &Scanner, sub_task: &SubTask) -> Result<Vec<Match>> {
+    let data = backend::execute_scan(sub_task).context("fetching sub task data")?;
+    scanner.scan(&data).context("scanning sub task data")
 }

pkg/collector/sharedlibrary/rustchecks/checks/datasecurity/src/backend/postgres.rs

@@ -1,7 +1,6 @@
 use anyhow::{Context, Result};
 use core::AgentCheck;
 use serde::Deserialize;
-use serde_json::Value;
 
 use crate::scanning::ScanningRule;
 
@@ -36,12 +35,43 @@ pub struct CheckConfig {
     pub scan_data: Vec<SubTask>,
 }
 
-/// A single scan sub task.
-#[derive(Debug, Deserialize)]
+/// A single scan sub task: a query to run against one data source.
+#[derive(Debug, Default, Deserialize)]
 pub struct SubTask {
     #[serde(default)]
     pub sub_task_id: String,
-    // TODO(DSEC-139): remove dummy response once the postgres backend lands.
+    /// Data-source platform, used to select the backend engine (e.g. `postgres`).
+    #[serde(default)]
+    pub platform: String,
+    #[serde(default)]
+    pub connection: Connection,
+    /// SQL query whose result columns are scanned.
+    #[serde(default)]
+    pub query: String,
+    /// Connect and query timeout in seconds. `0` disables the timeout.
+    #[serde(default = "default_timeout_seconds")]
+    pub timeout_seconds: u64,
+}
+
+fn default_timeout_seconds() -> u64 {
+    30
+}
+
+/// Database connection parameters for a sub task.
+#[derive(Debug, Default, Deserialize)]
+pub struct Connection {
+    #[serde(default)]
+    pub host: String,
+    #[serde(default = "default_port")]
+    pub port: u16,
+    #[serde(default)]
+    pub dbname: String,
     #[serde(default)]
-    pub dummy_response: Value,
+    pub username: String,
+    #[serde(default)]
+    pub password: String,
+}
+
+fn default_port() -> u16 {
+    5432
 }

pkg/collector/sharedlibrary/rustchecks/checks/datasecurity/src/check.rs

@@ -1,5 +1,6 @@
 use core::generate_ffi;
 
+mod backend;
 mod check;
 use check::check;
 

pkg/collector/sharedlibrary/rustchecks/checks/datasecurity/src/config.rs

@@ -1,15 +1,28 @@
 use serde::Serialize;
 
-/// JSON event payload emitted by the check. Kept close to the shape of the real
-/// SDS result so it can be grown into the full payload later.
+/// JSON event payload emitted by the check, one per sub task. Kept close to the
+/// shape of the real SDS result (see `sds_result.proto`) so it can be grown into
+/// the full payload later. On failure, `status` is `ERROR` and `failure_reason`
+/// carries the cause; on success it holds the scan `matches`.
 #[derive(Debug, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub struct ScanEventPayload {
     pub task_id: String,
     pub sub_task_id: String,
+    pub status: ScanStatus,
+    #[serde(skip_serializing_if = "String::is_empty")]
+    pub failure_reason: String,
     pub matches: Vec<Match>,
 }
 
+/// Outcome of a sub task scan, mirroring the proto `ScanStatus`.
+#[derive(Debug, Clone, Copy, Serialize)]
+#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
+pub enum ScanStatus {
+    Success,
+    Error,
+}
+
 #[derive(Debug, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub struct Match {

pkg/collector/sharedlibrary/rustchecks/checks/datasecurity/src/lib.rs

@@ -1,7 +1,7 @@
 checks:
   - id: datasecurity
     crate: datasecurity
-    include_in_build: false
+    include_in_build: true
     platforms:
       - linux
   - id: example