darkwing/server/services/sqlite_services/

mod.rs

mod chrome_time;
use crate::server::error::AppResult;
use anyhow::Context;
use async_trait::async_trait;
use chrome_time::ChromeTime;
use mockall::automock;
use rusqlite::{named_params, Connection, Result as SqliteResult};
use sha2::{Digest, Sha256};
use std::io::{Cursor, Read, Seek, SeekFrom, Write};
use std::sync::Arc;
use tempfile::NamedTempFile;
use tracing::info;
use zip::ZipArchive;

pub type DynSqliteService = Arc<dyn SqliteServiceTrait + Send + Sync>;

const SELECT_COOKIES_QUERY: &str = "SELECT host_key, name, CAST(encrypted_value AS BLOB), path, expires_utc, is_httponly FROM cookies";
const COOKIES_PATH: &str = "Default/Cookies";

#[derive(Debug, Clone)]
struct Cookie {
  pub name: String,
  pub value: String,
  pub domain: String,
  pub path: String,
  pub expiration_date: ChromeTime,
  pub http_only: Option<bool>,
}

#[automock]
#[async_trait]
pub trait SqliteServiceTrait {
  async fn merge_cookies(
    &self,
    old_archive: Option<String>,
    new_archive: String,
  ) -> AppResult<()>;
}

#[derive(Clone)]
pub struct SqliteService {}

impl SqliteService {
  pub fn new() -> Self {
    Self {}
  }

  fn insert_tables(&self, conn: &Connection) -> AppResult<()> {
    // Statement from V24
    conn.execute(
      "CREATE TABLE IF NOT EXISTS cookies(
      creation_utc INTEGER NOT NULL,
      host_key TEXT NOT NULL,
      top_frame_site_key TEXT NOT NULL,
      name TEXT NOT NULL,
      value TEXT NOT NULL,
      encrypted_value BLOB NOT NULL,
      path TEXT NOT NULL,
      expires_utc INTEGER NOT NULL,
      is_secure INTEGER NOT NULL,
      is_httponly INTEGER NOT NULL,
      last_access_utc INTEGER NOT NULL,
      has_expires INTEGER NOT NULL,
      is_persistent INTEGER NOT NULL,
      priority INTEGER NOT NULL,
      samesite INTEGER NOT NULL,
      source_scheme INTEGER NOT NULL,
      source_port INTEGER NOT NULL,
      last_update_utc INTEGER NOT NULL,
      source_type INTEGER NOT NULL,
      has_cross_site_ancestor INTEGER NOT NULL);",
      [],
    )?;

    conn.execute(
      "CREATE UNIQUE INDEX IF NOT EXISTS cookies_unique_index ON cookies(host_key, top_frame_site_key, has_cross_site_ancestor, name, path, source_scheme, source_port)",
      [],
    )?;

    conn.execute(
      "CREATE TABLE IF NOT EXISTS meta(key LONGVARCHAR NOT NULL UNIQUE PRIMARY KEY, value LONGVARCHAR)",
      [],
    )?;

    Ok(())
  }

  fn get_db_version(&self, conn: &Connection) -> SqliteResult<u8> {
    conn.query_row("SELECT value FROM meta WHERE key = 'version'", [], |row| {
      row.get::<_, String>(0).and_then(|v| {
        v.parse::<u8>()
          .map_err(parse_int_error_to_conversion_failure)
      })
    })
  }

  fn encode_cookie(&self, domain: &str, value: &str) -> Vec<u8> {
    let mut hasher = Sha256::new();
    hasher.update(domain.as_bytes());
    let mut res = hasher.finalize().to_vec();
    res.extend(value.as_bytes());
    res
  }

  fn read_cookies(
    &self,
    conn: &Connection,
    version: u8,
  ) -> SqliteResult<Vec<Cookie>> {
    let mut stmt = conn.prepare(SELECT_COOKIES_QUERY)?;

    let cookies = stmt.query_map([], |row| {
      let value = if version >= 24 {
        let encrypted_value: Vec<u8> = row.get(2)?;
        encrypted_value[32..].to_vec()
      } else {
        row.get::<_, Vec<u8>>(2)?
      };
      let value = String::from_utf8(value)
        .map_err(|err| rusqlite::Error::Utf8Error(err.utf8_error()))?;

      let res = Cookie {
        name: row.get(1)?,
        value,
        domain: row.get(0)?,
        path: row.get(3)?,
        expiration_date: row.get(4)?,
        http_only: row.get(5)?,
      };

      Ok(res)
    })?;

    cookies.collect()
  }

  fn write_cookies(
    &self,
    conn: &mut Connection,
    cookies: &[Cookie],
    version: u8,
  ) -> AppResult<()> {
    let tx = conn.transaction().context("Failed to start transaction")?;
    {
      let mut stmt = tx
        .prepare(self.get_insert_query(version))
        .context("Failed to prepare insert statement")?;

      for cookie in cookies {
        let encrypted_value = if version >= 24 {
          self.encode_cookie(&cookie.domain, &cookie.value)
        } else {
          cookie.value.as_bytes().to_vec()
        };

        stmt
          .execute(named_params![
            ":name": cookie.name,
            ":host_key": cookie.domain,
            ":path": cookie.path,
            ":expires_utc": cookie.expiration_date,
            ":encrypted_value": encrypted_value,
            ":is_httponly": cookie.http_only.unwrap_or(false),
          ])
          .context("Failed to execute insert statement")?;
      }
    }
    tx.commit().context("Failed to commit transaction")?;
    Ok(())
  }

  /// Returns an INSERT query for cookies.
  ///
  /// If there is conflict of `host_key, top_frame_site_key, name, path`,
  /// then the existing cookie is updated with new value and expiration date.
  ///
  /// Parameter order:
  /// - name
  /// - host_key
  /// - path
  /// - expires_utc
  /// - encrypted_value
  /// - is_httponly
  fn get_insert_query(&self, version: u8) -> &'static str {
    match version {
      ..=20 => "INSERT INTO cookies (name, host_key, path, expires_utc, encrypted_value, creation_utc, top_frame_site_key, value, is_secure, is_persistent, priority, is_httponly, last_access_utc, has_expires, samesite, source_scheme, source_port, is_same_party, last_update_utc) VALUES (:name, :host_key, :path, :expires_utc, :encrypted_value, 100000, '', '', 1, 1, 1, :is_httponly, 100000, 1, 0, 2, 443, 0, 0) ON CONFLICT (host_key, top_frame_site_key, name, path) DO UPDATE SET encrypted_value = excluded.encrypted_value, expires_utc = excluded.expires_utc",
      21..=22 => "INSERT INTO cookies (name, host_key, path, expires_utc, encrypted_value, creation_utc, top_frame_site_key, value, is_secure, is_persistent, priority, is_httponly, last_access_utc, has_expires, samesite, source_scheme, source_port, last_update_utc) VALUES (:name, :host_key, :path, :expires_utc, :encrypted_value, 100000, '', '', 1, 1, 1, :is_httponly, 100000, 1, 0, 2, 443, 0) ON CONFLICT (host_key, top_frame_site_key, name, path, source_scheme, source_port) DO UPDATE SET encrypted_value = excluded.encrypted_value, expires_utc = excluded.expires_utc",
      23..=100 => "INSERT INTO cookies (name, host_key, path, expires_utc, encrypted_value, creation_utc, top_frame_site_key, value, is_secure, is_persistent, priority, is_httponly, last_access_utc, has_expires, samesite, source_scheme, source_port, last_update_utc, source_type, has_cross_site_ancestor) VALUES (:name, :host_key, :path, :expires_utc, :encrypted_value, 100000, '', '', 1, 1, 1, :is_httponly, 100000, 1, 0, 2, 443, 0, 0, 0) ON CONFLICT (host_key, top_frame_site_key, has_cross_site_ancestor, name, path, source_scheme, source_port) DO UPDATE SET encrypted_value = excluded.encrypted_value, expires_utc = excluded.expires_utc",
      _ => unreachable!("Cannot be present in SQLite version 3"),
    }
  }

  fn extract_cookies(
    &self,
    zip: &mut ZipArchive<Cursor<Vec<u8>>>,
  ) -> AppResult<Vec<u8>> {
    let mut file = zip
      .by_name(COOKIES_PATH)
      .context("Failed to find Cookies file in archive")?;
    let mut content = Vec::new();
    file
      .read_to_end(&mut content)
      .context("Failed to read Cookies file content")?;
    Ok(content)
  }

  fn update_zip_archive(
    &self,
    zip: &mut ZipArchive<Cursor<Vec<u8>>>,
    path: &str,
    content: &[u8],
  ) -> AppResult<Cursor<Vec<u8>>> {
    let mut updated_zip = Cursor::new(Vec::new());
    {
      let mut writer = zip::ZipWriter::new(&mut updated_zip);
      for i in 0..zip.len() {
        let mut file = zip
          .by_index(i)
          .context("Failed to get file from zip archive")?;
        let options = zip::write::FileOptions::default()
          .compression_method(file.compression())
          .unix_permissions(file.unix_mode().unwrap_or(0));

        if file.name() == path {
          writer
            .start_file(path, options)
            .context("Failed to start writing Cookies file")?;
          writer
            .write_all(content)
            .context("Failed to write Cookies content")?;
        } else {
          writer
            .start_file(file.name(), options)
            .context("Failed to start writing file")?;
          std::io::copy(&mut file, &mut writer)
            .context("Failed to copy file content")?;
        }
      }
      writer
        .finish()
        .context("Failed to finish writing zip archive")?;
    }
    Seek::seek(&mut updated_zip, SeekFrom::Start(0))
      .context("Failed to seek to start of updated zip")?;
    Ok(updated_zip)
  }
}

#[async_trait]
impl SqliteServiceTrait for SqliteService {
  async fn merge_cookies(
    &self,
    old_archive: Option<String>,
    new_archive: String,
  ) -> AppResult<()> {
    info!("Merging cookies");

    // read file Default/Cookies from old_archive into memory
    // open connection to that file via rusqlite
    // read all rows from Cookies table
    // read file Default/Cookies from new_archive into memory
    // open connection to that file via rusqlite
    // read all rows from Cookies table
    // merge the two sets of rows, preferring the rows from the new_archive
    // write the merged rows back to the Cookies table
    // close both connections
    // return the new_archive file
    // caveats:
    // - the database is versioned and there is different versions of the
    //   database schema
    // - if any of DBs does not have a Cookies or Meta table, insert it via
    //   separate function (SqliteService::insert_tables
    // - before any insert and select, get the query for current version of
    //   database

    // No need to merge cookies if there is no old archive
    let old_archive = match old_archive {
      Some(path) => path,
      None => return Ok(()),
    };

    // Read zip archives into memory
    let old_content = tokio::fs::read(old_archive).await?;
    let new_content = tokio::fs::read(new_archive.clone()).await?;

    // Create ZipArchive from old and new content
    let mut old_zip = ZipArchive::new(Cursor::new(old_content))
      .context("Failed to create ZipArchive from old content")?;
    let mut new_zip = ZipArchive::new(Cursor::new(new_content.clone()))
      .context("Failed to create ZipArchive from new content")?;

    // Extract Default/Cookies from both archives
    let old_cookies = self
      .extract_cookies(&mut old_zip)
      .context("Failed to extract Cookies from old archive")?;
    let new_cookies = self
      .extract_cookies(&mut new_zip)
      .context("Failed to extract Cookies from new archive")?;

    // Create temporary files for SQLite connections
    let mut old_temp = NamedTempFile::new()
      .context("Failed to create temporary file for old archive")?;
    old_temp
      .write_all(&old_cookies)
      .context("Failed to write old Cookies content to temporary file")?;
    old_temp
      .seek(SeekFrom::Start(0))
      .context("Failed to seek to start of old temporary file")?;

    let mut new_temp = NamedTempFile::new()
      .context("Failed to create temporary file for new archive")?;
    new_temp
      .write_all(&new_cookies)
      .context("Failed to write new Cookies content to temporary file")?;
    new_temp
      .seek(SeekFrom::Start(0))
      .context("Failed to seek to start of new temporary file")?;

    // Open SQLite connections
    let old_conn = Connection::open(old_temp.path())
      .context("Failed to open old SQLite connection")?;
    let mut new_conn = Connection::open(new_temp.path())
      .context("Failed to open new SQLite connection")?;

    // Ensure tables exist in both databases
    self
      .insert_tables(&old_conn)
      .context("Failed to insert tables into old database")?;
    self
      .insert_tables(&new_conn)
      .context("Failed to insert tables into new database")?;

    // Get database versions
    let old_version = self
      .get_db_version(&old_conn)
      .context("Failed to get old database version")?;
    let new_version = self
      .get_db_version(&new_conn)
      .context("Failed to get new database version")?;

    // Read cookies from both databases
    let old_cookies = self
      .read_cookies(&old_conn, old_version)
      .context("Failed to read cookies from old database")?;
    let mut new_cookies = self
      .read_cookies(&new_conn, new_version)
      .context("Failed to read cookies from new database")?;

    // Merge cookies, preferring new ones
    let mut merged_cookies = Vec::new();
    for old_cookie in old_cookies {
      // if any cookie with the same domain and name exists in new_cookies,
      // replace value and expiration date in it
      new_cookies.iter_mut().for_each(|c| {
        if c.domain == old_cookie.domain && c.name == old_cookie.name {
          // c.value = old_cookie.value.clone();
          // c.expiration_date = old_cookie.expiration_date.clone();
          merged_cookies.push(c.clone());
        } else {
          merged_cookies.push(old_cookie.clone());
        }
      });
    }

    // Write merged cookies back to new database
    self
      .write_cookies(&mut new_conn, &merged_cookies, new_version)
      .context("Failed to write merged cookies to new database")?;

    // Close connections
    drop(old_conn);
    drop(new_conn);

    // Read updated content from new_temp
    new_temp
      .seek(SeekFrom::Start(0))
      .context("Failed to seek to start of new temporary file")?;
    let mut updated_content = Vec::new();
    new_temp
      .read_to_end(&mut updated_content)
      .context("Failed to read updated Cookies content from temporary file")?;

    // Update the Cookies file in the new zip archive
    let mut new_zip = ZipArchive::new(Cursor::new(new_content))
      .context("Failed to create ZipArchive from new content")?;
    let updated_zip = self
      .update_zip_archive(&mut new_zip, COOKIES_PATH, &updated_content)
      .context("Failed to update zip archive")?;

    tokio::fs::write(new_archive, updated_zip.into_inner()).await?;

    Ok(())
  }
}

fn parse_int_error_to_conversion_failure(
  err: impl std::error::Error + Send + Sync + 'static,
) -> rusqlite::Error {
  rusqlite::Error::FromSqlConversionFailure(
    0,
    rusqlite::types::Type::Text,
    Box::new(err),
  )
}

#[cfg(test)]
mod tests {
  use super::*;
  use pretty_assertions::assert_eq;
  use sqlx::Row;
  use sqlx::Value;
  use sqlx::ValueRef;
  use std::io::Write;

  #[tokio::test]
  async fn test_merge_cookies() -> AppResult<()> {
    let sqlite_service = Arc::new(SqliteService::new()) as DynSqliteService;

    // old_archive contains Cookies file with 3 cookies from .pixelscan.net, 1
    // cookie from .doubleclick.com (google tracking cookie) and 1 cookies from
    // `test` (which we added from our old anty-sync service). new_archive
    // contains Cookies file with 13 cookies from youtube, codesandbox, google
    // and 1 cookie from `test` (which we added from our old anty-sync service).
    // The merged archive should contain 14 cookies:
    // - 3 cookies from .pixelscan.net
    // - 13 cookies from youtube, codesandbox, google
    // - 1 cookie from .doubleclick.com
    // - 1 cookie from `test`, BUT with updated expiration date and value.
    let old_archive_path = "tests/data/datadir-v19-pixelscan.zip";
    let new_archive_path = "tests/data/datadir-v23-random.zip";
    let new_archive_copy_path =
      "tests/data/datadir-v23-random.zip.COPY_FOR_TESTS";

    // To do not accidentally change the new_archive, we will create a copy of
    // it. All files with COPY_FOR_TESTS suffix are ignored by .gitignore.
    let mut new_archive = tokio::fs::File::open(new_archive_path).await?;
    let mut new_archive_copy =
      tokio::fs::File::create(new_archive_copy_path).await?;
    tokio::io::copy(&mut new_archive, &mut new_archive_copy).await?;

    let _merged_archive = sqlite_service
      .merge_cookies(
        Some(old_archive_path.into()),
        new_archive_copy_path.into(),
      )
      .await
      .unwrap();

    // Ok, so now that our archive is merged, let's check that:
    // 1. The resulting file is valid zip archive
    // 2. The Cookies file is present and is not empty
    // 3. The Cookies file can be read and parsed correctly
    // 4. All cookies from old archive are present in the resulting archive
    // 5. All cookies from new archive are present in the resulting archive
    // 6. The cookies from new archive have new expiration date and value
    // 7. The cookies from old archive have not been changed
    // 8. The size of the resulting archive is more than the size of the old
    //    archive
    // 9. Total number of cookies is correct

    let old_archive_size = tokio::fs::metadata(old_archive_path).await?.len();

    let merged_archive_size =
      tokio::fs::metadata(new_archive_copy_path).await?.len();

    assert!(
      merged_archive_size > old_archive_size,
      "Merged archive is not larger than the old archive"
    );

    // read from "tests/data/datadir-v23-random.zip.COPY_FOR_TESTS"
    let file_content = tokio::fs::read(new_archive_copy_path).await?;
    let mut merged_zip = ZipArchive::new(Cursor::new(file_content))?;

    let mut cookies_file = merged_zip.by_name(COOKIES_PATH).unwrap();
    let mut cookies_content = Vec::new();
    cookies_file.read_to_end(&mut cookies_content)?;

    let mut temp_cookies_file = NamedTempFile::new()?;
    temp_cookies_file.write_all(&cookies_content)?;

    // We will read cookies via sqlx to reduce risk of just doubling faulty code
    // from above.

    let connection = sqlx::sqlite::SqlitePoolOptions::new()
      .connect(&format!(
        "sqlite://{}",
        temp_cookies_file.path().to_string_lossy()
      ))
      .await
      .context("Failed to connect to SQLite database")?;

    let cookies =
      sqlx::query("SELECT name, CAST(encrypted_value AS BLOB) as encrypted_value FROM cookies")
        .fetch_all(&connection)
        .await
        .context("Failed to fetch cookies from SQLite database")?;

    // assert_eq!(cookies.len(), 14, "Total number of cookies is incorrect");

    for cookie in &cookies {
      println!(
        "name: {}, value: {}",
        cookie
          .try_get_raw("name")
          .unwrap()
          .to_owned()
          .try_decode::<String>()
          .unwrap(),
        String::from_utf8(
          cookie
            .try_get_raw("encrypted_value")
            .unwrap()
            .to_owned()
            .try_decode::<Vec<u8>>()
            .unwrap(),
        )
        .unwrap(),
      );
    }

    assert!(
      cookies.iter().any(|c| {
        c.try_get_raw("name")
          .unwrap()
          .to_owned()
          .try_decode::<String>()
          .unwrap()
          == "test"
          && String::from_utf8(
            c.try_get_raw("encrypted_value")
              .unwrap()
              .to_owned()
              .try_decode::<Vec<u8>>()
              .unwrap(),
          )
          .unwrap()
            == "test"
      }),
      "test cookie is not present"
    );

    Ok(())
  }

  #[test]
  fn test_encode_cookie() {
    let service = SqliteService::new();
    let domain = "example.com";
    let value = "test_value";

    let encoded = service.encode_cookie(domain, value);

    // First 32 bytes should be SHA256 of domain
    let mut hasher = Sha256::new();
    hasher.update(domain.as_bytes());
    let expected_hash = hasher.finalize().to_vec();
    assert_eq!(&encoded[..32], &expected_hash[..]);

    // Remaining bytes should be the value
    assert_eq!(&encoded[32..], value.as_bytes());
  }

  #[test]
  fn test_insert_tables() -> AppResult<()> {
    let service = SqliteService::new();
    let temp_file =
      NamedTempFile::new().context("Failed to create temp file")?;
    let conn = Connection::open(temp_file.path())
      .context("Failed to open connection")?;

    service.insert_tables(&conn)?;

    // Verify tables exist
    let tables: Vec<String> = conn
      .prepare("SELECT name FROM sqlite_master WHERE type='table'")?
      .query_map([], |row| row.get(0))?
      .collect::<Result<_, _>>()?;

    assert!(tables.contains(&"cookies".to_string()));
    assert!(tables.contains(&"meta".to_string()));

    // Verify cookies table structure
    let columns: Vec<String> = conn
      .prepare("PRAGMA table_info(cookies)")?
      .query_map([], |row| row.get(1))?
      .collect::<Result<_, _>>()?;

    assert!(columns.contains(&"host_key".to_string()));
    assert!(columns.contains(&"name".to_string()));
    assert!(columns.contains(&"encrypted_value".to_string()));

    Ok(())
  }

  #[test]
  fn test_read_and_write_cookies() -> AppResult<()> {
    let service = SqliteService::new();
    let temp_file =
      NamedTempFile::new().context("Failed to create temp file")?;
    let mut conn = Connection::open(temp_file.path())
      .context("Failed to open connection")?;

    service.insert_tables(&conn)?;

    // Insert version
    conn
      .execute("INSERT INTO meta (key, value) VALUES ('version', '24')", [])?;

    let test_cookies = vec![
      Cookie {
        name: "cookie1".to_string(),
        value: "value1".to_string(),
        domain: "example.com".to_string(),
        path: "/".to_string(),
        expiration_date: ChromeTime::from_unix_timestamp(13317017516).unwrap(), /* 2392-01-01 */
        http_only: Some(true),
      },
      Cookie {
        name: "cookie2".to_string(),
        value: "value2".to_string(),
        domain: "test.com".to_string(),
        path: "/path".to_string(),
        expiration_date: ChromeTime::from_unix_timestamp(13317017516).unwrap(), /* 2392-01-01 */
        http_only: Some(false),
      },
    ];

    // Write cookies
    service.write_cookies(&mut conn, &test_cookies, 24)?;

    // Read cookies back
    let read_cookies = service.read_cookies(&conn, 24)?;

    assert_eq!(read_cookies.len(), test_cookies.len());

    // Compare cookies (excluding encrypted values which will be different due
    // to encoding)
    for (original, read) in test_cookies.iter().zip(read_cookies.iter()) {
      assert_eq!(original.name, read.name);
      assert_eq!(original.domain, read.domain);
      assert_eq!(original.path, read.path);
      assert_eq!(original.expiration_date, read.expiration_date);
      assert_eq!(original.http_only, read.http_only);
    }

    Ok(())
  }

  #[test]
  fn test_get_insert_query() {
    let service = SqliteService::new();

    // Test for different versions
    let v20_query = service.get_insert_query(20);
    assert!(v20_query.contains("is_same_party"));
    assert!(!v20_query.contains("source_type"));

    let v22_query = service.get_insert_query(22);
    assert!(!v22_query.contains("is_same_party"));
    assert!(!v22_query.contains("source_type"));

    let v24_query = service.get_insert_query(24);
    assert!(!v24_query.contains("is_same_party"));
    assert!(v24_query.contains("source_type"));
  }
}