Merge pull request #2 from renegade-fi/joey/compliance-db-schema

compliance: compliance-server: db: Define migrations and query db

Cargo.toml

@@ -29,3 +29,6 @@ renegade-circuits = { package = "circuits", git = "https://github.com/renegade-f
 renegade-circuit-types = { package = "circuit-types", git = "https://github.com/renegade-fi/renegade.git" }
 renegade-crypto = { git = "https://github.com/renegade-fi/renegade.git" }
 renegade-util = { package = "util", git = "https://github.com/renegade-fi/renegade.git" }
+
+# === Misc Dependencies === #
+tracing = "0.1"

compliance/compliance-api/src/lib.rs

@@ -1,5 +1,16 @@
+//! The API for the compliance server
+
+#![deny(missing_docs)]
+#![deny(clippy::missing_docs_in_private_items)]
+#![deny(unsafe_code)]
+#![deny(clippy::needless_pass_by_value)]
+#![deny(clippy::needless_pass_by_ref_mut)]
+
 use serde::{Deserialize, Serialize};
 
+/// The API endpoint for screening an address for compliance
+pub const WALLET_SCREEN_PATH: &str = "/v0/check-compliance";
+
 /// The response type for a compliance check
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct ComplianceCheckResponse {
@@ -13,5 +24,6 @@ pub enum ComplianceStatus {
     /// The wallet is compliant
     Compliant,
     /// The wallet is not compliant
-    NotCompliant,
+    #[allow(missing_docs)]
+    NotCompliant { reason: String },
 }

compliance/compliance-server/Cargo.toml

@@ -9,9 +9,16 @@ http-body-util = "0.1.0"
 warp = "0.3"
 compliance-api = { path = "../compliance-api" }
 
+# === Database === #
+diesel = { version = "2.2", features = ["postgres", "r2d2"] }
+
 # === Renegade Dependencies === #
 renegade-util = { workspace = true }
 
 # === Misc === #
 clap = { version = "4.5", features = ["derive"] }
+reqwest = { version = "0.12", features = ["json"] }
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
 tokio = { version = "1.37", features = ["full"] }
+tracing = { workspace = true }

compliance/compliance-server/diesel.toml

@@ -0,0 +1,9 @@
+# For documentation on how to configure this file,
+# see https://diesel.rs/guides/configuring-diesel-cli
+
+[print_schema]
+file = "src/schema.rs"
+custom_type_derives = ["diesel::query_builder::QueryId", "Clone"]
+
+[migrations_directory]
+dir = "./migrations"

compliance/compliance-server/migrations/00000000000000_diesel_initial_setup/down.sql

@@ -0,0 +1,6 @@
+-- This file was automatically created by Diesel to setup helper functions
+-- and other internal bookkeeping. This file is safe to edit, any future
+-- changes will be added to existing projects as new migrations.
+
+DROP FUNCTION IF EXISTS diesel_manage_updated_at(_tbl regclass);
+DROP FUNCTION IF EXISTS diesel_set_updated_at();

compliance/compliance-server/migrations/00000000000000_diesel_initial_setup/up.sql

@@ -0,0 +1,36 @@
+-- This file was automatically created by Diesel to setup helper functions
+-- and other internal bookkeeping. This file is safe to edit, any future
+-- changes will be added to existing projects as new migrations.
+
+
+
+
+-- Sets up a trigger for the given table to automatically set a column called
+-- `updated_at` whenever the row is modified (unless `updated_at` was included
+-- in the modified columns)
+--
+-- # Example
+--
+-- ```sql
+-- CREATE TABLE users (id SERIAL PRIMARY KEY, updated_at TIMESTAMP NOT NULL DEFAULT NOW());
+--
+-- SELECT diesel_manage_updated_at('users');
+-- ```
+CREATE OR REPLACE FUNCTION diesel_manage_updated_at(_tbl regclass) RETURNS VOID AS $$
+BEGIN
+    EXECUTE format('CREATE TRIGGER set_updated_at BEFORE UPDATE ON %s
+                    FOR EACH ROW EXECUTE PROCEDURE diesel_set_updated_at()', _tbl);
+END;
+$$ LANGUAGE plpgsql;
+
+CREATE OR REPLACE FUNCTION diesel_set_updated_at() RETURNS trigger AS $$
+BEGIN
+    IF (
+        NEW IS DISTINCT FROM OLD AND
+        NEW.updated_at IS NOT DISTINCT FROM OLD.updated_at
+    ) THEN
+        NEW.updated_at := current_timestamp;
+    END IF;
+    RETURN NEW;
+END;
+$$ LANGUAGE plpgsql;

compliance/compliance-server/migrations/2024-06-28-231756_create_compliance_table/down.sql

@@ -0,0 +1,2 @@
+-- Drop the wallet compliance table
+DROP TABLE IF EXISTS wallet_compliance;

compliance/compliance-server/migrations/2024-06-28-231756_create_compliance_table/up.sql

@@ -0,0 +1,9 @@
+-- Create a table for caching wallet compliance information
+CREATE TABLE IF NOT EXISTS wallet_compliance (
+    address TEXT PRIMARY KEY,
+    is_compliant BOOLEAN NOT NULL,
+    risk_level TEXT NOT NULL,
+    reason TEXT NOT NULL,
+    created_at TIMESTAMP NOT NULL DEFAULT NOW(),
+    expires_at TIMESTAMP NOT NULL DEFAULT NOW() + INTERVAL '1 year'
+);

compliance/compliance-server/src/chainalysis_api.rs

@@ -0,0 +1,107 @@
+//! Helpers for interacting with the chainalysis API
+
+use serde::{Deserialize, Serialize};
+use tracing::warn;
+
+use crate::{db::ComplianceEntry, error::ComplianceServerError};
+
+// -------------
+// | API Types |
+// -------------
+
+/// The base URL for the chainalysis entities API
+const CHAINALYSIS_API_BASE: &str = "https://api.chainalysis.com/api/risk/v2/entities";
+/// The header name for the auth token
+const TOKEN_HEADER: &str = "Token";
+
+/// The register address request body
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct RegisterAddressRequest {
+    /// The address to register
+    pub address: String,
+}
+
+/// The response to a risk assessment query
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct RiskAssessmentResponse {
+    /// The address that the assessment is for
+    pub address: String,
+    /// The risk assessment status
+    pub risk: String,
+    /// The reason for the risk assessment
+    #[serde(rename = "riskReason")]
+    pub risk_reason: Option<String>,
+}
+
+impl RiskAssessmentResponse {
+    /// Get a compliance entry from the risk assessment
+    pub fn as_compliance_entry(self) -> ComplianceEntry {
+        // We allow low and medium risk entries, high and severe are marked
+        // non-compliant
+        let compliant = match self.risk.as_str() {
+            "Low" | "Medium" => true,
+            "High" | "Severe" => false,
+            x => {
+                // For now we don't block on an unknown assessment, this should be unreachable
+                warn!("Unexpected risk assessment: {x}");
+                true
+            },
+        };
+
+        let risk_reason = self.risk_reason.unwrap_or_default();
+        ComplianceEntry::new(self.address, compliant, self.risk, risk_reason)
+    }
+}
+
+// ---------------
+// | Client Impl |
+// ---------------
+
+/// Query chainalysis for the compliance status of a wallet
+pub async fn query_chainalysis(
+    wallet_address: &str,
+    chainalysis_api_key: &str,
+) -> Result<ComplianceEntry, ComplianceServerError> {
+    // 1. Register the wallet
+    register_addr(wallet_address, chainalysis_api_key).await?;
+
+    // 2. Query the risk assessment
+    let risk_assessment = query_risk_assessment(wallet_address, chainalysis_api_key).await?;
+    Ok(risk_assessment.as_compliance_entry())
+}
+
+/// Register a wallet with chainalysis
+async fn register_addr(
+    wallet_address: &str,
+    chainalysis_api_key: &str,
+) -> Result<(), ComplianceServerError> {
+    let body = RegisterAddressRequest { address: wallet_address.to_string() };
+    let client = reqwest::Client::new();
+    client
+        .post(CHAINALYSIS_API_BASE)
+        .header(TOKEN_HEADER, chainalysis_api_key)
+        .json(&body)
+        .send()
+        .await?
+        .error_for_status()?;
+
+    Ok(())
+}
+
+/// Query the risk assessment from chainalysis
+async fn query_risk_assessment(
+    wallet_address: &str,
+    chainalysis_api_key: &str,
+) -> Result<RiskAssessmentResponse, ComplianceServerError> {
+    let url = format!("{CHAINALYSIS_API_BASE}/{wallet_address}");
+    let client = reqwest::Client::new();
+    let resp = client
+        .get(url)
+        .header(TOKEN_HEADER, chainalysis_api_key)
+        .send()
+        .await?
+        .error_for_status()?;
+
+    let risk_assessment: RiskAssessmentResponse = resp.json().await?;
+    Ok(risk_assessment)
+}

compliance/compliance-server/src/db.rs

@@ -0,0 +1,83 @@
+//! Database helpers for the server
+
+use std::time::{Duration, SystemTime};
+
+use compliance_api::ComplianceStatus;
+use diesel::{ExpressionMethods, Insertable, PgConnection, QueryDsl, Queryable, RunQueryDsl};
+use renegade_util::err_str;
+
+use crate::{
+    error::ComplianceServerError,
+    schema::{
+        wallet_compliance,
+        wallet_compliance::dsl::{address as address_col, wallet_compliance as compliance_table},
+    },
+};
+
+/// The default expiration duration for a compliance entry
+const DEFAULT_EXPIRATION_DURATION: Duration = Duration::from_days(365);
+
+// ----------
+// | Models |
+// ----------
+
+/// A compliance entry for a wallet
+#[derive(Debug, Clone, Queryable, Insertable)]
+#[table_name = "wallet_compliance"]
+#[allow(missing_docs)]
+pub struct ComplianceEntry {
+    pub address: String,
+    pub is_compliant: bool,
+    pub risk_level: String,
+    pub reason: String,
+    pub created_at: SystemTime,
+    pub expires_at: SystemTime,
+}
+
+impl ComplianceEntry {
+    /// Create a new entry from a risk assessment
+    pub fn new(address: String, is_compliant: bool, risk_level: String, reason: String) -> Self {
+        let created_at = SystemTime::now();
+        let expires_at = created_at + DEFAULT_EXPIRATION_DURATION;
+        ComplianceEntry { address, is_compliant, risk_level, reason, created_at, expires_at }
+    }
+
+    /// Get the compliance status for an entry
+    pub fn compliance_status(&self) -> ComplianceStatus {
+        if self.is_compliant {
+            ComplianceStatus::Compliant
+        } else {
+            ComplianceStatus::NotCompliant { reason: self.reason.clone() }
+        }
+    }
+}
+
+// -----------
+// | Queries |
+// -----------
+
+/// Get a compliance entry by address
+pub fn get_compliance_entry(
+    address: &str,
+    conn: &mut PgConnection,
+) -> Result<Option<ComplianceEntry>, ComplianceServerError> {
+    let query = compliance_table
+        .filter(address_col.eq(address))
+        .load::<ComplianceEntry>(conn)
+        .map_err(err_str!(ComplianceServerError::Db))?;
+
+    Ok(query.first().cloned())
+}
+
+/// Insert a compliance entry into the database
+pub fn insert_compliance_entry(
+    entry: ComplianceEntry,
+    conn: &mut PgConnection,
+) -> Result<(), ComplianceServerError> {
+    diesel::insert_into(compliance_table)
+        .values(entry)
+        .execute(conn)
+        .map_err(err_str!(ComplianceServerError::Db))?;
+
+    Ok(())
+}

compliance/compliance-server/src/error.rs

@@ -6,12 +6,26 @@ use warp::reject::Reject;
 
 /// The error type emitted by the compliance server
 #[derive(Debug, Clone)]
-pub enum ComplianceServerError {}
+pub enum ComplianceServerError {
+    /// An error with a database query
+    Db(String),
+    /// An error with the chainalysis API
+    Chainalysis(String),
+}
 
 impl Display for ComplianceServerError {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "ComplianceServerError")
+        match self {
+            ComplianceServerError::Db(e) => write!(f, "Database error: {}", e),
+            ComplianceServerError::Chainalysis(e) => write!(f, "Chainalysis error: {}", e),
+        }
     }
 }
 impl Error for ComplianceServerError {}
 impl Reject for ComplianceServerError {}
+
+impl From<reqwest::Error> for ComplianceServerError {
+    fn from(e: reqwest::Error) -> Self {
+        ComplianceServerError::Chainalysis(e.to_string())
+    }
+}