From 45833ac245f97fa4b08293466375e773d385e9fd Mon Sep 17 00:00:00 2001 From: benthecarman Date: Thu, 21 Sep 2023 16:24:30 -0500 Subject: [PATCH 1/7] Respond to tony's updates --- Cargo.lock | 22 +-- Cargo.toml | 3 +- migrations/2023-09-18-225828_baseline/up.sql | 8 +- .../down.sql | 3 - .../up.sql | 3 - src/auth.rs | 14 +- src/kv.rs | 87 ++++++++++++ src/main.rs | 55 ++++---- src/migration.rs | 22 ++- src/models/mod.rs | 91 +++++------- src/models/schema.rs | 2 +- src/routes.rs | 133 ++++++++++++------ 12 files changed, 266 insertions(+), 177 deletions(-) delete mode 100644 migrations/2023-09-20-043550_change-default-timestamp/down.sql delete mode 100644 migrations/2023-09-20-043550_change-default-timestamp/up.sql create mode 100644 src/kv.rs diff --git a/Cargo.lock b/Cargo.lock index b1ef5e7..f22c4f6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -294,7 +294,6 @@ dependencies = [ "num-integer", "num-traits", "pq-sys", - "r2d2", ] [[package]] @@ -986,17 +985,6 @@ dependencies = [ "proc-macro2", ] -[[package]] -name = "r2d2" -version = "0.8.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93" -dependencies = [ - "log", - "parking_lot", - "scheduled-thread-pool", -] - [[package]] name = "rand_core" version = "0.6.4" @@ -1119,15 +1107,6 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ad4cc8da4ef723ed60bced201181d83791ad433213d8c24efffda1eec85d741" -[[package]] -name = "scheduled-thread-pool" -version = "0.2.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19" -dependencies = [ - "parking_lot", -] - [[package]] name = "scopeguard" version = "1.2.0" @@ -1571,6 +1550,7 @@ version = "0.1.0" dependencies = [ "anyhow", "axum", + "base64 0.13.1", "chrono", "diesel", "diesel_migrations", diff --git a/Cargo.toml b/Cargo.toml index c7da229..27aec61 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -6,8 +6,9 @@ edition = "2021" [dependencies] anyhow = "1.0" axum = { version = "0.6.16", features = ["headers"] } +base64 = "0.13.1" chrono = { version = "0.4.26", features = ["serde"] } -diesel = { version = "2.1", features = ["postgres", "r2d2", "chrono", "numeric"] } +diesel = { version = "2.1", features = ["postgres", "chrono", "numeric"] } diesel_migrations = "2.1.0" dotenv = "0.15.0" futures = "0.3.28" diff --git a/migrations/2023-09-18-225828_baseline/up.sql b/migrations/2023-09-18-225828_baseline/up.sql index a901b6d..0acc8d0 100644 --- a/migrations/2023-09-18-225828_baseline/up.sql +++ b/migrations/2023-09-18-225828_baseline/up.sql @@ -2,10 +2,10 @@ CREATE TABLE vss_db ( store_id TEXT NOT NULL CHECK (store_id != ''), key TEXT NOT NULL, - value TEXT, + value bytea, version BIGINT NOT NULL, - created_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL, - updated_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL, + created_date TIMESTAMP DEFAULT '2023-07-13'::TIMESTAMP NOT NULL, + updated_date TIMESTAMP DEFAULT '2023-07-13'::TIMESTAMP NOT NULL, PRIMARY KEY (store_id, key) ); @@ -40,4 +40,4 @@ CREATE TRIGGER tr_set_dates_after_insert CREATE TRIGGER tr_set_dates_after_update BEFORE UPDATE ON vss_db FOR EACH ROW - EXECUTE FUNCTION set_updated_date(); \ No newline at end of file + EXECUTE FUNCTION set_updated_date(); diff --git a/migrations/2023-09-20-043550_change-default-timestamp/down.sql b/migrations/2023-09-20-043550_change-default-timestamp/down.sql deleted file mode 100644 index 96dcc8a..0000000 --- a/migrations/2023-09-20-043550_change-default-timestamp/down.sql +++ /dev/null @@ -1,3 +0,0 @@ -ALTER TABLE vss_db - ALTER COLUMN created_date SET DEFAULT CURRENT_TIMESTAMP, - ALTER COLUMN updated_date SET DEFAULT CURRENT_TIMESTAMP; diff --git a/migrations/2023-09-20-043550_change-default-timestamp/up.sql b/migrations/2023-09-20-043550_change-default-timestamp/up.sql deleted file mode 100644 index 1df7ce1..0000000 --- a/migrations/2023-09-20-043550_change-default-timestamp/up.sql +++ /dev/null @@ -1,3 +0,0 @@ -ALTER TABLE vss_db - ALTER COLUMN created_date SET DEFAULT '2023-07-13'::TIMESTAMP, - ALTER COLUMN updated_date SET DEFAULT '2023-07-13'::TIMESTAMP; diff --git a/src/auth.rs b/src/auth.rs index 79a42db..6928465 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -1,5 +1,5 @@ use crate::State; -use axum::http::StatusCode; +use axum::http::{HeaderMap, StatusCode}; use jwt_compact::alg::Es256k; use jwt_compact::{AlgorithmExt, TimeOptions, Token, UntrustedToken}; use log::error; @@ -7,12 +7,20 @@ use secp256k1::PublicKey; use serde::{Deserialize, Serialize}; use sha2::Sha256; -pub(crate) fn verify_token(token: &str, state: &State) -> Result { +pub(crate) fn verify_token( + token: &str, + state: &State, + headers: &HeaderMap, +) -> Result { let es256k1 = Es256k::::new(state.secp.clone()); validate_jwt_from_user(token, state.auth_key, &es256k1).map_err(|e| { error!("Unauthorized: {e}"); - (StatusCode::UNAUTHORIZED, format!("Unauthorized: {e}")) + ( + StatusCode::UNAUTHORIZED, + headers.clone(), + format!("Unauthorized: {e}"), + ) }) } diff --git a/src/kv.rs b/src/kv.rs new file mode 100644 index 0000000..3eba388 --- /dev/null +++ b/src/kv.rs @@ -0,0 +1,87 @@ +use core::fmt; +use serde::de::Visitor; +use serde::*; + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct KeyValue { + pub key: String, + pub value: ByteData, + pub version: i64, +} + +impl KeyValue { + pub fn new(key: String, value: Vec, version: i64) -> KeyValue { + KeyValue { + key, + value: ByteData(value), + version, + } + } +} + +#[derive(Debug, Clone)] +pub struct ByteData(pub Vec); + +impl Serialize for ByteData { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + self.0.serialize(serializer) + } +} + +impl<'de> Deserialize<'de> for ByteData { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + struct ByteDataVisitor; + + impl<'de> Visitor<'de> for ByteDataVisitor { + type Value = ByteData; + + fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result { + formatter.write_str("a Vec or a base64 encoded string") + } + + fn visit_str(self, v: &str) -> Result + where + E: de::Error, + { + let decoded = + base64::decode(v).map_err(|err| de::Error::custom(err.to_string()))?; + Ok(ByteData(decoded)) + } + + fn visit_seq(self, seq: S) -> Result + where + S: de::SeqAccess<'de>, + { + let vec = Vec::::deserialize(de::value::SeqAccessDeserializer::new(seq))?; + Ok(ByteData(vec)) + } + } + + deserializer.deserialize_any(ByteDataVisitor) + } +} + +// need this for backwards compat for now + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct KeyValueOld { + pub key: String, + pub value: String, + pub version: i64, +} + +impl From for KeyValueOld { + fn from(kv: KeyValue) -> Self { + KeyValueOld { + key: kv.key, + value: base64::encode(kv.value.0), + version: kv.version, + } + } +} diff --git a/src/main.rs b/src/main.rs index 874454d..c41efb7 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,15 +1,15 @@ use crate::models::MIGRATIONS; use crate::routes::*; -use axum::http::{Method, StatusCode, Uri}; +use axum::headers::Origin; +use axum::http::{HeaderMap, StatusCode, Uri}; use axum::routing::{get, post, put}; -use axum::{http, Extension, Router}; -use diesel::r2d2::{ConnectionManager, Pool}; -use diesel::PgConnection; +use axum::{Extension, Router, TypedHeader}; +use diesel::{Connection, PgConnection}; use diesel_migrations::MigrationHarness; use secp256k1::{All, PublicKey, Secp256k1}; -use tower_http::cors::{Any, CorsLayer}; mod auth; +mod kv; mod migration; mod models; mod routes; @@ -28,7 +28,7 @@ const ALLOWED_LOCALHOST: &str = "http://127.0.0.1:"; #[derive(Clone)] pub struct State { - db_pool: Pool>, + pg_url: String, auth_key: PublicKey, secp: Secp256k1, } @@ -52,15 +52,10 @@ async fn main() -> anyhow::Result<()> { let auth_key = PublicKey::from_slice(&auth_key_bytes)?; // DB management - let manager = ConnectionManager::::new(&pg_url); - let db_pool = Pool::builder() - .max_size(16) - .test_on_check_out(true) - .build(manager) - .expect("Could not build connection pool"); + let mut connection = PgConnection::establish(&pg_url).unwrap(); + // TODO not sure if code should handle the migration, could be dangerous with multiple instances // run migrations - let mut connection = db_pool.get()?; connection .run_pending_migrations(MIGRATIONS) .expect("migrations could not run"); @@ -68,7 +63,7 @@ async fn main() -> anyhow::Result<()> { let secp = Secp256k1::new(); let state = State { - db_pool, + pg_url, auth_key, secp, }; @@ -80,20 +75,14 @@ async fn main() -> anyhow::Result<()> { let server_router = Router::new() .route("/health-check", get(health_check)) .route("/getObject", post(get_object)) + .route("/v2/getObject", post(get_object_v2)) .route("/putObjects", put(put_objects)) + .route("/v2/putObjects", put(put_objects)) .route("/listKeyVersions", post(list_key_versions)) + .route("/v2/listKeyVersions", post(list_key_versions)) .route("/migration", get(migration::migration)) .fallback(fallback) - .layer(Extension(state.clone())) - .layer( - CorsLayer::new() - .allow_origin(Any) - .allow_headers(vec![ - http::header::CONTENT_TYPE, - http::header::AUTHORIZATION, - ]) - .allow_methods([Method::GET, Method::POST, Method::PUT, Method::OPTIONS]), - ); + .layer(Extension(state)); let server = axum::Server::bind(&addr).serve(server_router.into_make_service()); @@ -113,6 +102,20 @@ async fn main() -> anyhow::Result<()> { Ok(()) } -async fn fallback(uri: Uri) -> (StatusCode, String) { - (StatusCode::NOT_FOUND, format!("No route for {uri}")) +async fn fallback( + origin: Option>, + uri: Uri, +) -> (StatusCode, HeaderMap, String) { + let origin = match validate_cors(origin) { + Ok(origin) => origin, + Err((status, headers, msg)) => return (status, headers, msg), + }; + + let headers = create_cors_headers(&origin); + + ( + StatusCode::NOT_FOUND, + headers, + format!("No route for {uri}"), + ) } diff --git a/src/migration.rs b/src/migration.rs index 444761b..f5b6c50 100644 --- a/src/migration.rs +++ b/src/migration.rs @@ -6,7 +6,7 @@ use axum::headers::Authorization; use axum::http::StatusCode; use axum::{Extension, Json, TypedHeader}; use chrono::{DateTime, NaiveDateTime, Utc}; -use diesel::Connection; +use diesel::{Connection, PgConnection}; use log::{error, info}; use serde::{Deserialize, Deserializer}; use serde_json::json; @@ -18,7 +18,7 @@ pub struct Item { pub key: String, #[serde(default)] pub value: String, - pub version: u32, + pub version: i64, #[serde(default)] #[serde(deserialize_with = "deserialize_datetime_opt")] @@ -66,7 +66,7 @@ pub async fn migration_impl(admin_key: String, state: &State) -> anyhow::Result< let mut finished = false; - let mut conn = state.db_pool.get()?; + let mut conn = PgConnection::establish(&state.pg_url).unwrap(); info!("Starting migration"); while !finished { @@ -78,24 +78,20 @@ pub async fn migration_impl(admin_key: String, state: &State) -> anyhow::Result< .post(&url) .set("x-api-key", &admin_key) .send_string(&payload.to_string())?; - let values: Vec = resp.into_json()?; + let items: Vec = resp.into_json()?; // Insert values into DB conn.transaction::<_, anyhow::Error, _>(|conn| { - for value in values.iter() { - VssItem::put_item( - conn, - &value.store_id, - &value.key, - &value.value, - value.version as u64, - )?; + for item in items.iter() { + if let Ok(value) = base64::decode(&item.value) { + VssItem::put_item(conn, &item.store_id, &item.key, &value, item.version)?; + } } Ok(()) })?; - if values.len() < limit { + if items.len() < limit { finished = true; } else { offset += limit; diff --git a/src/models/mod.rs b/src/models/mod.rs index 5d7e831..2dd8fec 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -1,7 +1,7 @@ -use crate::routes::KeyValue; +use crate::kv::KeyValue; use diesel::prelude::*; use diesel::sql_query; -use diesel::sql_types::{BigInt, Text}; +use diesel::sql_types::{BigInt, Bytea, Text}; use diesel_migrations::{embed_migrations, EmbeddedMigrations}; use schema::vss_db; use serde::{Deserialize, Serialize}; @@ -26,29 +26,17 @@ pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!(); pub struct VssItem { pub store_id: String, pub key: String, - pub value: Option, + pub value: Option>, pub version: i64, created_date: chrono::NaiveDateTime, updated_date: chrono::NaiveDateTime, } -#[derive(Insertable, AsChangeset)] -#[diesel(table_name = vss_db)] -pub struct NewVssItem { - pub store_id: String, - pub key: String, - pub value: Option, - pub version: i64, -} - impl VssItem { pub fn into_kv(self) -> Option { - self.value.map(|value| KeyValue { - key: self.key, - value, - version: self.version as u64, - }) + self.value + .map(|value| KeyValue::new(self.key, value, self.version)) } pub fn get_item( @@ -67,20 +55,13 @@ impl VssItem { conn: &mut PgConnection, store_id: &str, key: &str, - value: &str, - version: u64, + value: &[u8], + version: i64, ) -> anyhow::Result<()> { - // safely convert u64 to i64 - let version = if version >= i64::MAX as u64 { - i64::MAX - } else { - version as i64 - }; - sql_query(include_str!("put_item.sql")) .bind::(store_id) .bind::(key) - .bind::(value) + .bind::(value) .bind::(version) .execute(conn)?; @@ -111,7 +92,6 @@ impl VssItem { mod test { use super::*; use crate::State; - use diesel::r2d2::{ConnectionManager, Pool}; use diesel::{Connection, PgConnection, RunQueryDsl}; use diesel_migrations::MigrationHarness; use secp256k1::Secp256k1; @@ -121,16 +101,9 @@ mod test { fn init_state() -> State { dotenv::dotenv().ok(); - let url = std::env::var("DATABASE_URL").expect("DATABASE_URL must be set"); - let manager = ConnectionManager::::new(url); - let db_pool = Pool::builder() - .max_size(16) - .test_on_check_out(true) - .build(manager) - .expect("Could not build connection pool"); - - // run migrations - let mut connection = db_pool.get().unwrap(); + let pg_url = std::env::var("DATABASE_URL").expect("DATABASE_URL must be set"); + let mut connection = PgConnection::establish(&pg_url).unwrap(); + connection .run_pending_migrations(MIGRATIONS) .expect("migrations could not run"); @@ -140,14 +113,14 @@ mod test { let secp = Secp256k1::new(); State { - db_pool, + pg_url, auth_key, secp, } } fn clear_database(state: &State) { - let conn = &mut state.db_pool.get().unwrap(); + let mut conn = PgConnection::establish(&state.pg_url).unwrap(); conn.transaction::<_, anyhow::Error, _>(|conn| { diesel::delete(vss_db::table).execute(conn)?; @@ -163,22 +136,22 @@ mod test { let store_id = "test_store_id"; let key = "test"; - let value = "test_value"; + let value = [1, 2, 3, 4, 5]; let version = 0; - let mut conn = state.db_pool.get().unwrap(); - VssItem::put_item(&mut conn, store_id, key, value, version).unwrap(); + let mut conn = PgConnection::establish(&state.pg_url).unwrap(); + VssItem::put_item(&mut conn, store_id, key, &value, version).unwrap(); let versions = VssItem::list_key_versions(&mut conn, store_id, None).unwrap(); assert_eq!(versions.len(), 1); assert_eq!(versions[0].0, key); - assert_eq!(versions[0].1, version as i64); + assert_eq!(versions[0].1, version); - let new_value = "new_value"; + let new_value = [6, 7, 8, 9, 10]; let new_version = version + 1; - VssItem::put_item(&mut conn, store_id, key, new_value, new_version).unwrap(); + VssItem::put_item(&mut conn, store_id, key, &new_value, new_version).unwrap(); let item = VssItem::get_item(&mut conn, store_id, key) .unwrap() @@ -187,7 +160,7 @@ mod test { assert_eq!(item.store_id, store_id); assert_eq!(item.key, key); assert_eq!(item.value.unwrap(), new_value); - assert_eq!(item.version, new_version as i64); + assert_eq!(item.version, new_version); clear_database(&state); } @@ -199,11 +172,11 @@ mod test { let store_id = "max_test_store_id"; let key = "max_test"; - let value = "test_value"; - let version = u32::MAX as u64; + let value = [1, 2, 3, 4, 5]; + let version = u32::MAX as i64; - let mut conn = state.db_pool.get().unwrap(); - VssItem::put_item(&mut conn, store_id, key, value, version).unwrap(); + let mut conn = PgConnection::establish(&state.pg_url).unwrap(); + VssItem::put_item(&mut conn, store_id, key, &value, version).unwrap(); let item = VssItem::get_item(&mut conn, store_id, key) .unwrap() @@ -213,9 +186,9 @@ mod test { assert_eq!(item.key, key); assert_eq!(item.value.unwrap(), value); - let new_value = "new_value"; + let new_value = [6, 7, 8, 9, 10]; - VssItem::put_item(&mut conn, store_id, key, new_value, version).unwrap(); + VssItem::put_item(&mut conn, store_id, key, &new_value, version).unwrap(); let item = VssItem::get_item(&mut conn, store_id, key) .unwrap() @@ -236,13 +209,13 @@ mod test { let store_id = "list_kv_test_store_id"; let key = "kv_test"; let key1 = "other_kv_test"; - let value = "test_value"; + let value = [1, 2, 3, 4, 5]; let version = 0; - let mut conn = state.db_pool.get().unwrap(); - VssItem::put_item(&mut conn, store_id, key, value, version).unwrap(); + let mut conn = PgConnection::establish(&state.pg_url).unwrap(); + VssItem::put_item(&mut conn, store_id, key, &value, version).unwrap(); - VssItem::put_item(&mut conn, store_id, key1, value, version).unwrap(); + VssItem::put_item(&mut conn, store_id, key1, &value, version).unwrap(); let versions = VssItem::list_key_versions(&mut conn, store_id, None).unwrap(); assert_eq!(versions.len(), 2); @@ -250,12 +223,12 @@ mod test { let versions = VssItem::list_key_versions(&mut conn, store_id, Some("kv")).unwrap(); assert_eq!(versions.len(), 1); assert_eq!(versions[0].0, key); - assert_eq!(versions[0].1, version as i64); + assert_eq!(versions[0].1, version); let versions = VssItem::list_key_versions(&mut conn, store_id, Some("other")).unwrap(); assert_eq!(versions.len(), 1); assert_eq!(versions[0].0, key1); - assert_eq!(versions[0].1, version as i64); + assert_eq!(versions[0].1, version); clear_database(&state); } diff --git a/src/models/schema.rs b/src/models/schema.rs index 0857dac..de1aed0 100644 --- a/src/models/schema.rs +++ b/src/models/schema.rs @@ -4,7 +4,7 @@ diesel::table! { vss_db (store_id, key) { store_id -> Text, key -> Text, - value -> Nullable, + value -> Nullable, version -> Int8, created_date -> Timestamp, updated_date -> Timestamp, diff --git a/src/routes.rs b/src/routes.rs index d9bfc87..31f4ec2 100644 --- a/src/routes.rs +++ b/src/routes.rs @@ -1,30 +1,29 @@ use crate::auth::verify_token; +use crate::kv::{KeyValue, KeyValueOld}; use crate::models::VssItem; use crate::{State, ALLOWED_LOCALHOST, ALLOWED_ORIGINS, ALLOWED_SUBDOMAIN}; use axum::headers::authorization::Bearer; -use axum::headers::{Authorization, Origin}; +use axum::headers::{Authorization, HeaderMap, Origin}; +use axum::http::header::{ + ACCESS_CONTROL_ALLOW_HEADERS, ACCESS_CONTROL_ALLOW_METHODS, ACCESS_CONTROL_ALLOW_ORIGIN, + CONTENT_TYPE, +}; use axum::http::StatusCode; use axum::{Extension, Json, TypedHeader}; -use diesel::Connection; +use diesel::{Connection, PgConnection}; use log::{debug, error, trace}; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; -#[derive(Debug, Clone, Serialize, Deserialize)] -pub struct KeyValue { - pub key: String, - pub value: String, - pub version: u64, -} - -macro_rules! check_store_id { - ($payload:ident, $store_id:expr) => { +macro_rules! ensure_store_id { + ($payload:ident, $store_id:expr, $headers:ident) => { match $payload.store_id { None => $payload.store_id = Some($store_id), Some(ref id) => { if id != &$store_id { return Err(( StatusCode::UNAUTHORIZED, + $headers, format!("Unauthorized: store_id mismatch"), )); } @@ -43,11 +42,10 @@ pub async fn get_object_impl( req: GetObjectRequest, state: &State, ) -> anyhow::Result> { - let mut conn = state.db_pool.get()?; - trace!("get_object_impl: {req:?}"); let store_id = req.store_id.expect("must have"); + let mut conn = PgConnection::establish(&state.pg_url).unwrap(); let item = VssItem::get_item(&mut conn, &store_id, &req.key)?; Ok(item.and_then(|i| i.into_kv())) @@ -58,16 +56,41 @@ pub async fn get_object( TypedHeader(token): TypedHeader>, Extension(state): Extension, Json(mut payload): Json, -) -> Result>, (StatusCode, String)> { +) -> Result<(HeaderMap, Json>), (StatusCode, HeaderMap, String)> { debug!("get_object: {payload:?}"); - validate_cors(origin)?; - let store_id = verify_token(token.token(), &state)?; + let origin = validate_cors(origin)?; + let mut headers = create_cors_headers(&origin); + headers.insert(CONTENT_TYPE, "application/json".parse().unwrap()); + + let store_id = verify_token(token.token(), &state, &headers)?; - check_store_id!(payload, store_id); + ensure_store_id!(payload, store_id, headers); + + match get_object_impl(payload, &state).await { + Ok(Some(res)) => Ok((headers, Json(Some(res.into())))), + Ok(None) => Ok((headers, Json(None))), + Err(e) => Err(handle_anyhow_error(e, headers)), + } +} + +pub async fn get_object_v2( + origin: Option>, + TypedHeader(token): TypedHeader>, + Extension(state): Extension, + Json(mut payload): Json, +) -> Result>, (StatusCode, HeaderMap, String)> { + debug!("get_object v2: {payload:?}"); + let origin = validate_cors(origin)?; + let mut headers = create_cors_headers(&origin); + headers.insert(CONTENT_TYPE, "application/json".parse().unwrap()); + + let store_id = verify_token(token.token(), &state, &headers)?; + + ensure_store_id!(payload, store_id, headers); match get_object_impl(payload, &state).await { Ok(res) => Ok(Json(res)), - Err(e) => Err(handle_anyhow_error(e)), + Err(e) => Err(handle_anyhow_error(e, headers)), } } @@ -87,10 +110,10 @@ pub async fn put_objects_impl(req: PutObjectsRequest, state: &State) -> anyhow:: let store_id = req.store_id.expect("must have"); - let mut conn = state.db_pool.get()?; + let mut conn = PgConnection::establish(&state.pg_url).unwrap(); conn.transaction(|conn| { for kv in req.transaction_items { - VssItem::put_item(conn, &store_id, &kv.key, &kv.value, kv.version)?; + VssItem::put_item(conn, &store_id, &kv.key, &kv.value.0, kv.version)?; } Ok(()) @@ -102,15 +125,18 @@ pub async fn put_objects( TypedHeader(token): TypedHeader>, Extension(state): Extension, Json(mut payload): Json, -) -> Result, (StatusCode, String)> { - validate_cors(origin)?; - let store_id = verify_token(token.token(), &state)?; +) -> Result<(HeaderMap, Json<()>), (StatusCode, HeaderMap, String)> { + let origin = validate_cors(origin)?; + let mut headers = create_cors_headers(&origin); + headers.insert(CONTENT_TYPE, "application/json".parse().unwrap()); + + let store_id = verify_token(token.token(), &state, &headers)?; - check_store_id!(payload, store_id); + ensure_store_id!(payload, store_id, headers); match put_objects_impl(payload, &state).await { - Ok(res) => Ok(Json(res)), - Err(e) => Err(handle_anyhow_error(e)), + Ok(res) => Ok((headers, Json(res))), + Err(e) => Err(handle_anyhow_error(e, headers)), } } @@ -126,11 +152,10 @@ pub async fn list_key_versions_impl( req: ListKeyVersionsRequest, state: &State, ) -> anyhow::Result> { - let mut conn = state.db_pool.get()?; - // todo pagination let store_id = req.store_id.expect("must have"); + let mut conn = PgConnection::establish(&state.pg_url).unwrap(); let versions = VssItem::list_key_versions(&mut conn, &store_id, req.key_prefix.as_deref())?; let json = versions @@ -151,15 +176,18 @@ pub async fn list_key_versions( TypedHeader(token): TypedHeader>, Extension(state): Extension, Json(mut payload): Json, -) -> Result>, (StatusCode, String)> { - validate_cors(origin)?; - let store_id = verify_token(token.token(), &state)?; +) -> Result<(HeaderMap, Json>), (StatusCode, HeaderMap, String)> { + let origin = validate_cors(origin)?; + let mut headers = create_cors_headers(&origin); + headers.insert(CONTENT_TYPE, "application/json".parse().unwrap()); + + let store_id = verify_token(token.token(), &state, &headers)?; - check_store_id!(payload, store_id); + ensure_store_id!(payload, store_id, headers); match list_key_versions_impl(payload, &state).await { - Ok(res) => Ok(Json(res)), - Err(e) => Err(handle_anyhow_error(e)), + Ok(res) => Ok((headers, Json(res))), + Err(e) => Err(handle_anyhow_error(e, headers)), } } @@ -167,26 +195,45 @@ pub async fn health_check() -> Result, (StatusCode, String)> { Ok(Json(())) } -fn validate_cors(origin: Option>) -> Result<(), (StatusCode, String)> { +pub fn validate_cors( + origin: Option>, +) -> Result { if let Some(TypedHeader(origin)) = origin { if origin.is_null() { - return Ok(()); + return Ok("*".to_string()); } let origin_str = origin.to_string(); - if !ALLOWED_ORIGINS.contains(&origin_str.as_str()) - && !origin_str.ends_with(ALLOWED_SUBDOMAIN) - && !origin_str.starts_with(ALLOWED_LOCALHOST) + if ALLOWED_ORIGINS.contains(&origin_str.as_str()) + || origin_str.ends_with(ALLOWED_SUBDOMAIN) + || origin_str.starts_with(ALLOWED_LOCALHOST) { + return Ok(origin_str); + } else { + let headers = create_cors_headers("*"); // The origin is not in the allowed list block the request - return Err((StatusCode::NOT_FOUND, String::new())); + return Err((StatusCode::NOT_FOUND, headers, String::new())); } } - Ok(()) + Ok("*".to_string()) +} + +pub fn create_cors_headers(origin: &str) -> HeaderMap { + let mut headers = HeaderMap::new(); + headers.insert(ACCESS_CONTROL_ALLOW_ORIGIN, origin.parse().unwrap()); + headers.insert( + ACCESS_CONTROL_ALLOW_METHODS, + "GET, POST, PUT, DELETE, OPTIONS".parse().unwrap(), + ); + headers.insert(ACCESS_CONTROL_ALLOW_HEADERS, "*".parse().unwrap()); + headers } -pub(crate) fn handle_anyhow_error(err: anyhow::Error) -> (StatusCode, String) { +pub(crate) fn handle_anyhow_error( + err: anyhow::Error, + headers: HeaderMap, +) -> (StatusCode, HeaderMap, String) { error!("Error: {err:?}"); - (StatusCode::BAD_REQUEST, format!("{err}")) + (StatusCode::BAD_REQUEST, headers, format!("{err}")) } From c65735f696e4a0837b88e253710087364ad1545c Mon Sep 17 00:00:00 2001 From: benthecarman Date: Fri, 22 Sep 2023 20:07:49 -0500 Subject: [PATCH 2/7] wip --- Cargo.lock | 458 ++++++++++-------- Cargo.toml | 5 +- Dockerfile | 4 +- migrations/.keep | 0 .../down.sql | 6 - .../up.sql | 36 -- .../2023-09-18-225828_baseline/down.sql | 9 - migrations/2023-09-18-225828_baseline/up.sql | 43 -- src/auth.rs | 14 +- src/main.rs | 77 +-- src/migration.rs | 22 +- src/models/migration_baseline.sql | 76 +++ src/models/mod.rs | 284 ++++++----- src/models/schema.rs | 12 - src/routes.rs | 129 ++--- 15 files changed, 629 insertions(+), 546 deletions(-) delete mode 100644 migrations/.keep delete mode 100644 migrations/00000000000000_diesel_initial_setup/down.sql delete mode 100644 migrations/00000000000000_diesel_initial_setup/up.sql delete mode 100644 migrations/2023-09-18-225828_baseline/down.sql delete mode 100644 migrations/2023-09-18-225828_baseline/up.sql create mode 100644 src/models/migration_baseline.sql delete mode 100644 src/models/schema.rs diff --git a/Cargo.lock b/Cargo.lock index f22c4f6..715eadf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -19,9 +19,9 @@ checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe" [[package]] name = "aho-corasick" -version = "1.1.0" +version = "1.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f2135563fb5c609d2b2b87c1e8ce7bc41b0b45430fa9661f457981503dd5bf0" +checksum = "ea5d730647d4fadd988536d06fecce94b7b4f2a7efdae548f1cf4b63205518ab" dependencies = [ "memchr", ] @@ -147,19 +147,6 @@ version = "1.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8c3c1a368f70d6cf7302d78f8f7093da241fb8e8807c05cc9e51a125895a6d5b" -[[package]] -name = "bigdecimal" -version = "0.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "454bca3db10617b88b566f205ed190aedb0e0e6dd4cad61d3988a72e8c5594cb" -dependencies = [ - "autocfg", - "libm", - "num-bigint", - "num-integer", - "num-traits", -] - [[package]] name = "bitcoin-private" version = "0.1.0" @@ -278,56 +265,6 @@ dependencies = [ "typenum", ] -[[package]] -name = "diesel" -version = "2.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d98235fdc2f355d330a8244184ab6b4b33c28679c0b4158f63138e51d6cf7e88" -dependencies = [ - "bigdecimal", - "bitflags 2.4.0", - "byteorder", - "chrono", - "diesel_derives", - "itoa", - "num-bigint", - "num-integer", - "num-traits", - "pq-sys", -] - -[[package]] -name = "diesel_derives" -version = "2.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e054665eaf6d97d1e7125512bb2d35d07c73ac86cc6920174cb42d1ab697a554" -dependencies = [ - "diesel_table_macro_syntax", - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "diesel_migrations" -version = "2.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6036b3f0120c5961381b570ee20a02432d7e2d27ea60de9578799cf9156914ac" -dependencies = [ - "diesel", - "migrations_internals", - "migrations_macros", -] - -[[package]] -name = "diesel_table_macro_syntax" -version = "0.1.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fc5557efc453706fed5e4fa85006fe9817c224c3f480a34c7e5959fd700921c5" -dependencies = [ - "syn", -] - [[package]] name = "digest" version = "0.10.7" @@ -358,12 +295,6 @@ dependencies = [ "termcolor", ] -[[package]] -name = "equivalent" -version = "1.0.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5" - [[package]] name = "errno" version = "0.3.3" @@ -385,6 +316,18 @@ dependencies = [ "libc", ] +[[package]] +name = "fallible-iterator" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4443176a9f2c162692bd3d352d745ef9413eec5782a80d8fd6f8a1ac692a07f7" + +[[package]] +name = "finl_unicode" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fcfdc7a0362c9f4444381a9e697c79d435fe65b52a37466fc2c1184cee9edc6" + [[package]] name = "flate2" version = "1.0.27" @@ -401,6 +344,21 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" +[[package]] +name = "foreign-types" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6f339eb8adc052cd2ca78910fda869aefa38d22d5cb648e6485e4d3fc06f3b1" +dependencies = [ + "foreign-types-shared", +] + +[[package]] +name = "foreign-types-shared" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00b0228411908ca8685dba7fc2cdd70ec9990a6e753e89b6ac91a84c40fbaf4b" + [[package]] name = "form_urlencoded" version = "1.2.0" @@ -509,6 +467,17 @@ dependencies = [ "version_check", ] +[[package]] +name = "getrandom" +version = "0.2.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be4136b2a15dd319360be1c07d9933517ccf0be8f16bf62a3bee4f0d618df427" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] + [[package]] name = "gimli" version = "0.28.0" @@ -521,20 +490,13 @@ version = "1.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eabb4a44450da02c90444cf74558da904edde8fb4e9035a9a6a4e15445af0bd7" -[[package]] -name = "hashbrown" -version = "0.14.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2c6201b9ff9fd90a5a3bac2e56a830d0caa509576f0e503818ee82c181b3437a" - [[package]] name = "headers" -version = "0.3.8" +version = "0.3.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f3e372db8e5c0d213e0cd0b9be18be2aca3d44cf2fe30a9d46a65581cd454584" +checksum = "06683b93020a07e3dbcf5f8c0f6d40080d725bea7936fc01ad345c01b97dc270" dependencies = [ - "base64 0.13.1", - "bitflags 1.3.2", + "base64 0.21.4", "bytes", "headers-core", "http", @@ -554,9 +516,9 @@ dependencies = [ [[package]] name = "hermit-abi" -version = "0.3.2" +version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "443144c8cdadd93ebf52ddb4056d257f5b52c04d3c804e657d19eb73fc33668b" +checksum = "d77f7ec81a6d05a3abb01ab6eb7590f6083d08449fe5a1c8b1e620283546ccb7" [[package]] name = "hex" @@ -675,16 +637,6 @@ dependencies = [ "unicode-normalization", ] -[[package]] -name = "indexmap" -version = "2.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d5477fe2230a79769d8dc68e0eabf5437907c0457a5614a9e8dddb67f65eb65d" -dependencies = [ - "equivalent", - "hashbrown", -] - [[package]] name = "is-terminal" version = "0.4.9" @@ -745,12 +697,6 @@ version = "0.2.148" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9cdc71e17332e86d2e1d38c1f99edcb6288ee11b815fb1a4b049eaa2114d369b" -[[package]] -name = "libm" -version = "0.2.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f7012b1bbb0719e1097c47611d3898568c546d597c2e74d66f6087edd5233ff4" - [[package]] name = "linux-raw-sys" version = "0.4.7" @@ -775,36 +721,25 @@ checksum = "b5e6163cb8c49088c2c36f57875e58ccd8c87c7427f7fbd50ea6710b2f3f2e8f" [[package]] name = "matchit" -version = "0.7.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ed1202b2a6f884ae56f04cff409ab315c5ce26b5e58d7412e484f01fd52f52ef" - -[[package]] -name = "memchr" -version = "2.6.3" +version = "0.7.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f232d6ef707e1956a43342693d2a31e72989554d58299d7a88738cc95b0d35c" +checksum = "0e7465ac9959cc2b1404e8e2367b43684a6d13790fe23056cc8c6c5a6b7bcb94" [[package]] -name = "migrations_internals" -version = "2.1.0" +name = "md-5" +version = "0.10.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f23f71580015254b020e856feac3df5878c2c7a8812297edd6c0a485ac9dada" +checksum = "d89e7ee0cfbedfc4da3340218492196241d89eefb6dab27de5df917a6d2e78cf" dependencies = [ - "serde", - "toml", + "cfg-if", + "digest", ] [[package]] -name = "migrations_macros" -version = "2.1.0" +name = "memchr" +version = "2.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cce3325ac70e67bbab5bd837a31cae01f1a6db64e0e744a33cb03a543469ef08" -dependencies = [ - "migrations_internals", - "proc-macro2", - "quote", -] +checksum = "8f232d6ef707e1956a43342693d2a31e72989554d58299d7a88738cc95b0d35c" [[package]] name = "mime" @@ -832,27 +767,6 @@ dependencies = [ "windows-sys", ] -[[package]] -name = "num-bigint" -version = "0.4.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "608e7659b5c3d7cba262d894801b9ec9d00de989e8a82bd4bef91d08da45cdc0" -dependencies = [ - "autocfg", - "num-integer", - "num-traits", -] - -[[package]] -name = "num-integer" -version = "0.1.45" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "225d3389fb3509a24c93f5c29eb6bde2586b98d9f016636dff58d7c6f7569cd9" -dependencies = [ - "autocfg", - "num-traits", -] - [[package]] name = "num-traits" version = "0.2.16" @@ -887,6 +801,44 @@ version = "1.18.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dd8b5dd2ae5ed71462c540258bedcb51965123ad7e7ccf4b9a8cafaa4a63576d" +[[package]] +name = "openssl" +version = "0.10.57" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bac25ee399abb46215765b1cb35bc0212377e58a061560d8b29b024fd0430e7c" +dependencies = [ + "bitflags 2.4.0", + "cfg-if", + "foreign-types", + "libc", + "once_cell", + "openssl-macros", + "openssl-sys", +] + +[[package]] +name = "openssl-macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "openssl-sys" +version = "0.9.93" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db4d56a4c0478783083cfafcc42493dd4a981d41669da64b4572a2a089b51b1d" +dependencies = [ + "cc", + "libc", + "pkg-config", + "vcpkg", +] + [[package]] name = "parking_lot" version = "0.12.1" @@ -916,6 +868,24 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b2a4787296e9989611394c33f193f676704af1686e70b8f8033ab5ba9a35a94" +[[package]] +name = "phf" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ade2d8b8f33c7333b51bcf0428d37e217e9f32192ae4772156f65063b8ce03dc" +dependencies = [ + "phf_shared", +] + +[[package]] +name = "phf_shared" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90fcb95eef784c2ac79119d1dd819e162b5da872ce6f3c3abe1e8ca1c082f72b" +dependencies = [ + "siphasher", +] + [[package]] name = "pin-project" version = "1.1.3" @@ -949,14 +919,60 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" [[package]] -name = "pq-sys" -version = "0.4.8" +name = "pkg-config" +version = "0.3.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26072860ba924cbfa98ea39c8c19b4dd6a4a25423dbdf219c1eca91aa0cf6964" + +[[package]] +name = "postgres-openssl" +version = "0.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "31c0052426df997c0cbd30789eb44ca097e3541717a7b8fa36b1c464ee7edebd" +checksum = "1de0ea6504e07ca78355a6fb88ad0f36cafe9e696cbc6717f16a207f3a60be72" dependencies = [ - "vcpkg", + "futures", + "openssl", + "tokio", + "tokio-openssl", + "tokio-postgres", +] + +[[package]] +name = "postgres-protocol" +version = "0.6.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49b6c5ef183cd3ab4ba005f1ca64c21e8bd97ce4699cfea9e8d9a2c4958ca520" +dependencies = [ + "base64 0.21.4", + "byteorder", + "bytes", + "fallible-iterator", + "hmac", + "md-5", + "memchr", + "rand", + "sha2", + "stringprep", +] + +[[package]] +name = "postgres-types" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8d2234cdee9408b523530a9b6d2d6b373d1db34f6a8e51dc03ded1828d7fb67c" +dependencies = [ + "bytes", + "chrono", + "fallible-iterator", + "postgres-protocol", ] +[[package]] +name = "ppv-lite86" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de" + [[package]] name = "pretty_env_logger" version = "0.5.0" @@ -985,11 +1001,35 @@ dependencies = [ "proc-macro2", ] +[[package]] +name = "rand" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +dependencies = [ + "libc", + "rand_chacha", + "rand_core", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core", +] + [[package]] name = "rand_core" version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom", +] [[package]] name = "redox_syscall" @@ -1052,9 +1092,9 @@ checksum = "d626bb9dae77e28219937af045c257c28bfd3f69333c512553507f5f9798cb76" [[package]] name = "rustix" -version = "0.38.13" +version = "0.38.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7db8590df6dfcd144d22afd1b83b36c21a18d7cbc1dc4bb5295a8712e9eb662" +checksum = "747c788e9ce8e92b12cd485c49ddf90723550b654b32508f979b71a7b1ecda4f" dependencies = [ "bitflags 2.4.0", "errno", @@ -1071,7 +1111,7 @@ checksum = "cd8d6c9f025a446bc4d18ad9632e69aec8f287aa84499ee335599fabd20c3fd8" dependencies = [ "log", "ring", - "rustls-webpki 0.101.5", + "rustls-webpki 0.101.6", "sct", ] @@ -1087,9 +1127,9 @@ dependencies = [ [[package]] name = "rustls-webpki" -version = "0.101.5" +version = "0.101.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "45a27e3b59326c16e23d30aeb7a36a24cc0d29e71d68ff611cdfb4a01d013bed" +checksum = "3c7d5dece342910d9ba34d259310cae3e0154b873b35408b787b59bce53d34fe" dependencies = [ "ring", "untrusted", @@ -1193,15 +1233,6 @@ dependencies = [ "serde", ] -[[package]] -name = "serde_spanned" -version = "0.6.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "96426c9936fd7a0124915f9185ea1d20aa9445cc9821142f0a73bc9207a2e186" -dependencies = [ - "serde", -] - [[package]] name = "serde_urlencoded" version = "0.7.1" @@ -1216,9 +1247,9 @@ dependencies = [ [[package]] name = "sha1" -version = "0.10.5" +version = "0.10.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f04293dc80c3993519f2d7f6f511707ee7094fe0c6d3406feb330cdb3540eba3" +checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" dependencies = [ "cfg-if", "cpufeatures", @@ -1245,6 +1276,12 @@ dependencies = [ "libc", ] +[[package]] +name = "siphasher" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38b58827f4464d87d377d175e90bf58eb00fd8716ff0a62f80356b5e61555d0d" + [[package]] name = "slab" version = "0.4.9" @@ -1256,9 +1293,9 @@ dependencies = [ [[package]] name = "smallvec" -version = "1.11.0" +version = "1.11.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "62bb4feee49fdd9f707ef802e22365a35de4b7b299de4763d44bfea899442ff9" +checksum = "942b4a808e05215192e39f4ab80813e599068285906cc91aa64f923db842bd5a" [[package]] name = "socket2" @@ -1286,6 +1323,17 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6e63cff320ae2c57904679ba7cb63280a3dc4613885beafb148ee7bf9aa9042d" +[[package]] +name = "stringprep" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb41d74e231a107a1b4ee36bd1214b11285b77768d2e3824aedafa988fd36ee6" +dependencies = [ + "finl_unicode", + "unicode-bidi", + "unicode-normalization", +] + [[package]] name = "subtle" version = "2.5.0" @@ -1364,37 +1412,55 @@ dependencies = [ ] [[package]] -name = "toml" -version = "0.7.8" +name = "tokio-openssl" +version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dd79e69d3b627db300ff956027cc6c3798cef26d22526befdfcd12feeb6d2257" +checksum = "c08f9ffb7809f1b20c1b398d92acf4cc719874b3b2b2d9ea2f09b4a80350878a" dependencies = [ - "serde", - "serde_spanned", - "toml_datetime", - "toml_edit", + "futures-util", + "openssl", + "openssl-sys", + "tokio", ] [[package]] -name = "toml_datetime" -version = "0.6.3" +name = "tokio-postgres" +version = "0.7.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7cda73e2f1397b1262d6dfdcef8aafae14d1de7748d66822d3bfeeb6d03e5e4b" +checksum = "d340244b32d920260ae7448cb72b6e238bddc3d4f7603394e7dd46ed8e48f5b8" dependencies = [ - "serde", + "async-trait", + "byteorder", + "bytes", + "fallible-iterator", + "futures-channel", + "futures-util", + "log", + "parking_lot", + "percent-encoding", + "phf", + "pin-project-lite", + "postgres-protocol", + "postgres-types", + "rand", + "socket2 0.5.4", + "tokio", + "tokio-util", + "whoami", ] [[package]] -name = "toml_edit" -version = "0.19.15" +name = "tokio-util" +version = "0.7.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1b5bb770da30e5cbfde35a2d7b9b8a2c4b8ef89548a7a6aeab5c9a576e3e7421" +checksum = "1d68074620f57a0b21594d9735eb2e98ab38b17f80d3fcb189fca266771ca60d" dependencies = [ - "indexmap", - "serde", - "serde_spanned", - "toml_datetime", - "winnow", + "bytes", + "futures-core", + "futures-sink", + "pin-project-lite", + "tokio", + "tracing", ] [[package]] @@ -1552,19 +1618,20 @@ dependencies = [ "axum", "base64 0.13.1", "chrono", - "diesel", - "diesel_migrations", "dotenv", "futures", "hex", "jwt-compact", "log", + "openssl", + "postgres-openssl", "pretty_env_logger", "secp256k1", "serde", "serde_json", "sha2", "tokio", + "tokio-postgres", "tower-http", "ureq", ] @@ -1657,6 +1724,16 @@ dependencies = [ "rustls-webpki 0.100.3", ] +[[package]] +name = "whoami" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22fc3756b8a9133049b26c7f61ab35416c130e8c09b660f5b3958b446f52cc50" +dependencies = [ + "wasm-bindgen", + "web-sys", +] + [[package]] name = "winapi" version = "0.3.9" @@ -1675,9 +1752,9 @@ checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" [[package]] name = "winapi-util" -version = "0.1.5" +version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "70ec6ce85bb158151cae5e5c87f95a8e97d2c0c4b001223f33a334e3ce5de178" +checksum = "f29e6f9198ba0d26b4c9f07dbe6f9ed633e1f3d5b8b414090084349e46a52596" dependencies = [ "winapi", ] @@ -1763,15 +1840,6 @@ version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" -[[package]] -name = "winnow" -version = "0.5.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7c2e3184b9c4e92ad5167ca73039d0c42476302ab603e2fec4487511f38ccefc" -dependencies = [ - "memchr", -] - [[package]] name = "zeroize" version = "1.6.0" diff --git a/Cargo.toml b/Cargo.toml index 27aec61..a3928dd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,8 +8,6 @@ anyhow = "1.0" axum = { version = "0.6.16", features = ["headers"] } base64 = "0.13.1" chrono = { version = "0.4.26", features = ["serde"] } -diesel = { version = "2.1", features = ["postgres", "chrono", "numeric"] } -diesel_migrations = "2.1.0" dotenv = "0.15.0" futures = "0.3.28" hex = "0.4.3" @@ -21,6 +19,9 @@ sha2 = { version = "0.10", default-features = false } serde = { version = "^1.0", features = ["derive"] } serde_json = "1.0.67" tokio = { version = "1.12.0", features = ["full"] } +tokio-postgres = { version = "0.7.10", features = ["with-chrono-0_4"] } tower-http = { version = "0.4.0", features = ["cors"] } ureq = { version = "2.5.0", features = ["json"] } +postgres-openssl = "0.5.0" +openssl = "0.10.57" diff --git a/Dockerfile b/Dockerfile index 326562d..e96c19c 100644 --- a/Dockerfile +++ b/Dockerfile @@ -8,9 +8,9 @@ RUN --mount=type=cache,target=/usr/local/cargo,from=rust:latest,source=/usr/loca cargo build --release && mv ./target/release/vss-rs ./vss-rs # Runtime image -FROM debian:bookworm-slim +FROM debian:bullseye-slim -RUN apt update && apt install -y openssl libpq-dev pkg-config libc6 +RUN apt update && apt install -y openssl libpq-dev pkg-config libc6 openssl libssl-dev libpq5 ca-certificates # Run as "app" user RUN useradd -ms /bin/bash app diff --git a/migrations/.keep b/migrations/.keep deleted file mode 100644 index e69de29..0000000 diff --git a/migrations/00000000000000_diesel_initial_setup/down.sql b/migrations/00000000000000_diesel_initial_setup/down.sql deleted file mode 100644 index a9f5260..0000000 --- a/migrations/00000000000000_diesel_initial_setup/down.sql +++ /dev/null @@ -1,6 +0,0 @@ --- 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(); diff --git a/migrations/00000000000000_diesel_initial_setup/up.sql b/migrations/00000000000000_diesel_initial_setup/up.sql deleted file mode 100644 index d68895b..0000000 --- a/migrations/00000000000000_diesel_initial_setup/up.sql +++ /dev/null @@ -1,36 +0,0 @@ --- 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; diff --git a/migrations/2023-09-18-225828_baseline/down.sql b/migrations/2023-09-18-225828_baseline/down.sql deleted file mode 100644 index 46feaba..0000000 --- a/migrations/2023-09-18-225828_baseline/down.sql +++ /dev/null @@ -1,9 +0,0 @@ --- 1. Drop the triggers -DROP TRIGGER IF EXISTS tr_set_dates_after_insert ON vss_db; -DROP TRIGGER IF EXISTS tr_set_dates_after_update ON vss_db; - --- 2. Drop the trigger functions -DROP FUNCTION IF EXISTS set_created_date(); -DROP FUNCTION IF EXISTS set_updated_date(); - -DROP TABLE IF EXISTS vss_db; diff --git a/migrations/2023-09-18-225828_baseline/up.sql b/migrations/2023-09-18-225828_baseline/up.sql deleted file mode 100644 index 0acc8d0..0000000 --- a/migrations/2023-09-18-225828_baseline/up.sql +++ /dev/null @@ -1,43 +0,0 @@ -CREATE TABLE vss_db -( - store_id TEXT NOT NULL CHECK (store_id != ''), - key TEXT NOT NULL, - value bytea, - version BIGINT NOT NULL, - created_date TIMESTAMP DEFAULT '2023-07-13'::TIMESTAMP NOT NULL, - updated_date TIMESTAMP DEFAULT '2023-07-13'::TIMESTAMP NOT NULL, - PRIMARY KEY (store_id, key) -); - --- triggers to set dates automatically, generated by ChatGPT - --- Function to set created_date and updated_date during INSERT -CREATE OR REPLACE FUNCTION set_created_date() -RETURNS TRIGGER AS $$ -BEGIN - NEW.created_date := CURRENT_TIMESTAMP; - NEW.updated_date := CURRENT_TIMESTAMP; -RETURN NEW; -END; -$$ LANGUAGE plpgsql; - --- Function to set updated_date during UPDATE -CREATE OR REPLACE FUNCTION set_updated_date() -RETURNS TRIGGER AS $$ -BEGIN - NEW.updated_date := CURRENT_TIMESTAMP; -RETURN NEW; -END; -$$ LANGUAGE plpgsql; - --- Trigger for INSERT operation on vss_db -CREATE TRIGGER tr_set_dates_after_insert - BEFORE INSERT ON vss_db - FOR EACH ROW - EXECUTE FUNCTION set_created_date(); - --- Trigger for UPDATE operation on vss_db -CREATE TRIGGER tr_set_dates_after_update - BEFORE UPDATE ON vss_db - FOR EACH ROW - EXECUTE FUNCTION set_updated_date(); diff --git a/src/auth.rs b/src/auth.rs index 6928465..79a42db 100644 --- a/src/auth.rs +++ b/src/auth.rs @@ -1,5 +1,5 @@ use crate::State; -use axum::http::{HeaderMap, StatusCode}; +use axum::http::StatusCode; use jwt_compact::alg::Es256k; use jwt_compact::{AlgorithmExt, TimeOptions, Token, UntrustedToken}; use log::error; @@ -7,20 +7,12 @@ use secp256k1::PublicKey; use serde::{Deserialize, Serialize}; use sha2::Sha256; -pub(crate) fn verify_token( - token: &str, - state: &State, - headers: &HeaderMap, -) -> Result { +pub(crate) fn verify_token(token: &str, state: &State) -> Result { let es256k1 = Es256k::::new(state.secp.clone()); validate_jwt_from_user(token, state.auth_key, &es256k1).map_err(|e| { error!("Unauthorized: {e}"); - ( - StatusCode::UNAUTHORIZED, - headers.clone(), - format!("Unauthorized: {e}"), - ) + (StatusCode::UNAUTHORIZED, format!("Unauthorized: {e}")) }) } diff --git a/src/main.rs b/src/main.rs index c41efb7..109a2f7 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,12 +1,14 @@ -use crate::models::MIGRATIONS; use crate::routes::*; use axum::headers::Origin; -use axum::http::{HeaderMap, StatusCode, Uri}; +use axum::http::{request::Parts, HeaderValue, Method, StatusCode, Uri}; use axum::routing::{get, post, put}; -use axum::{Extension, Router, TypedHeader}; -use diesel::{Connection, PgConnection}; -use diesel_migrations::MigrationHarness; +use axum::{http, Extension, Router, TypedHeader}; +use openssl::ssl::{SslConnector, SslMethod}; +use postgres_openssl::MakeTlsConnector; use secp256k1::{All, PublicKey, Secp256k1}; +use std::sync::Arc; +use tokio_postgres::Client; +use tower_http::cors::{AllowOrigin, CorsLayer}; mod auth; mod kv; @@ -28,9 +30,9 @@ const ALLOWED_LOCALHOST: &str = "http://127.0.0.1:"; #[derive(Clone)] pub struct State { - pg_url: String, - auth_key: PublicKey, - secp: Secp256k1, + pub client: Arc, + pub auth_key: PublicKey, + pub secp: Secp256k1, } #[tokio::main] @@ -51,19 +53,24 @@ async fn main() -> anyhow::Result<()> { let auth_key_bytes = hex::decode(auth_key)?; let auth_key = PublicKey::from_slice(&auth_key_bytes)?; - // DB management - let mut connection = PgConnection::establish(&pg_url).unwrap(); + let builder = SslConnector::builder(SslMethod::tls())?; + let connector = MakeTlsConnector::new(builder.build()); - // TODO not sure if code should handle the migration, could be dangerous with multiple instances - // run migrations - connection - .run_pending_migrations(MIGRATIONS) - .expect("migrations could not run"); + // Connect to the database. + let (client, connection) = tokio_postgres::connect(&pg_url, connector).await?; + + // The connection object performs the actual communication with the database, + // so spawn it off to run on its own. + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("db connection error: {e}"); + } + }); let secp = Secp256k1::new(); let state = State { - pg_url, + client: Arc::new(client), auth_key, secp, }; @@ -82,6 +89,26 @@ async fn main() -> anyhow::Result<()> { .route("/v2/listKeyVersions", post(list_key_versions)) .route("/migration", get(migration::migration)) .fallback(fallback) + .layer( + CorsLayer::new() + .allow_origin(AllowOrigin::predicate( + |origin: &HeaderValue, _request_parts: &Parts| { + let Ok(origin) = origin.to_str() else { + return false; + }; + + valid_origin(origin) + }, + )) + .allow_headers([http::header::CONTENT_TYPE, http::header::AUTHORIZATION]) + .allow_methods([ + Method::GET, + Method::POST, + Method::PUT, + Method::DELETE, + Method::OPTIONS, + ]), + ) .layer(Extension(state)); let server = axum::Server::bind(&addr).serve(server_router.into_make_service()); @@ -102,20 +129,10 @@ async fn main() -> anyhow::Result<()> { Ok(()) } -async fn fallback( - origin: Option>, - uri: Uri, -) -> (StatusCode, HeaderMap, String) { - let origin = match validate_cors(origin) { - Ok(origin) => origin, - Err((status, headers, msg)) => return (status, headers, msg), +async fn fallback(origin: Option>, uri: Uri) -> (StatusCode, String) { + if let Err((status, msg)) = validate_cors(origin) { + return (status, msg); }; - let headers = create_cors_headers(&origin); - - ( - StatusCode::NOT_FOUND, - headers, - format!("No route for {uri}"), - ) + (StatusCode::NOT_FOUND, format!("No route for {uri}")) } diff --git a/src/migration.rs b/src/migration.rs index f5b6c50..c18d061 100644 --- a/src/migration.rs +++ b/src/migration.rs @@ -6,7 +6,6 @@ use axum::headers::Authorization; use axum::http::StatusCode; use axum::{Extension, Json, TypedHeader}; use chrono::{DateTime, NaiveDateTime, Utc}; -use diesel::{Connection, PgConnection}; use log::{error, info}; use serde::{Deserialize, Deserializer}; use serde_json::json; @@ -66,8 +65,6 @@ pub async fn migration_impl(admin_key: String, state: &State) -> anyhow::Result< let mut finished = false; - let mut conn = PgConnection::establish(&state.pg_url).unwrap(); - info!("Starting migration"); while !finished { info!("Fetching {limit} items from offset {offset}"); @@ -81,15 +78,18 @@ pub async fn migration_impl(admin_key: String, state: &State) -> anyhow::Result< let items: Vec = resp.into_json()?; // Insert values into DB - conn.transaction::<_, anyhow::Error, _>(|conn| { - for item in items.iter() { - if let Ok(value) = base64::decode(&item.value) { - VssItem::put_item(conn, &item.store_id, &item.key, &value, item.version)?; - } + for item in items.iter() { + if let Ok(value) = base64::decode(&item.value) { + VssItem::put_item( + &state.client, + &item.store_id, + &item.key, + &value, + item.version, + ) + .await?; } - - Ok(()) - })?; + } if items.len() < limit { finished = true; diff --git a/src/models/migration_baseline.sql b/src/models/migration_baseline.sql new file mode 100644 index 0000000..965c311 --- /dev/null +++ b/src/models/migration_baseline.sql @@ -0,0 +1,76 @@ +CREATE TABLE vss_db +( + store_id TEXT NOT NULL CHECK (store_id != ''), + key TEXT NOT NULL, + value bytea, + version BIGINT NOT NULL, + created_date TIMESTAMP DEFAULT '2023-07-13'::TIMESTAMP NOT NULL, + updated_date TIMESTAMP DEFAULT '2023-07-13'::TIMESTAMP NOT NULL, + PRIMARY KEY (store_id, key) +); + +-- triggers to set dates automatically, generated by ChatGPT + +-- Function to set created_date and updated_date during INSERT +CREATE OR REPLACE FUNCTION set_created_date() + RETURNS TRIGGER AS $$ +BEGIN + NEW.created_date := CURRENT_TIMESTAMP; + NEW.updated_date := CURRENT_TIMESTAMP; + RETURN NEW; +END; +$$ LANGUAGE plpgsql; + +-- Function to set updated_date during UPDATE +CREATE OR REPLACE FUNCTION set_updated_date() + RETURNS TRIGGER AS $$ +BEGIN + NEW.updated_date := CURRENT_TIMESTAMP; + RETURN NEW; +END; +$$ LANGUAGE plpgsql; + +-- Trigger for INSERT operation on vss_db +CREATE TRIGGER tr_set_dates_after_insert + BEFORE INSERT ON vss_db + FOR EACH ROW +EXECUTE FUNCTION set_created_date(); + +-- Trigger for UPDATE operation on vss_db +CREATE TRIGGER tr_set_dates_after_update + BEFORE UPDATE ON vss_db + FOR EACH ROW +EXECUTE FUNCTION set_updated_date(); + +-- upsert function +CREATE OR REPLACE FUNCTION upsert_vss_db( + p_store_id TEXT, + p_key TEXT, + p_value bytea, + p_version BIGINT +) RETURNS VOID AS $$ +BEGIN + + WITH new_values (store_id, key, value, version) AS (VALUES (p_store_id, p_key, p_value, p_version)) + INSERT + INTO vss_db + (store_id, key, value, version) + SELECT new_values.store_id, + new_values.key, + new_values.value, + new_values.version + FROM new_values + LEFT JOIN vss_db AS existing + ON new_values.store_id = existing.store_id + AND new_values.key = existing.key + WHERE CASE + WHEN new_values.version >= 4294967295 THEN new_values.version >= COALESCE(existing.version, -1) + ELSE new_values.version > COALESCE(existing.version, -1) + END + ON CONFLICT (store_id, key) + DO UPDATE SET value = excluded.value, + version = excluded.version; + +END; +$$ LANGUAGE plpgsql; + diff --git a/src/models/mod.rs b/src/models/mod.rs index 2dd8fec..6968e6b 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -1,28 +1,8 @@ use crate::kv::KeyValue; -use diesel::prelude::*; -use diesel::sql_query; -use diesel::sql_types::{BigInt, Bytea, Text}; -use diesel_migrations::{embed_migrations, EmbeddedMigrations}; -use schema::vss_db; use serde::{Deserialize, Serialize}; +use tokio_postgres::{Client, Row}; -pub mod schema; - -pub const MIGRATIONS: EmbeddedMigrations = embed_migrations!(); - -#[derive( - QueryableByName, - Queryable, - Insertable, - AsChangeset, - Serialize, - Deserialize, - Debug, - Clone, - PartialEq, -)] -#[diesel(check_for_backend(diesel::pg::Pg))] -#[diesel(table_name = vss_db)] +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] pub struct VssItem { pub store_id: String, pub key: String, @@ -39,121 +19,193 @@ impl VssItem { .map(|value| KeyValue::new(self.key, value, self.version)) } - pub fn get_item( - conn: &mut PgConnection, - store_id: &str, - key: &str, + pub async fn get_item( + client: &Client, + store_id: &String, + key: &String, ) -> anyhow::Result> { - Ok(vss_db::table - .filter(vss_db::store_id.eq(store_id)) - .filter(vss_db::key.eq(key)) - .first::(conn) - .optional()?) + let stmt = client + .prepare("SELECT * FROM vss_db WHERE store_id = $1 AND key = $2") + .await?; + + client + .query_opt(&stmt, &[store_id, key]) + .await? + .map(row_to_vss_item) + .transpose() } - pub fn put_item( - conn: &mut PgConnection, - store_id: &str, - key: &str, - value: &[u8], + pub async fn put_item( + client: &Client, + store_id: &String, + key: &String, + value: &Vec, version: i64, ) -> anyhow::Result<()> { - sql_query(include_str!("put_item.sql")) - .bind::(store_id) - .bind::(key) - .bind::(value) - .bind::(version) - .execute(conn)?; + // Use Postgres built-in functions for safe escaping + let store_id_escaped = format!("quote_literal('{}')", store_id); + let key_escaped = format!("quote_literal('{}')", key); + + // For bytea data, using E'' syntax + let value_escaped = format!("E'\\\\x{}'", hex::encode(value)); + + let sql = format!( + "SELECT upsert_vss_db({}, {}, {}, {});", + store_id_escaped, key_escaped, value_escaped, version + ); + client.simple_query(&sql).await?; Ok(()) } - pub fn list_key_versions( - conn: &mut PgConnection, - store_id: &str, - prefix: Option<&str>, + pub async fn list_key_versions( + client: &Client, + store_id: &String, + prefix: Option<&String>, ) -> anyhow::Result> { - let table = vss_db::table - .filter(vss_db::store_id.eq(store_id)) - .select((vss_db::key, vss_db::version)); - - let res = match prefix { - None => table.load::<(String, i64)>(conn)?, - Some(prefix) => table - .filter(vss_db::key.ilike(format!("{prefix}%"))) - .load::<(String, i64)>(conn)?, + let store_id_escaped = format!("quote_literal('{}')", store_id); + let rows = match prefix { + Some(prefix) => { + // Safely escape the inputs using quote_literal + let prefix_escaped = format!("quote_literal('{}') || '%'", prefix); + let sql = format!( + "SELECT key, version FROM vss_db WHERE store_id = {} AND key ILIKE {}", + store_id_escaped, prefix_escaped + ); + client.simple_query(&sql).await? + } + None => { + let sql = format!( + "SELECT key, version FROM vss_db WHERE store_id = {}", + store_id_escaped + ); + client.simple_query(&sql).await? + } }; + let mut res = Vec::new(); + // Parse results + for message in rows { + if let tokio_postgres::SimpleQueryMessage::Row(row) = message { + let key: String = row + .get("key") + .ok_or(anyhow::anyhow!("key not found"))? + .to_string(); + + let version: i64 = row + .get("version") + .ok_or(anyhow::anyhow!("version not found"))? + .parse()?; + + res.push((key, version)); + } + } + Ok(res) } } +fn row_to_vss_item(row: Row) -> anyhow::Result { + let store_id: String = row.get(0); + let key: String = row.get(1); + let value: Option> = row.get(2); + let version: i64 = row.get(3); + let created_date: chrono::NaiveDateTime = row.get(4); + let updated_date: chrono::NaiveDateTime = row.get(5); + + Ok(VssItem { + store_id, + key, + value, + version, + created_date, + updated_date, + }) +} + #[cfg(test)] mod test { use super::*; use crate::State; - use diesel::{Connection, PgConnection, RunQueryDsl}; - use diesel_migrations::MigrationHarness; use secp256k1::Secp256k1; use std::str::FromStr; + use std::sync::Arc; + use tokio_postgres::NoTls; const PUBKEY: &str = "04547d92b618856f4eda84a64ec32f1694c9608a3f9dc73e91f08b5daa087260164fbc9e2a563cf4c5ef9f4c614fd9dfca7582f8de429a4799a4b202fbe80a7db5"; - fn init_state() -> State { + async fn init_state() -> State { dotenv::dotenv().ok(); let pg_url = std::env::var("DATABASE_URL").expect("DATABASE_URL must be set"); - let mut connection = PgConnection::establish(&pg_url).unwrap(); - - connection - .run_pending_migrations(MIGRATIONS) - .expect("migrations could not run"); + // Connect to the database. + let (client, connection) = tokio_postgres::connect(&pg_url, NoTls).await.unwrap(); + + // The connection object performs the actual communication with the database, + // so spawn it off to run on its own. + tokio::spawn(async move { + if let Err(e) = connection.await { + eprintln!("db connection error: {e}"); + } + }); + + client + .simple_query("DROP TABLE IF EXISTS vss_db") + .await + .unwrap(); + client + .simple_query(include_str!("migration_baseline.sql")) + .await + .unwrap(); let auth_key = secp256k1::PublicKey::from_str(PUBKEY).unwrap(); let secp = Secp256k1::new(); State { - pg_url, + client: Arc::new(client), auth_key, secp, } } - fn clear_database(state: &State) { - let mut conn = PgConnection::establish(&state.pg_url).unwrap(); - - conn.transaction::<_, anyhow::Error, _>(|conn| { - diesel::delete(vss_db::table).execute(conn)?; - Ok(()) - }) - .unwrap(); + async fn clear_database(state: &State) { + state + .client + .execute("DROP TABLE vss_db", &[]) + .await + .unwrap(); } #[tokio::test] async fn test_vss_flow() { - let state = init_state(); - clear_database(&state); + let state = init_state().await; - let store_id = "test_store_id"; - let key = "test"; - let value = [1, 2, 3, 4, 5]; + let store_id = "test_store_id".to_string(); + let key = "test".to_string(); + let value: Vec = vec![1, 2, 3, 4, 5]; let version = 0; - let mut conn = PgConnection::establish(&state.pg_url).unwrap(); - VssItem::put_item(&mut conn, store_id, key, &value, version).unwrap(); + VssItem::put_item(&state.client, &store_id, &key, &value, version) + .await + .unwrap(); - let versions = VssItem::list_key_versions(&mut conn, store_id, None).unwrap(); + let versions = VssItem::list_key_versions(&state.client, &store_id, None) + .await + .unwrap(); assert_eq!(versions.len(), 1); assert_eq!(versions[0].0, key); assert_eq!(versions[0].1, version); - let new_value = [6, 7, 8, 9, 10]; + let new_value = vec![6, 7, 8, 9, 10]; let new_version = version + 1; - VssItem::put_item(&mut conn, store_id, key, &new_value, new_version).unwrap(); + VssItem::put_item(&state.client, &store_id, &key, &new_value, new_version) + .await + .unwrap(); - let item = VssItem::get_item(&mut conn, store_id, key) + let item = VssItem::get_item(&state.client, &store_id, &key) + .await .unwrap() .unwrap(); @@ -162,23 +214,24 @@ mod test { assert_eq!(item.value.unwrap(), new_value); assert_eq!(item.version, new_version); - clear_database(&state); + clear_database(&state).await; } #[tokio::test] async fn test_max_version_number() { - let state = init_state(); - clear_database(&state); + let state = init_state().await; - let store_id = "max_test_store_id"; - let key = "max_test"; - let value = [1, 2, 3, 4, 5]; + let store_id = "max_test_store_id".to_string(); + let key = "max_test".to_string(); + let value = vec![1, 2, 3, 4, 5]; let version = u32::MAX as i64; - let mut conn = PgConnection::establish(&state.pg_url).unwrap(); - VssItem::put_item(&mut conn, store_id, key, &value, version).unwrap(); + VssItem::put_item(&state.client, &store_id, &key, &value, version) + .await + .unwrap(); - let item = VssItem::get_item(&mut conn, store_id, key) + let item = VssItem::get_item(&state.client, &store_id, &key) + .await .unwrap() .unwrap(); @@ -186,11 +239,14 @@ mod test { assert_eq!(item.key, key); assert_eq!(item.value.unwrap(), value); - let new_value = [6, 7, 8, 9, 10]; + let new_value = vec![6, 7, 8, 9, 10]; - VssItem::put_item(&mut conn, store_id, key, &new_value, version).unwrap(); + VssItem::put_item(&state.client, &store_id, &key, &new_value, version) + .await + .unwrap(); - let item = VssItem::get_item(&mut conn, store_id, key) + let item = VssItem::get_item(&state.client, &store_id, &key) + .await .unwrap() .unwrap(); @@ -198,38 +254,48 @@ mod test { assert_eq!(item.key, key); assert_eq!(item.value.unwrap(), new_value); - clear_database(&state); + clear_database(&state).await; } #[tokio::test] async fn test_list_key_versions() { - let state = init_state(); - clear_database(&state); + let state = init_state().await; - let store_id = "list_kv_test_store_id"; - let key = "kv_test"; - let key1 = "other_kv_test"; - let value = [1, 2, 3, 4, 5]; + let store_id = "list_kv_test_store_id".to_string(); + let key = "kv_test".to_string(); + let key1 = "other_kv_test".to_string(); + let value = vec![1, 2, 3, 4, 5]; let version = 0; - let mut conn = PgConnection::establish(&state.pg_url).unwrap(); - VssItem::put_item(&mut conn, store_id, key, &value, version).unwrap(); + VssItem::put_item(&state.client, &store_id, &key, &value, version) + .await + .unwrap(); - VssItem::put_item(&mut conn, store_id, key1, &value, version).unwrap(); + VssItem::put_item(&state.client, &store_id, &key1, &value, version) + .await + .unwrap(); - let versions = VssItem::list_key_versions(&mut conn, store_id, None).unwrap(); + let versions = VssItem::list_key_versions(&state.client, &store_id, None) + .await + .unwrap(); assert_eq!(versions.len(), 2); - let versions = VssItem::list_key_versions(&mut conn, store_id, Some("kv")).unwrap(); + let versions = + VssItem::list_key_versions(&state.client, &store_id, Some(&"kv".to_string())) + .await + .unwrap(); assert_eq!(versions.len(), 1); assert_eq!(versions[0].0, key); assert_eq!(versions[0].1, version); - let versions = VssItem::list_key_versions(&mut conn, store_id, Some("other")).unwrap(); + let versions = + VssItem::list_key_versions(&state.client, &store_id, Some(&"other".to_string())) + .await + .unwrap(); assert_eq!(versions.len(), 1); assert_eq!(versions[0].0, key1); assert_eq!(versions[0].1, version); - clear_database(&state); + clear_database(&state).await; } } diff --git a/src/models/schema.rs b/src/models/schema.rs deleted file mode 100644 index de1aed0..0000000 --- a/src/models/schema.rs +++ /dev/null @@ -1,12 +0,0 @@ -// @generated automatically by Diesel CLI. - -diesel::table! { - vss_db (store_id, key) { - store_id -> Text, - key -> Text, - value -> Nullable, - version -> Int8, - created_date -> Timestamp, - updated_date -> Timestamp, - } -} diff --git a/src/routes.rs b/src/routes.rs index 31f4ec2..a113b7e 100644 --- a/src/routes.rs +++ b/src/routes.rs @@ -3,27 +3,21 @@ use crate::kv::{KeyValue, KeyValueOld}; use crate::models::VssItem; use crate::{State, ALLOWED_LOCALHOST, ALLOWED_ORIGINS, ALLOWED_SUBDOMAIN}; use axum::headers::authorization::Bearer; -use axum::headers::{Authorization, HeaderMap, Origin}; -use axum::http::header::{ - ACCESS_CONTROL_ALLOW_HEADERS, ACCESS_CONTROL_ALLOW_METHODS, ACCESS_CONTROL_ALLOW_ORIGIN, - CONTENT_TYPE, -}; +use axum::headers::{Authorization, Origin}; use axum::http::StatusCode; use axum::{Extension, Json, TypedHeader}; -use diesel::{Connection, PgConnection}; use log::{debug, error, trace}; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; macro_rules! ensure_store_id { - ($payload:ident, $store_id:expr, $headers:ident) => { + ($payload:ident, $store_id:expr) => { match $payload.store_id { None => $payload.store_id = Some($store_id), Some(ref id) => { if id != &$store_id { return Err(( StatusCode::UNAUTHORIZED, - $headers, format!("Unauthorized: store_id mismatch"), )); } @@ -45,8 +39,7 @@ pub async fn get_object_impl( trace!("get_object_impl: {req:?}"); let store_id = req.store_id.expect("must have"); - let mut conn = PgConnection::establish(&state.pg_url).unwrap(); - let item = VssItem::get_item(&mut conn, &store_id, &req.key)?; + let item = VssItem::get_item(&state.client, &store_id, &req.key).await?; Ok(item.and_then(|i| i.into_kv())) } @@ -56,20 +49,18 @@ pub async fn get_object( TypedHeader(token): TypedHeader>, Extension(state): Extension, Json(mut payload): Json, -) -> Result<(HeaderMap, Json>), (StatusCode, HeaderMap, String)> { +) -> Result>, (StatusCode, String)> { debug!("get_object: {payload:?}"); - let origin = validate_cors(origin)?; - let mut headers = create_cors_headers(&origin); - headers.insert(CONTENT_TYPE, "application/json".parse().unwrap()); + validate_cors(origin)?; - let store_id = verify_token(token.token(), &state, &headers)?; + let store_id = verify_token(token.token(), &state)?; - ensure_store_id!(payload, store_id, headers); + ensure_store_id!(payload, store_id); match get_object_impl(payload, &state).await { - Ok(Some(res)) => Ok((headers, Json(Some(res.into())))), - Ok(None) => Ok((headers, Json(None))), - Err(e) => Err(handle_anyhow_error(e, headers)), + Ok(Some(res)) => Ok(Json(Some(res.into()))), + Ok(None) => Ok(Json(None)), + Err(e) => Err(handle_anyhow_error("get_object", e)), } } @@ -78,19 +69,17 @@ pub async fn get_object_v2( TypedHeader(token): TypedHeader>, Extension(state): Extension, Json(mut payload): Json, -) -> Result>, (StatusCode, HeaderMap, String)> { +) -> Result>, (StatusCode, String)> { debug!("get_object v2: {payload:?}"); - let origin = validate_cors(origin)?; - let mut headers = create_cors_headers(&origin); - headers.insert(CONTENT_TYPE, "application/json".parse().unwrap()); + validate_cors(origin)?; - let store_id = verify_token(token.token(), &state, &headers)?; + let store_id = verify_token(token.token(), &state)?; - ensure_store_id!(payload, store_id, headers); + ensure_store_id!(payload, store_id); match get_object_impl(payload, &state).await { Ok(res) => Ok(Json(res)), - Err(e) => Err(handle_anyhow_error(e, headers)), + Err(e) => Err(handle_anyhow_error("get_object_v2", e)), } } @@ -110,14 +99,12 @@ pub async fn put_objects_impl(req: PutObjectsRequest, state: &State) -> anyhow:: let store_id = req.store_id.expect("must have"); - let mut conn = PgConnection::establish(&state.pg_url).unwrap(); - conn.transaction(|conn| { - for kv in req.transaction_items { - VssItem::put_item(conn, &store_id, &kv.key, &kv.value.0, kv.version)?; - } + // todo use transaction + for kv in req.transaction_items { + VssItem::put_item(&state.client, &store_id, &kv.key, &kv.value.0, kv.version).await?; + } - Ok(()) - }) + Ok(()) } pub async fn put_objects( @@ -125,18 +112,16 @@ pub async fn put_objects( TypedHeader(token): TypedHeader>, Extension(state): Extension, Json(mut payload): Json, -) -> Result<(HeaderMap, Json<()>), (StatusCode, HeaderMap, String)> { - let origin = validate_cors(origin)?; - let mut headers = create_cors_headers(&origin); - headers.insert(CONTENT_TYPE, "application/json".parse().unwrap()); +) -> Result, (StatusCode, String)> { + validate_cors(origin)?; - let store_id = verify_token(token.token(), &state, &headers)?; + let store_id = verify_token(token.token(), &state)?; - ensure_store_id!(payload, store_id, headers); + ensure_store_id!(payload, store_id); match put_objects_impl(payload, &state).await { - Ok(res) => Ok((headers, Json(res))), - Err(e) => Err(handle_anyhow_error(e, headers)), + Ok(res) => Ok(Json(res)), + Err(e) => Err(handle_anyhow_error("put_objects", e)), } } @@ -155,8 +140,8 @@ pub async fn list_key_versions_impl( // todo pagination let store_id = req.store_id.expect("must have"); - let mut conn = PgConnection::establish(&state.pg_url).unwrap(); - let versions = VssItem::list_key_versions(&mut conn, &store_id, req.key_prefix.as_deref())?; + let versions = + VssItem::list_key_versions(&state.client, &store_id, req.key_prefix.as_ref()).await?; let json = versions .into_iter() @@ -176,18 +161,16 @@ pub async fn list_key_versions( TypedHeader(token): TypedHeader>, Extension(state): Extension, Json(mut payload): Json, -) -> Result<(HeaderMap, Json>), (StatusCode, HeaderMap, String)> { - let origin = validate_cors(origin)?; - let mut headers = create_cors_headers(&origin); - headers.insert(CONTENT_TYPE, "application/json".parse().unwrap()); +) -> Result>, (StatusCode, String)> { + validate_cors(origin)?; - let store_id = verify_token(token.token(), &state, &headers)?; + let store_id = verify_token(token.token(), &state)?; - ensure_store_id!(payload, store_id, headers); + ensure_store_id!(payload, store_id); match list_key_versions_impl(payload, &state).await { - Ok(res) => Ok((headers, Json(res))), - Err(e) => Err(handle_anyhow_error(e, headers)), + Ok(res) => Ok(Json(res)), + Err(e) => Err(handle_anyhow_error("list_key_versions", e)), } } @@ -195,45 +178,31 @@ pub async fn health_check() -> Result, (StatusCode, String)> { Ok(Json(())) } -pub fn validate_cors( - origin: Option>, -) -> Result { +pub fn valid_origin(origin: &str) -> bool { + ALLOWED_ORIGINS.contains(&origin) + || origin.ends_with(ALLOWED_SUBDOMAIN) + || origin.starts_with(ALLOWED_LOCALHOST) +} + +pub fn validate_cors(origin: Option>) -> Result<(), (StatusCode, String)> { if let Some(TypedHeader(origin)) = origin { if origin.is_null() { - return Ok("*".to_string()); + return Ok(()); } let origin_str = origin.to_string(); - if ALLOWED_ORIGINS.contains(&origin_str.as_str()) - || origin_str.ends_with(ALLOWED_SUBDOMAIN) - || origin_str.starts_with(ALLOWED_LOCALHOST) - { - return Ok(origin_str); + if valid_origin(&origin_str) { + return Ok(()); } else { - let headers = create_cors_headers("*"); // The origin is not in the allowed list block the request - return Err((StatusCode::NOT_FOUND, headers, String::new())); + return Err((StatusCode::NOT_FOUND, String::new())); } } - Ok("*".to_string()) -} - -pub fn create_cors_headers(origin: &str) -> HeaderMap { - let mut headers = HeaderMap::new(); - headers.insert(ACCESS_CONTROL_ALLOW_ORIGIN, origin.parse().unwrap()); - headers.insert( - ACCESS_CONTROL_ALLOW_METHODS, - "GET, POST, PUT, DELETE, OPTIONS".parse().unwrap(), - ); - headers.insert(ACCESS_CONTROL_ALLOW_HEADERS, "*".parse().unwrap()); - headers + Ok(()) } -pub(crate) fn handle_anyhow_error( - err: anyhow::Error, - headers: HeaderMap, -) -> (StatusCode, HeaderMap, String) { - error!("Error: {err:?}"); - (StatusCode::BAD_REQUEST, headers, format!("{err}")) +pub(crate) fn handle_anyhow_error(function: &str, err: anyhow::Error) -> (StatusCode, String) { + error!("Error in {function}: {err:?}"); + (StatusCode::BAD_REQUEST, format!("{err}")) } From 9371a01b6dbfd32f5dffbc5b4b733ccdc305fa32 Mon Sep 17 00:00:00 2001 From: benthecarman Date: Fri, 22 Sep 2023 21:29:34 -0500 Subject: [PATCH 3/7] maybe working --- Cargo.lock | 138 ++++++++++++++++++++++-------- Cargo.toml | 6 +- src/main.rs | 17 ++-- src/models/migration_baseline.sql | 2 +- src/models/mod.rs | 88 ++++++++----------- 5 files changed, 152 insertions(+), 99 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 715eadf..c940fac 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -231,6 +231,16 @@ dependencies = [ "windows-targets", ] +[[package]] +name = "core-foundation" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "194a7a9e6de53fa55116934067c844d9d749312f75c6f6d0980e8c252f8c2146" +dependencies = [ + "core-foundation-sys", + "libc", +] + [[package]] name = "core-foundation-sys" version = "0.8.4" @@ -322,6 +332,12 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4443176a9f2c162692bd3d352d745ef9413eec5782a80d8fd6f8a1ac692a07f7" +[[package]] +name = "fastrand" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6999dc1837253364c2ebb0704ba97994bd874e8f195d665c50b7548f6ea92764" + [[package]] name = "finl_unicode" version = "1.2.0" @@ -767,6 +783,24 @@ dependencies = [ "windows-sys", ] +[[package]] +name = "native-tls" +version = "0.2.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07226173c32f2926027b63cce4bcd8076c3552846cbe7925f3aaffeac0a3b92e" +dependencies = [ + "lazy_static", + "libc", + "log", + "openssl", + "openssl-probe", + "openssl-sys", + "schannel", + "security-framework", + "security-framework-sys", + "tempfile", +] + [[package]] name = "num-traits" version = "0.2.16" @@ -827,6 +861,12 @@ dependencies = [ "syn", ] +[[package]] +name = "openssl-probe" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff011a302c396a5197692431fc1948019154afc178baf7d8e37367442a4601cf" + [[package]] name = "openssl-sys" version = "0.9.93" @@ -925,25 +965,22 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "26072860ba924cbfa98ea39c8c19b4dd6a4a25423dbdf219c1eca91aa0cf6964" [[package]] -name = "postgres-openssl" +name = "postgres-native-tls" version = "0.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1de0ea6504e07ca78355a6fb88ad0f36cafe9e696cbc6717f16a207f3a60be72" +source = "git+https://github.com/prisma/rust-postgres?branch=pgbouncer-mode#a1a2dc6d9584deaf70a14293c428e7b6ca614d98" dependencies = [ - "futures", - "openssl", + "native-tls", "tokio", - "tokio-openssl", + "tokio-native-tls", "tokio-postgres", ] [[package]] name = "postgres-protocol" -version = "0.6.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49b6c5ef183cd3ab4ba005f1ca64c21e8bd97ce4699cfea9e8d9a2c4958ca520" +version = "0.6.4" +source = "git+https://github.com/prisma/rust-postgres?branch=pgbouncer-mode#a1a2dc6d9584deaf70a14293c428e7b6ca614d98" dependencies = [ - "base64 0.21.4", + "base64 0.13.1", "byteorder", "bytes", "fallible-iterator", @@ -957,9 +994,8 @@ dependencies = [ [[package]] name = "postgres-types" -version = "0.2.6" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8d2234cdee9408b523530a9b6d2d6b373d1db34f6a8e51dc03ded1828d7fb67c" +version = "0.2.4" +source = "git+https://github.com/prisma/rust-postgres?branch=pgbouncer-mode#a1a2dc6d9584deaf70a14293c428e7b6ca614d98" dependencies = [ "bytes", "chrono", @@ -1147,6 +1183,15 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ad4cc8da4ef723ed60bced201181d83791ad433213d8c24efffda1eec85d741" +[[package]] +name = "schannel" +version = "0.1.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c3733bf4cf7ea0880754e19cb5a462007c4a8c1914bff372ccc95b464f1df88" +dependencies = [ + "windows-sys", +] + [[package]] name = "scopeguard" version = "1.2.0" @@ -1182,6 +1227,29 @@ dependencies = [ "cc", ] +[[package]] +name = "security-framework" +version = "2.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05b64fb303737d99b81884b2c63433e9ae28abebe5eb5045dcdd175dc2ecf4de" +dependencies = [ + "bitflags 1.3.2", + "core-foundation", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + +[[package]] +name = "security-framework-sys" +version = "2.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e932934257d3b408ed8f30db49d85ea163bfe74961f017f405b025af298f0c7a" +dependencies = [ + "core-foundation-sys", + "libc", +] + [[package]] name = "serde" version = "1.0.188" @@ -1357,6 +1425,19 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2047c6ded9c721764247e62cd3b03c09ffc529b2ba5b10ec482ae507a4a70160" +[[package]] +name = "tempfile" +version = "3.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb94d2f3cc536af71caac6b6fcebf65860b347e7ce0cc9ebe8f70d3e521054ef" +dependencies = [ + "cfg-if", + "fastrand", + "redox_syscall", + "rustix", + "windows-sys", +] + [[package]] name = "termcolor" version = "1.3.0" @@ -1412,22 +1493,19 @@ dependencies = [ ] [[package]] -name = "tokio-openssl" -version = "0.6.3" +name = "tokio-native-tls" +version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c08f9ffb7809f1b20c1b398d92acf4cc719874b3b2b2d9ea2f09b4a80350878a" +checksum = "bbae76ab933c85776efabc971569dd6119c580d8f5d448769dec1764bf796ef2" dependencies = [ - "futures-util", - "openssl", - "openssl-sys", + "native-tls", "tokio", ] [[package]] name = "tokio-postgres" -version = "0.7.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d340244b32d920260ae7448cb72b6e238bddc3d4f7603394e7dd46ed8e48f5b8" +version = "0.7.7" +source = "git+https://github.com/prisma/rust-postgres?branch=pgbouncer-mode#a1a2dc6d9584deaf70a14293c428e7b6ca614d98" dependencies = [ "async-trait", "byteorder", @@ -1442,11 +1520,9 @@ dependencies = [ "pin-project-lite", "postgres-protocol", "postgres-types", - "rand", "socket2 0.5.4", "tokio", "tokio-util", - "whoami", ] [[package]] @@ -1623,8 +1699,8 @@ dependencies = [ "hex", "jwt-compact", "log", - "openssl", - "postgres-openssl", + "native-tls", + "postgres-native-tls", "pretty_env_logger", "secp256k1", "serde", @@ -1724,16 +1800,6 @@ dependencies = [ "rustls-webpki 0.100.3", ] -[[package]] -name = "whoami" -version = "1.4.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "22fc3756b8a9133049b26c7f61ab35416c130e8c09b660f5b3958b446f52cc50" -dependencies = [ - "wasm-bindgen", - "web-sys", -] - [[package]] name = "winapi" version = "0.3.9" diff --git a/Cargo.toml b/Cargo.toml index a3928dd..4476964 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -19,9 +19,9 @@ sha2 = { version = "0.10", default-features = false } serde = { version = "^1.0", features = ["derive"] } serde_json = "1.0.67" tokio = { version = "1.12.0", features = ["full"] } -tokio-postgres = { version = "0.7.10", features = ["with-chrono-0_4"] } +tokio-postgres = { git = "https://github.com/prisma/rust-postgres", branch = "pgbouncer-mode", features = ["with-chrono-0_4"] } tower-http = { version = "0.4.0", features = ["cors"] } ureq = { version = "2.5.0", features = ["json"] } -postgres-openssl = "0.5.0" -openssl = "0.10.57" +postgres-native-tls = { git = "https://github.com/prisma/rust-postgres", branch = "pgbouncer-mode" } +native-tls = "0.2" diff --git a/src/main.rs b/src/main.rs index 109a2f7..c7fdc53 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,11 +3,12 @@ use axum::headers::Origin; use axum::http::{request::Parts, HeaderValue, Method, StatusCode, Uri}; use axum::routing::{get, post, put}; use axum::{http, Extension, Router, TypedHeader}; -use openssl::ssl::{SslConnector, SslMethod}; -use postgres_openssl::MakeTlsConnector; +use native_tls::TlsConnector; +use postgres_native_tls::MakeTlsConnector; use secp256k1::{All, PublicKey, Secp256k1}; +use std::str::FromStr; use std::sync::Arc; -use tokio_postgres::Client; +use tokio_postgres::{Client, Config}; use tower_http::cors::{AllowOrigin, CorsLayer}; mod auth; @@ -53,17 +54,19 @@ async fn main() -> anyhow::Result<()> { let auth_key_bytes = hex::decode(auth_key)?; let auth_key = PublicKey::from_slice(&auth_key_bytes)?; - let builder = SslConnector::builder(SslMethod::tls())?; - let connector = MakeTlsConnector::new(builder.build()); + let tls = TlsConnector::new()?; + let connector = MakeTlsConnector::new(tls); // Connect to the database. - let (client, connection) = tokio_postgres::connect(&pg_url, connector).await?; + let mut config = Config::from_str(&pg_url).unwrap(); + config.pgbouncer_mode(true); + let (client, connection) = config.connect(connector).await?; // The connection object performs the actual communication with the database, // so spawn it off to run on its own. tokio::spawn(async move { if let Err(e) = connection.await { - eprintln!("db connection error: {e}"); + panic!("db connection error: {e}"); } }); diff --git a/src/models/migration_baseline.sql b/src/models/migration_baseline.sql index 965c311..37574b7 100644 --- a/src/models/migration_baseline.sql +++ b/src/models/migration_baseline.sql @@ -1,4 +1,4 @@ -CREATE TABLE vss_db +CREATE TABLE IF NOT EXISTS vss_db ( store_id TEXT NOT NULL CHECK (store_id != ''), key TEXT NOT NULL, diff --git a/src/models/mod.rs b/src/models/mod.rs index 6968e6b..5c474f7 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -24,12 +24,11 @@ impl VssItem { store_id: &String, key: &String, ) -> anyhow::Result> { - let stmt = client - .prepare("SELECT * FROM vss_db WHERE store_id = $1 AND key = $2") - .await?; - client - .query_opt(&stmt, &[store_id, key]) + .query_opt( + "SELECT * FROM vss_db WHERE store_id = $1 AND key = $2", + &[store_id, key], + ) .await? .map(row_to_vss_item) .transpose() @@ -42,18 +41,12 @@ impl VssItem { value: &Vec, version: i64, ) -> anyhow::Result<()> { - // Use Postgres built-in functions for safe escaping - let store_id_escaped = format!("quote_literal('{}')", store_id); - let key_escaped = format!("quote_literal('{}')", key); - - // For bytea data, using E'' syntax - let value_escaped = format!("E'\\\\x{}'", hex::encode(value)); - - let sql = format!( - "SELECT upsert_vss_db({}, {}, {}, {});", - store_id_escaped, key_escaped, value_escaped, version - ); - client.simple_query(&sql).await?; + client + .execute( + "SELECT upsert_vss_db($1, $2, $3, $4)", + &[store_id, key, value, &version], + ) + .await?; Ok(()) } @@ -63,43 +56,32 @@ impl VssItem { store_id: &String, prefix: Option<&String>, ) -> anyhow::Result> { - let store_id_escaped = format!("quote_literal('{}')", store_id); let rows = match prefix { - Some(prefix) => { - // Safely escape the inputs using quote_literal - let prefix_escaped = format!("quote_literal('{}') || '%'", prefix); - let sql = format!( - "SELECT key, version FROM vss_db WHERE store_id = {} AND key ILIKE {}", - store_id_escaped, prefix_escaped - ); - client.simple_query(&sql).await? - } + Some(prefix) => client + .query( + "SELECT key, version FROM vss_db WHERE store_id = $1 AND key ILIKE $2 || '%'", + &[store_id, prefix], + ) + .await?, None => { - let sql = format!( - "SELECT key, version FROM vss_db WHERE store_id = {}", - store_id_escaped - ); - client.simple_query(&sql).await? + client + .query( + "SELECT key, version FROM vss_db WHERE store_id = $1", + &[store_id], + ) + .await? } }; - let mut res = Vec::new(); - // Parse results - for message in rows { - if let tokio_postgres::SimpleQueryMessage::Row(row) = message { - let key: String = row - .get("key") - .ok_or(anyhow::anyhow!("key not found"))? - .to_string(); - - let version: i64 = row - .get("version") - .ok_or(anyhow::anyhow!("version not found"))? - .parse()?; - - res.push((key, version)); - } - } + let res = rows + .into_iter() + .map(|row| { + let key: String = row.get(0); + let version: i64 = row.get(1); + + (key, version) + }) + .collect(); Ok(res) } @@ -130,15 +112,17 @@ mod test { use secp256k1::Secp256k1; use std::str::FromStr; use std::sync::Arc; - use tokio_postgres::NoTls; + use tokio_postgres::{Config, NoTls}; const PUBKEY: &str = "04547d92b618856f4eda84a64ec32f1694c9608a3f9dc73e91f08b5daa087260164fbc9e2a563cf4c5ef9f4c614fd9dfca7582f8de429a4799a4b202fbe80a7db5"; async fn init_state() -> State { dotenv::dotenv().ok(); let pg_url = std::env::var("DATABASE_URL").expect("DATABASE_URL must be set"); + let mut config = Config::from_str(&pg_url).unwrap(); + config.pgbouncer_mode(true); // Connect to the database. - let (client, connection) = tokio_postgres::connect(&pg_url, NoTls).await.unwrap(); + let (client, connection) = config.connect(NoTls).await.unwrap(); // The connection object performs the actual communication with the database, // so spawn it off to run on its own. @@ -171,7 +155,7 @@ mod test { async fn clear_database(state: &State) { state .client - .execute("DROP TABLE vss_db", &[]) + .simple_query("DROP TABLE vss_db") .await .unwrap(); } From 9f09f6512ca3fd5e737873c86d82212d282a593f Mon Sep 17 00:00:00 2001 From: benthecarman Date: Fri, 22 Sep 2023 22:31:16 -0500 Subject: [PATCH 4/7] Go back to pooled connection --- Cargo.lock | 476 +++++++----------- Cargo.toml | 7 +- Dockerfile | 4 +- migrations/.keep | 0 .../down.sql | 6 + .../up.sql | 36 ++ .../2023-09-23-030518_baseline/down.sql | 9 + migrations/2023-09-23-030518_baseline/up.sql | 79 +++ src/main.rs | 33 +- src/migration.rs | 22 +- src/models/mod.rs | 275 +++++----- src/models/put_item.sql | 17 - src/models/schema.rs | 12 + src/routes.rs | 23 +- 14 files changed, 482 insertions(+), 517 deletions(-) create mode 100644 migrations/.keep create mode 100644 migrations/00000000000000_diesel_initial_setup/down.sql create mode 100644 migrations/00000000000000_diesel_initial_setup/up.sql create mode 100644 migrations/2023-09-23-030518_baseline/down.sql create mode 100644 migrations/2023-09-23-030518_baseline/up.sql delete mode 100644 src/models/put_item.sql create mode 100644 src/models/schema.rs diff --git a/Cargo.lock b/Cargo.lock index c940fac..d5597d3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -147,6 +147,19 @@ version = "1.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8c3c1a368f70d6cf7302d78f8f7093da241fb8e8807c05cc9e51a125895a6d5b" +[[package]] +name = "bigdecimal" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "454bca3db10617b88b566f205ed190aedb0e0e6dd4cad61d3988a72e8c5594cb" +dependencies = [ + "autocfg", + "libm", + "num-bigint", + "num-integer", + "num-traits", +] + [[package]] name = "bitcoin-private" version = "0.1.0" @@ -231,16 +244,6 @@ dependencies = [ "windows-targets", ] -[[package]] -name = "core-foundation" -version = "0.9.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "194a7a9e6de53fa55116934067c844d9d749312f75c6f6d0980e8c252f8c2146" -dependencies = [ - "core-foundation-sys", - "libc", -] - [[package]] name = "core-foundation-sys" version = "0.8.4" @@ -275,6 +278,57 @@ dependencies = [ "typenum", ] +[[package]] +name = "diesel" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d98235fdc2f355d330a8244184ab6b4b33c28679c0b4158f63138e51d6cf7e88" +dependencies = [ + "bigdecimal", + "bitflags 2.4.0", + "byteorder", + "chrono", + "diesel_derives", + "itoa", + "num-bigint", + "num-integer", + "num-traits", + "pq-sys", + "r2d2", +] + +[[package]] +name = "diesel_derives" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e054665eaf6d97d1e7125512bb2d35d07c73ac86cc6920174cb42d1ab697a554" +dependencies = [ + "diesel_table_macro_syntax", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "diesel_migrations" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6036b3f0120c5961381b570ee20a02432d7e2d27ea60de9578799cf9156914ac" +dependencies = [ + "diesel", + "migrations_internals", + "migrations_macros", +] + +[[package]] +name = "diesel_table_macro_syntax" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc5557efc453706fed5e4fa85006fe9817c224c3f480a34c7e5959fd700921c5" +dependencies = [ + "syn", +] + [[package]] name = "digest" version = "0.10.7" @@ -305,6 +359,12 @@ dependencies = [ "termcolor", ] +[[package]] +name = "equivalent" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5443807d6dff69373d433ab9ef5378ad8df50ca6298caf15de6e52e24aaf54d5" + [[package]] name = "errno" version = "0.3.3" @@ -326,24 +386,6 @@ dependencies = [ "libc", ] -[[package]] -name = "fallible-iterator" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4443176a9f2c162692bd3d352d745ef9413eec5782a80d8fd6f8a1ac692a07f7" - -[[package]] -name = "fastrand" -version = "2.0.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6999dc1837253364c2ebb0704ba97994bd874e8f195d665c50b7548f6ea92764" - -[[package]] -name = "finl_unicode" -version = "1.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fcfdc7a0362c9f4444381a9e697c79d435fe65b52a37466fc2c1184cee9edc6" - [[package]] name = "flate2" version = "1.0.27" @@ -360,21 +402,6 @@ version = "1.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" -[[package]] -name = "foreign-types" -version = "0.3.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f6f339eb8adc052cd2ca78910fda869aefa38d22d5cb648e6485e4d3fc06f3b1" -dependencies = [ - "foreign-types-shared", -] - -[[package]] -name = "foreign-types-shared" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "00b0228411908ca8685dba7fc2cdd70ec9990a6e753e89b6ac91a84c40fbaf4b" - [[package]] name = "form_urlencoded" version = "1.2.0" @@ -483,17 +510,6 @@ dependencies = [ "version_check", ] -[[package]] -name = "getrandom" -version = "0.2.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "be4136b2a15dd319360be1c07d9933517ccf0be8f16bf62a3bee4f0d618df427" -dependencies = [ - "cfg-if", - "libc", - "wasi", -] - [[package]] name = "gimli" version = "0.28.0" @@ -506,6 +522,12 @@ version = "1.8.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eabb4a44450da02c90444cf74558da904edde8fb4e9035a9a6a4e15445af0bd7" +[[package]] +name = "hashbrown" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c6201b9ff9fd90a5a3bac2e56a830d0caa509576f0e503818ee82c181b3437a" + [[package]] name = "headers" version = "0.3.9" @@ -653,6 +675,16 @@ dependencies = [ "unicode-normalization", ] +[[package]] +name = "indexmap" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d5477fe2230a79769d8dc68e0eabf5437907c0457a5614a9e8dddb67f65eb65d" +dependencies = [ + "equivalent", + "hashbrown", +] + [[package]] name = "is-terminal" version = "0.4.9" @@ -713,6 +745,12 @@ version = "0.2.148" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9cdc71e17332e86d2e1d38c1f99edcb6288ee11b815fb1a4b049eaa2114d369b" +[[package]] +name = "libm" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f7012b1bbb0719e1097c47611d3898568c546d597c2e74d66f6087edd5233ff4" + [[package]] name = "linux-raw-sys" version = "0.4.7" @@ -742,20 +780,31 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0e7465ac9959cc2b1404e8e2367b43684a6d13790fe23056cc8c6c5a6b7bcb94" [[package]] -name = "md-5" -version = "0.10.6" +name = "memchr" +version = "2.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d89e7ee0cfbedfc4da3340218492196241d89eefb6dab27de5df917a6d2e78cf" +checksum = "8f232d6ef707e1956a43342693d2a31e72989554d58299d7a88738cc95b0d35c" + +[[package]] +name = "migrations_internals" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f23f71580015254b020e856feac3df5878c2c7a8812297edd6c0a485ac9dada" dependencies = [ - "cfg-if", - "digest", + "serde", + "toml", ] [[package]] -name = "memchr" -version = "2.6.3" +name = "migrations_macros" +version = "2.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f232d6ef707e1956a43342693d2a31e72989554d58299d7a88738cc95b0d35c" +checksum = "cce3325ac70e67bbab5bd837a31cae01f1a6db64e0e744a33cb03a543469ef08" +dependencies = [ + "migrations_internals", + "proc-macro2", + "quote", +] [[package]] name = "mime" @@ -784,21 +833,24 @@ dependencies = [ ] [[package]] -name = "native-tls" -version = "0.2.11" +name = "num-bigint" +version = "0.4.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "07226173c32f2926027b63cce4bcd8076c3552846cbe7925f3aaffeac0a3b92e" +checksum = "608e7659b5c3d7cba262d894801b9ec9d00de989e8a82bd4bef91d08da45cdc0" dependencies = [ - "lazy_static", - "libc", - "log", - "openssl", - "openssl-probe", - "openssl-sys", - "schannel", - "security-framework", - "security-framework-sys", - "tempfile", + "autocfg", + "num-integer", + "num-traits", +] + +[[package]] +name = "num-integer" +version = "0.1.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "225d3389fb3509a24c93f5c29eb6bde2586b98d9f016636dff58d7c6f7569cd9" +dependencies = [ + "autocfg", + "num-traits", ] [[package]] @@ -835,50 +887,6 @@ version = "1.18.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "dd8b5dd2ae5ed71462c540258bedcb51965123ad7e7ccf4b9a8cafaa4a63576d" -[[package]] -name = "openssl" -version = "0.10.57" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bac25ee399abb46215765b1cb35bc0212377e58a061560d8b29b024fd0430e7c" -dependencies = [ - "bitflags 2.4.0", - "cfg-if", - "foreign-types", - "libc", - "once_cell", - "openssl-macros", - "openssl-sys", -] - -[[package]] -name = "openssl-macros" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c" -dependencies = [ - "proc-macro2", - "quote", - "syn", -] - -[[package]] -name = "openssl-probe" -version = "0.1.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ff011a302c396a5197692431fc1948019154afc178baf7d8e37367442a4601cf" - -[[package]] -name = "openssl-sys" -version = "0.9.93" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "db4d56a4c0478783083cfafcc42493dd4a981d41669da64b4572a2a089b51b1d" -dependencies = [ - "cc", - "libc", - "pkg-config", - "vcpkg", -] - [[package]] name = "parking_lot" version = "0.12.1" @@ -908,24 +916,6 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b2a4787296e9989611394c33f193f676704af1686e70b8f8033ab5ba9a35a94" -[[package]] -name = "phf" -version = "0.11.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ade2d8b8f33c7333b51bcf0428d37e217e9f32192ae4772156f65063b8ce03dc" -dependencies = [ - "phf_shared", -] - -[[package]] -name = "phf_shared" -version = "0.11.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "90fcb95eef784c2ac79119d1dd819e162b5da872ce6f3c3abe1e8ca1c082f72b" -dependencies = [ - "siphasher", -] - [[package]] name = "pin-project" version = "1.1.3" @@ -959,56 +949,14 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" [[package]] -name = "pkg-config" -version = "0.3.27" +name = "pq-sys" +version = "0.4.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "26072860ba924cbfa98ea39c8c19b4dd6a4a25423dbdf219c1eca91aa0cf6964" - -[[package]] -name = "postgres-native-tls" -version = "0.5.0" -source = "git+https://github.com/prisma/rust-postgres?branch=pgbouncer-mode#a1a2dc6d9584deaf70a14293c428e7b6ca614d98" -dependencies = [ - "native-tls", - "tokio", - "tokio-native-tls", - "tokio-postgres", -] - -[[package]] -name = "postgres-protocol" -version = "0.6.4" -source = "git+https://github.com/prisma/rust-postgres?branch=pgbouncer-mode#a1a2dc6d9584deaf70a14293c428e7b6ca614d98" +checksum = "31c0052426df997c0cbd30789eb44ca097e3541717a7b8fa36b1c464ee7edebd" dependencies = [ - "base64 0.13.1", - "byteorder", - "bytes", - "fallible-iterator", - "hmac", - "md-5", - "memchr", - "rand", - "sha2", - "stringprep", -] - -[[package]] -name = "postgres-types" -version = "0.2.4" -source = "git+https://github.com/prisma/rust-postgres?branch=pgbouncer-mode#a1a2dc6d9584deaf70a14293c428e7b6ca614d98" -dependencies = [ - "bytes", - "chrono", - "fallible-iterator", - "postgres-protocol", + "vcpkg", ] -[[package]] -name = "ppv-lite86" -version = "0.2.17" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b40af805b3121feab8a3c29f04d8ad262fa8e0561883e7653e024ae4479e6de" - [[package]] name = "pretty_env_logger" version = "0.5.0" @@ -1038,24 +986,14 @@ dependencies = [ ] [[package]] -name = "rand" -version = "0.8.5" +name = "r2d2" +version = "0.8.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93" dependencies = [ - "libc", - "rand_chacha", - "rand_core", -] - -[[package]] -name = "rand_chacha" -version = "0.3.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" -dependencies = [ - "ppv-lite86", - "rand_core", + "log", + "parking_lot", + "scheduled-thread-pool", ] [[package]] @@ -1063,9 +1001,6 @@ name = "rand_core" version = "0.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" -dependencies = [ - "getrandom", -] [[package]] name = "redox_syscall" @@ -1184,12 +1119,12 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ad4cc8da4ef723ed60bced201181d83791ad433213d8c24efffda1eec85d741" [[package]] -name = "schannel" -version = "0.1.22" +name = "scheduled-thread-pool" +version = "0.2.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c3733bf4cf7ea0880754e19cb5a462007c4a8c1914bff372ccc95b464f1df88" +checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19" dependencies = [ - "windows-sys", + "parking_lot", ] [[package]] @@ -1227,29 +1162,6 @@ dependencies = [ "cc", ] -[[package]] -name = "security-framework" -version = "2.9.2" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "05b64fb303737d99b81884b2c63433e9ae28abebe5eb5045dcdd175dc2ecf4de" -dependencies = [ - "bitflags 1.3.2", - "core-foundation", - "core-foundation-sys", - "libc", - "security-framework-sys", -] - -[[package]] -name = "security-framework-sys" -version = "2.9.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e932934257d3b408ed8f30db49d85ea163bfe74961f017f405b025af298f0c7a" -dependencies = [ - "core-foundation-sys", - "libc", -] - [[package]] name = "serde" version = "1.0.188" @@ -1301,6 +1213,15 @@ dependencies = [ "serde", ] +[[package]] +name = "serde_spanned" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96426c9936fd7a0124915f9185ea1d20aa9445cc9821142f0a73bc9207a2e186" +dependencies = [ + "serde", +] + [[package]] name = "serde_urlencoded" version = "0.7.1" @@ -1344,12 +1265,6 @@ dependencies = [ "libc", ] -[[package]] -name = "siphasher" -version = "0.3.11" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "38b58827f4464d87d377d175e90bf58eb00fd8716ff0a62f80356b5e61555d0d" - [[package]] name = "slab" version = "0.4.9" @@ -1391,17 +1306,6 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6e63cff320ae2c57904679ba7cb63280a3dc4613885beafb148ee7bf9aa9042d" -[[package]] -name = "stringprep" -version = "0.1.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bb41d74e231a107a1b4ee36bd1214b11285b77768d2e3824aedafa988fd36ee6" -dependencies = [ - "finl_unicode", - "unicode-bidi", - "unicode-normalization", -] - [[package]] name = "subtle" version = "2.5.0" @@ -1425,19 +1329,6 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2047c6ded9c721764247e62cd3b03c09ffc529b2ba5b10ec482ae507a4a70160" -[[package]] -name = "tempfile" -version = "3.8.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cb94d2f3cc536af71caac6b6fcebf65860b347e7ce0cc9ebe8f70d3e521054ef" -dependencies = [ - "cfg-if", - "fastrand", - "redox_syscall", - "rustix", - "windows-sys", -] - [[package]] name = "termcolor" version = "1.3.0" @@ -1493,50 +1384,37 @@ dependencies = [ ] [[package]] -name = "tokio-native-tls" -version = "0.3.1" +name = "toml" +version = "0.7.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bbae76ab933c85776efabc971569dd6119c580d8f5d448769dec1764bf796ef2" +checksum = "dd79e69d3b627db300ff956027cc6c3798cef26d22526befdfcd12feeb6d2257" dependencies = [ - "native-tls", - "tokio", + "serde", + "serde_spanned", + "toml_datetime", + "toml_edit", ] [[package]] -name = "tokio-postgres" -version = "0.7.7" -source = "git+https://github.com/prisma/rust-postgres?branch=pgbouncer-mode#a1a2dc6d9584deaf70a14293c428e7b6ca614d98" +name = "toml_datetime" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cda73e2f1397b1262d6dfdcef8aafae14d1de7748d66822d3bfeeb6d03e5e4b" dependencies = [ - "async-trait", - "byteorder", - "bytes", - "fallible-iterator", - "futures-channel", - "futures-util", - "log", - "parking_lot", - "percent-encoding", - "phf", - "pin-project-lite", - "postgres-protocol", - "postgres-types", - "socket2 0.5.4", - "tokio", - "tokio-util", + "serde", ] [[package]] -name = "tokio-util" -version = "0.7.9" +name = "toml_edit" +version = "0.19.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1d68074620f57a0b21594d9735eb2e98ab38b17f80d3fcb189fca266771ca60d" +checksum = "1b5bb770da30e5cbfde35a2d7b9b8a2c4b8ef89548a7a6aeab5c9a576e3e7421" dependencies = [ - "bytes", - "futures-core", - "futures-sink", - "pin-project-lite", - "tokio", - "tracing", + "indexmap", + "serde", + "serde_spanned", + "toml_datetime", + "winnow", ] [[package]] @@ -1694,20 +1572,19 @@ dependencies = [ "axum", "base64 0.13.1", "chrono", + "diesel", + "diesel_migrations", "dotenv", "futures", "hex", "jwt-compact", "log", - "native-tls", - "postgres-native-tls", "pretty_env_logger", "secp256k1", "serde", "serde_json", "sha2", "tokio", - "tokio-postgres", "tower-http", "ureq", ] @@ -1906,6 +1783,15 @@ version = "0.48.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed94fce61571a4006852b7389a063ab983c02eb1bb37b47f8272ce92d06d9538" +[[package]] +name = "winnow" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c2e3184b9c4e92ad5167ca73039d0c42476302ab603e2fec4487511f38ccefc" +dependencies = [ + "memchr", +] + [[package]] name = "zeroize" version = "1.6.0" diff --git a/Cargo.toml b/Cargo.toml index 4476964..4d921a7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -8,6 +8,7 @@ anyhow = "1.0" axum = { version = "0.6.16", features = ["headers"] } base64 = "0.13.1" chrono = { version = "0.4.26", features = ["serde"] } +diesel = { version = "2.1", features = ["postgres", "r2d2", "chrono", "numeric"] } dotenv = "0.15.0" futures = "0.3.28" hex = "0.4.3" @@ -19,9 +20,9 @@ sha2 = { version = "0.10", default-features = false } serde = { version = "^1.0", features = ["derive"] } serde_json = "1.0.67" tokio = { version = "1.12.0", features = ["full"] } -tokio-postgres = { git = "https://github.com/prisma/rust-postgres", branch = "pgbouncer-mode", features = ["with-chrono-0_4"] } tower-http = { version = "0.4.0", features = ["cors"] } ureq = { version = "2.5.0", features = ["json"] } -postgres-native-tls = { git = "https://github.com/prisma/rust-postgres", branch = "pgbouncer-mode" } -native-tls = "0.2" + +[dev-dependencies] +diesel_migrations = "2.1.0" \ No newline at end of file diff --git a/Dockerfile b/Dockerfile index e96c19c..326562d 100644 --- a/Dockerfile +++ b/Dockerfile @@ -8,9 +8,9 @@ RUN --mount=type=cache,target=/usr/local/cargo,from=rust:latest,source=/usr/loca cargo build --release && mv ./target/release/vss-rs ./vss-rs # Runtime image -FROM debian:bullseye-slim +FROM debian:bookworm-slim -RUN apt update && apt install -y openssl libpq-dev pkg-config libc6 openssl libssl-dev libpq5 ca-certificates +RUN apt update && apt install -y openssl libpq-dev pkg-config libc6 # Run as "app" user RUN useradd -ms /bin/bash app diff --git a/migrations/.keep b/migrations/.keep new file mode 100644 index 0000000..e69de29 diff --git a/migrations/00000000000000_diesel_initial_setup/down.sql b/migrations/00000000000000_diesel_initial_setup/down.sql new file mode 100644 index 0000000..a9f5260 --- /dev/null +++ b/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(); diff --git a/migrations/00000000000000_diesel_initial_setup/up.sql b/migrations/00000000000000_diesel_initial_setup/up.sql new file mode 100644 index 0000000..d68895b --- /dev/null +++ b/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; diff --git a/migrations/2023-09-23-030518_baseline/down.sql b/migrations/2023-09-23-030518_baseline/down.sql new file mode 100644 index 0000000..4b78a0e --- /dev/null +++ b/migrations/2023-09-23-030518_baseline/down.sql @@ -0,0 +1,9 @@ +-- 1. Drop the triggers +DROP TRIGGER IF EXISTS tr_set_dates_after_insert ON vss_db; +DROP TRIGGER IF EXISTS tr_set_dates_after_update ON vss_db; + +-- 2. Drop the trigger functions +DROP FUNCTION IF EXISTS set_created_date(); +DROP FUNCTION IF EXISTS set_updated_date(); + +DROP TABLE IF EXISTS vss_db; \ No newline at end of file diff --git a/migrations/2023-09-23-030518_baseline/up.sql b/migrations/2023-09-23-030518_baseline/up.sql new file mode 100644 index 0000000..63722d3 --- /dev/null +++ b/migrations/2023-09-23-030518_baseline/up.sql @@ -0,0 +1,79 @@ +CREATE TABLE vss_db +( + store_id TEXT NOT NULL CHECK (store_id != ''), + key TEXT NOT NULL, + value bytea, + version BIGINT NOT NULL, + created_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL, + updated_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP NOT NULL, + PRIMARY KEY (store_id, key) +); + +-- triggers to set dates automatically, generated by ChatGPT + +-- Function to set created_date and updated_date during INSERT +CREATE OR REPLACE FUNCTION set_created_date() + RETURNS TRIGGER AS +$$ +BEGIN + NEW.created_date := CURRENT_TIMESTAMP; + NEW.updated_date := CURRENT_TIMESTAMP; + RETURN NEW; +END; +$$ LANGUAGE plpgsql; + +-- Function to set updated_date during UPDATE +CREATE OR REPLACE FUNCTION set_updated_date() + RETURNS TRIGGER AS +$$ +BEGIN + NEW.updated_date := CURRENT_TIMESTAMP; + RETURN NEW; +END; +$$ LANGUAGE plpgsql; + +-- Trigger for INSERT operation on vss_db +CREATE TRIGGER tr_set_dates_after_insert + BEFORE INSERT + ON vss_db + FOR EACH ROW +EXECUTE FUNCTION set_created_date(); + +-- Trigger for UPDATE operation on vss_db +CREATE TRIGGER tr_set_dates_after_update + BEFORE UPDATE + ON vss_db + FOR EACH ROW +EXECUTE FUNCTION set_updated_date(); + +CREATE OR REPLACE FUNCTION upsert_vss_db( + p_store_id TEXT, + p_key TEXT, + p_value bytea, + p_version BIGINT +) RETURNS VOID AS +$$ +BEGIN + + WITH new_values (store_id, key, value, version) AS (VALUES (p_store_id, p_key, p_value, p_version)) + INSERT + INTO vss_db + (store_id, key, value, version) + SELECT new_values.store_id, + new_values.key, + new_values.value, + new_values.version + FROM new_values + LEFT JOIN vss_db AS existing + ON new_values.store_id = existing.store_id + AND new_values.key = existing.key + WHERE CASE + WHEN new_values.version >= 4294967295 THEN new_values.version >= COALESCE(existing.version, -1) + ELSE new_values.version > COALESCE(existing.version, -1) + END + ON CONFLICT (store_id, key) + DO UPDATE SET value = excluded.value, + version = excluded.version; + +END; +$$ LANGUAGE plpgsql; diff --git a/src/main.rs b/src/main.rs index c7fdc53..68722b7 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,12 +3,9 @@ use axum::headers::Origin; use axum::http::{request::Parts, HeaderValue, Method, StatusCode, Uri}; use axum::routing::{get, post, put}; use axum::{http, Extension, Router, TypedHeader}; -use native_tls::TlsConnector; -use postgres_native_tls::MakeTlsConnector; +use diesel::r2d2::{ConnectionManager, Pool}; +use diesel::PgConnection; use secp256k1::{All, PublicKey, Secp256k1}; -use std::str::FromStr; -use std::sync::Arc; -use tokio_postgres::{Client, Config}; use tower_http::cors::{AllowOrigin, CorsLayer}; mod auth; @@ -31,7 +28,7 @@ const ALLOWED_LOCALHOST: &str = "http://127.0.0.1:"; #[derive(Clone)] pub struct State { - pub client: Arc, + db_pool: Pool>, pub auth_key: PublicKey, pub secp: Secp256k1, } @@ -54,26 +51,18 @@ async fn main() -> anyhow::Result<()> { let auth_key_bytes = hex::decode(auth_key)?; let auth_key = PublicKey::from_slice(&auth_key_bytes)?; - let tls = TlsConnector::new()?; - let connector = MakeTlsConnector::new(tls); - - // Connect to the database. - let mut config = Config::from_str(&pg_url).unwrap(); - config.pgbouncer_mode(true); - let (client, connection) = config.connect(connector).await?; - - // The connection object performs the actual communication with the database, - // so spawn it off to run on its own. - tokio::spawn(async move { - if let Err(e) = connection.await { - panic!("db connection error: {e}"); - } - }); + // DB management + let manager = ConnectionManager::::new(&pg_url); + let db_pool = Pool::builder() + .max_size(10) // should be a multiple of 100, our database connection limit + .test_on_check_out(true) + .build(manager) + .expect("Could not build connection pool"); let secp = Secp256k1::new(); let state = State { - client: Arc::new(client), + db_pool, auth_key, secp, }; diff --git a/src/migration.rs b/src/migration.rs index c18d061..52acad6 100644 --- a/src/migration.rs +++ b/src/migration.rs @@ -6,6 +6,7 @@ use axum::headers::Authorization; use axum::http::StatusCode; use axum::{Extension, Json, TypedHeader}; use chrono::{DateTime, NaiveDateTime, Utc}; +use diesel::Connection; use log::{error, info}; use serde::{Deserialize, Deserializer}; use serde_json::json; @@ -77,19 +78,18 @@ pub async fn migration_impl(admin_key: String, state: &State) -> anyhow::Result< .send_string(&payload.to_string())?; let items: Vec = resp.into_json()?; + let mut conn = state.db_pool.get().unwrap(); + // Insert values into DB - for item in items.iter() { - if let Ok(value) = base64::decode(&item.value) { - VssItem::put_item( - &state.client, - &item.store_id, - &item.key, - &value, - item.version, - ) - .await?; + conn.transaction::<_, anyhow::Error, _>(|conn| { + for item in items.iter() { + if let Ok(value) = base64::decode(&item.value) { + VssItem::put_item(conn, &item.store_id, &item.key, &value, item.version)?; + } } - } + + Ok(()) + })?; if items.len() < limit { finished = true; diff --git a/src/models/mod.rs b/src/models/mod.rs index 5c474f7..d7f82a5 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -1,8 +1,25 @@ use crate::kv::KeyValue; +use diesel::prelude::*; +use diesel::sql_query; +use diesel::sql_types::{BigInt, Bytea, Text}; +use schema::vss_db; use serde::{Deserialize, Serialize}; -use tokio_postgres::{Client, Row}; -#[derive(Serialize, Deserialize, Debug, Clone, PartialEq)] +mod schema; + +#[derive( + QueryableByName, + Queryable, + Insertable, + AsChangeset, + Serialize, + Deserialize, + Debug, + Clone, + PartialEq, +)] +#[diesel(check_for_backend(diesel::pg::Pg))] +#[diesel(table_name = vss_db)] pub struct VssItem { pub store_id: String, pub key: String, @@ -19,177 +36,129 @@ impl VssItem { .map(|value| KeyValue::new(self.key, value, self.version)) } - pub async fn get_item( - client: &Client, - store_id: &String, - key: &String, + pub fn get_item( + conn: &mut PgConnection, + store_id: &str, + key: &str, ) -> anyhow::Result> { - client - .query_opt( - "SELECT * FROM vss_db WHERE store_id = $1 AND key = $2", - &[store_id, key], - ) - .await? - .map(row_to_vss_item) - .transpose() + Ok(vss_db::table + .filter(vss_db::store_id.eq(store_id)) + .filter(vss_db::key.eq(key)) + .first::(conn) + .optional()?) } - pub async fn put_item( - client: &Client, - store_id: &String, - key: &String, - value: &Vec, + pub fn put_item( + conn: &mut PgConnection, + store_id: &str, + key: &str, + value: &[u8], version: i64, ) -> anyhow::Result<()> { - client - .execute( - "SELECT upsert_vss_db($1, $2, $3, $4)", - &[store_id, key, value, &version], - ) - .await?; + sql_query("SELECT upsert_vss_db($1, $2, $3, $4)") + .bind::(store_id) + .bind::(key) + .bind::(value) + .bind::(version) + .execute(conn)?; Ok(()) } - pub async fn list_key_versions( - client: &Client, - store_id: &String, - prefix: Option<&String>, + pub fn list_key_versions( + conn: &mut PgConnection, + store_id: &str, + prefix: Option<&str>, ) -> anyhow::Result> { - let rows = match prefix { - Some(prefix) => client - .query( - "SELECT key, version FROM vss_db WHERE store_id = $1 AND key ILIKE $2 || '%'", - &[store_id, prefix], - ) - .await?, - None => { - client - .query( - "SELECT key, version FROM vss_db WHERE store_id = $1", - &[store_id], - ) - .await? - } + let table = vss_db::table + .filter(vss_db::store_id.eq(store_id)) + .select((vss_db::key, vss_db::version)); + + let res = match prefix { + None => table.load::<(String, i64)>(conn)?, + Some(prefix) => table + .filter(vss_db::key.ilike(format!("{prefix}%"))) + .load::<(String, i64)>(conn)?, }; - let res = rows - .into_iter() - .map(|row| { - let key: String = row.get(0); - let version: i64 = row.get(1); - - (key, version) - }) - .collect(); - Ok(res) } } -fn row_to_vss_item(row: Row) -> anyhow::Result { - let store_id: String = row.get(0); - let key: String = row.get(1); - let value: Option> = row.get(2); - let version: i64 = row.get(3); - let created_date: chrono::NaiveDateTime = row.get(4); - let updated_date: chrono::NaiveDateTime = row.get(5); - - Ok(VssItem { - store_id, - key, - value, - version, - created_date, - updated_date, - }) -} - #[cfg(test)] mod test { use super::*; use crate::State; + use diesel::r2d2::{ConnectionManager, Pool}; + use diesel_migrations::{embed_migrations, EmbeddedMigrations, MigrationHarness}; use secp256k1::Secp256k1; use std::str::FromStr; - use std::sync::Arc; - use tokio_postgres::{Config, NoTls}; const PUBKEY: &str = "04547d92b618856f4eda84a64ec32f1694c9608a3f9dc73e91f08b5daa087260164fbc9e2a563cf4c5ef9f4c614fd9dfca7582f8de429a4799a4b202fbe80a7db5"; + const MIGRATIONS: EmbeddedMigrations = embed_migrations!(); - async fn init_state() -> State { + fn init_state() -> State { dotenv::dotenv().ok(); - let pg_url = std::env::var("DATABASE_URL").expect("DATABASE_URL must be set"); - let mut config = Config::from_str(&pg_url).unwrap(); - config.pgbouncer_mode(true); - // Connect to the database. - let (client, connection) = config.connect(NoTls).await.unwrap(); - - // The connection object performs the actual communication with the database, - // so spawn it off to run on its own. - tokio::spawn(async move { - if let Err(e) = connection.await { - eprintln!("db connection error: {e}"); - } - }); - - client - .simple_query("DROP TABLE IF EXISTS vss_db") - .await - .unwrap(); - client - .simple_query(include_str!("migration_baseline.sql")) - .await - .unwrap(); + let url = std::env::var("DATABASE_URL").expect("DATABASE_URL must be set"); + let manager = ConnectionManager::::new(url); + let db_pool = Pool::builder() + .max_size(10) + .test_on_check_out(true) + .build(manager) + .expect("Could not build connection pool"); + + // run migrations + let mut connection = db_pool.get().unwrap(); + connection + .run_pending_migrations(MIGRATIONS) + .expect("migrations could not run"); let auth_key = secp256k1::PublicKey::from_str(PUBKEY).unwrap(); let secp = Secp256k1::new(); State { - client: Arc::new(client), + db_pool, auth_key, secp, } } - async fn clear_database(state: &State) { - state - .client - .simple_query("DROP TABLE vss_db") - .await - .unwrap(); + fn clear_database(state: &State) { + let conn = &mut state.db_pool.get().unwrap(); + + conn.transaction::<_, anyhow::Error, _>(|conn| { + diesel::delete(vss_db::table).execute(conn)?; + Ok(()) + }) + .unwrap(); } #[tokio::test] async fn test_vss_flow() { - let state = init_state().await; + let state = init_state(); + clear_database(&state); - let store_id = "test_store_id".to_string(); - let key = "test".to_string(); - let value: Vec = vec![1, 2, 3, 4, 5]; + let store_id = "test_store_id"; + let key = "test"; + let value = [1, 2, 3]; let version = 0; - VssItem::put_item(&state.client, &store_id, &key, &value, version) - .await - .unwrap(); + let mut conn = state.db_pool.get().unwrap(); + VssItem::put_item(&mut conn, store_id, key, &value, version).unwrap(); - let versions = VssItem::list_key_versions(&state.client, &store_id, None) - .await - .unwrap(); + let versions = VssItem::list_key_versions(&mut conn, store_id, None).unwrap(); assert_eq!(versions.len(), 1); assert_eq!(versions[0].0, key); assert_eq!(versions[0].1, version); - let new_value = vec![6, 7, 8, 9, 10]; + let new_value = [4, 5, 6]; let new_version = version + 1; - VssItem::put_item(&state.client, &store_id, &key, &new_value, new_version) - .await - .unwrap(); + VssItem::put_item(&mut conn, store_id, key, &new_value, new_version).unwrap(); - let item = VssItem::get_item(&state.client, &store_id, &key) - .await + let item = VssItem::get_item(&mut conn, store_id, key) .unwrap() .unwrap(); @@ -198,24 +167,23 @@ mod test { assert_eq!(item.value.unwrap(), new_value); assert_eq!(item.version, new_version); - clear_database(&state).await; + clear_database(&state); } #[tokio::test] async fn test_max_version_number() { - let state = init_state().await; + let state = init_state(); + clear_database(&state); - let store_id = "max_test_store_id".to_string(); - let key = "max_test".to_string(); - let value = vec![1, 2, 3, 4, 5]; + let store_id = "max_test_store_id"; + let key = "max_test"; + let value = [1, 2, 3]; let version = u32::MAX as i64; - VssItem::put_item(&state.client, &store_id, &key, &value, version) - .await - .unwrap(); + let mut conn = state.db_pool.get().unwrap(); + VssItem::put_item(&mut conn, store_id, key, &value, version).unwrap(); - let item = VssItem::get_item(&state.client, &store_id, &key) - .await + let item = VssItem::get_item(&mut conn, store_id, key) .unwrap() .unwrap(); @@ -223,14 +191,11 @@ mod test { assert_eq!(item.key, key); assert_eq!(item.value.unwrap(), value); - let new_value = vec![6, 7, 8, 9, 10]; + let new_value = [4, 5, 6]; - VssItem::put_item(&state.client, &store_id, &key, &new_value, version) - .await - .unwrap(); + VssItem::put_item(&mut conn, store_id, key, &new_value, version).unwrap(); - let item = VssItem::get_item(&state.client, &store_id, &key) - .await + let item = VssItem::get_item(&mut conn, store_id, key) .unwrap() .unwrap(); @@ -238,48 +203,38 @@ mod test { assert_eq!(item.key, key); assert_eq!(item.value.unwrap(), new_value); - clear_database(&state).await; + clear_database(&state); } #[tokio::test] async fn test_list_key_versions() { - let state = init_state().await; + let state = init_state(); + clear_database(&state); - let store_id = "list_kv_test_store_id".to_string(); - let key = "kv_test".to_string(); - let key1 = "other_kv_test".to_string(); - let value = vec![1, 2, 3, 4, 5]; + let store_id = "list_kv_test_store_id"; + let key = "kv_test"; + let key1 = "other_kv_test"; + let value = [1, 2, 3]; let version = 0; - VssItem::put_item(&state.client, &store_id, &key, &value, version) - .await - .unwrap(); + let mut conn = state.db_pool.get().unwrap(); + VssItem::put_item(&mut conn, store_id, key, &value, version).unwrap(); - VssItem::put_item(&state.client, &store_id, &key1, &value, version) - .await - .unwrap(); + VssItem::put_item(&mut conn, store_id, key1, &value, version).unwrap(); - let versions = VssItem::list_key_versions(&state.client, &store_id, None) - .await - .unwrap(); + let versions = VssItem::list_key_versions(&mut conn, store_id, None).unwrap(); assert_eq!(versions.len(), 2); - let versions = - VssItem::list_key_versions(&state.client, &store_id, Some(&"kv".to_string())) - .await - .unwrap(); + let versions = VssItem::list_key_versions(&mut conn, store_id, Some("kv")).unwrap(); assert_eq!(versions.len(), 1); assert_eq!(versions[0].0, key); assert_eq!(versions[0].1, version); - let versions = - VssItem::list_key_versions(&state.client, &store_id, Some(&"other".to_string())) - .await - .unwrap(); + let versions = VssItem::list_key_versions(&mut conn, store_id, Some("other")).unwrap(); assert_eq!(versions.len(), 1); assert_eq!(versions[0].0, key1); assert_eq!(versions[0].1, version); - clear_database(&state).await; + clear_database(&state); } } diff --git a/src/models/put_item.sql b/src/models/put_item.sql deleted file mode 100644 index feb7b63..0000000 --- a/src/models/put_item.sql +++ /dev/null @@ -1,17 +0,0 @@ -WITH new_values (store_id, key, value, version) AS (VALUES ($1, $2, $3, $4)) -INSERT -INTO vss_db - (store_id, key, value, version) -SELECT new_values.store_id, - new_values.key, - new_values.value, - new_values.version -FROM new_values - LEFT JOIN vss_db AS existing ON new_values.store_id = existing.store_id AND new_values.key = existing.key -WHERE CASE - WHEN new_values.version >= 4294967295 THEN new_values.version >= COALESCE(existing.version, -1) - ELSE new_values.version > COALESCE(existing.version, -1) - END -ON CONFLICT (store_id, key) - DO UPDATE SET value = excluded.value, - version = excluded.version; diff --git a/src/models/schema.rs b/src/models/schema.rs new file mode 100644 index 0000000..de1aed0 --- /dev/null +++ b/src/models/schema.rs @@ -0,0 +1,12 @@ +// @generated automatically by Diesel CLI. + +diesel::table! { + vss_db (store_id, key) { + store_id -> Text, + key -> Text, + value -> Nullable, + version -> Int8, + created_date -> Timestamp, + updated_date -> Timestamp, + } +} diff --git a/src/routes.rs b/src/routes.rs index a113b7e..7003a60 100644 --- a/src/routes.rs +++ b/src/routes.rs @@ -6,6 +6,7 @@ use axum::headers::authorization::Bearer; use axum::headers::{Authorization, Origin}; use axum::http::StatusCode; use axum::{Extension, Json, TypedHeader}; +use diesel::Connection; use log::{debug, error, trace}; use serde::{Deserialize, Serialize}; use serde_json::{json, Value}; @@ -39,7 +40,9 @@ pub async fn get_object_impl( trace!("get_object_impl: {req:?}"); let store_id = req.store_id.expect("must have"); - let item = VssItem::get_item(&state.client, &store_id, &req.key).await?; + let mut conn = state.db_pool.get()?; + + let item = VssItem::get_item(&mut conn, &store_id, &req.key)?; Ok(item.and_then(|i| i.into_kv())) } @@ -99,10 +102,15 @@ pub async fn put_objects_impl(req: PutObjectsRequest, state: &State) -> anyhow:: let store_id = req.store_id.expect("must have"); - // todo use transaction - for kv in req.transaction_items { - VssItem::put_item(&state.client, &store_id, &kv.key, &kv.value.0, kv.version).await?; - } + let mut conn = state.db_pool.get()?; + + conn.transaction::<_, anyhow::Error, _>(|conn| { + for kv in req.transaction_items { + VssItem::put_item(conn, &store_id, &kv.key, &kv.value.0, kv.version)?; + } + + Ok(()) + })?; Ok(()) } @@ -140,8 +148,9 @@ pub async fn list_key_versions_impl( // todo pagination let store_id = req.store_id.expect("must have"); - let versions = - VssItem::list_key_versions(&state.client, &store_id, req.key_prefix.as_ref()).await?; + let mut conn = state.db_pool.get()?; + + let versions = VssItem::list_key_versions(&mut conn, &store_id, req.key_prefix.as_deref())?; let json = versions .into_iter() From 10124de30d3a736100b1c0731661583add5c56f4 Mon Sep 17 00:00:00 2001 From: benthecarman Date: Sat, 23 Sep 2023 02:36:26 -0500 Subject: [PATCH 5/7] drill config --- drill.yml | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) create mode 100644 drill.yml diff --git a/drill.yml b/drill.yml new file mode 100644 index 0000000..a9e55e1 --- /dev/null +++ b/drill.yml @@ -0,0 +1,27 @@ +base: 'https://vss-staging.fly.dev' +concurrency: 250 +iterations: 2000 + +plan: + - name: Health Check + request: + method: GET + url: /health-check + + - name: Put Objects + request: + method: PUT + url: /putObjects + body: '{"transaction_items": [{"key": "key", "value": [0, 1, 2], "version": 0}]}' + headers: + Content-Type: 'application/json' + Authorization: Bearer eyJhbGciOiJFUzI1NksiLCJraWQiOiIwMzU0N2Q5MmI2MTg4NTZmNGVkYTg0YTY0ZWMzMmYxNjk0Yzk2MDhhM2Y5ZGM3M2U5MWYwOGI1ZGFhMDg3MjYwMTYifQ.eyJleHAiOjE2OTU0NTc4MjgsIm5iZiI6MTY5NTQ1MDYyOCwiaWF0IjoxNjk1NDU0MjI4LCJzdWIiOiIwMjlhOGYwY2MxZWNmZDU5NWNjNzQyYWU5OGU4NDZlNDRjODdmMWJjYWVjNzcxOTZhOTc3MzFjNzllMmJmZDI1ODUifQ.llJlMub-FWU9tgmQRMchziyg6jLGOgKYPq5DOm4dOGUqMCtDyRQX--ILBLhgkHTVhTy0EFyYu0x4clVcb7kV0A + + - name: Get Object + request: + method: POST + url: /getObject + body: '{"key": "key"}' + headers: + Content-Type: 'application/json' + Authorization: Bearer eyJhbGciOiJFUzI1NksiLCJraWQiOiIwMzU0N2Q5MmI2MTg4NTZmNGVkYTg0YTY0ZWMzMmYxNjk0Yzk2MDhhM2Y5ZGM3M2U5MWYwOGI1ZGFhMDg3MjYwMTYifQ.eyJleHAiOjE2OTU0NTc4MjgsIm5iZiI6MTY5NTQ1MDYyOCwiaWF0IjoxNjk1NDU0MjI4LCJzdWIiOiIwMjlhOGYwY2MxZWNmZDU5NWNjNzQyYWU5OGU4NDZlNDRjODdmMWJjYWVjNzcxOTZhOTc3MzFjNzllMmJmZDI1ODUifQ.llJlMub-FWU9tgmQRMchziyg6jLGOgKYPq5DOm4dOGUqMCtDyRQX--ILBLhgkHTVhTy0EFyYu0x4clVcb7kV0A From 44e9fcfdb14f6bc99606876321ff9125876e63c2 Mon Sep 17 00:00:00 2001 From: benthecarman Date: Sat, 23 Sep 2023 16:47:46 -0500 Subject: [PATCH 6/7] Remove old migration file --- src/models/migration_baseline.sql | 76 ------------------------------- 1 file changed, 76 deletions(-) delete mode 100644 src/models/migration_baseline.sql diff --git a/src/models/migration_baseline.sql b/src/models/migration_baseline.sql deleted file mode 100644 index 37574b7..0000000 --- a/src/models/migration_baseline.sql +++ /dev/null @@ -1,76 +0,0 @@ -CREATE TABLE IF NOT EXISTS vss_db -( - store_id TEXT NOT NULL CHECK (store_id != ''), - key TEXT NOT NULL, - value bytea, - version BIGINT NOT NULL, - created_date TIMESTAMP DEFAULT '2023-07-13'::TIMESTAMP NOT NULL, - updated_date TIMESTAMP DEFAULT '2023-07-13'::TIMESTAMP NOT NULL, - PRIMARY KEY (store_id, key) -); - --- triggers to set dates automatically, generated by ChatGPT - --- Function to set created_date and updated_date during INSERT -CREATE OR REPLACE FUNCTION set_created_date() - RETURNS TRIGGER AS $$ -BEGIN - NEW.created_date := CURRENT_TIMESTAMP; - NEW.updated_date := CURRENT_TIMESTAMP; - RETURN NEW; -END; -$$ LANGUAGE plpgsql; - --- Function to set updated_date during UPDATE -CREATE OR REPLACE FUNCTION set_updated_date() - RETURNS TRIGGER AS $$ -BEGIN - NEW.updated_date := CURRENT_TIMESTAMP; - RETURN NEW; -END; -$$ LANGUAGE plpgsql; - --- Trigger for INSERT operation on vss_db -CREATE TRIGGER tr_set_dates_after_insert - BEFORE INSERT ON vss_db - FOR EACH ROW -EXECUTE FUNCTION set_created_date(); - --- Trigger for UPDATE operation on vss_db -CREATE TRIGGER tr_set_dates_after_update - BEFORE UPDATE ON vss_db - FOR EACH ROW -EXECUTE FUNCTION set_updated_date(); - --- upsert function -CREATE OR REPLACE FUNCTION upsert_vss_db( - p_store_id TEXT, - p_key TEXT, - p_value bytea, - p_version BIGINT -) RETURNS VOID AS $$ -BEGIN - - WITH new_values (store_id, key, value, version) AS (VALUES (p_store_id, p_key, p_value, p_version)) - INSERT - INTO vss_db - (store_id, key, value, version) - SELECT new_values.store_id, - new_values.key, - new_values.value, - new_values.version - FROM new_values - LEFT JOIN vss_db AS existing - ON new_values.store_id = existing.store_id - AND new_values.key = existing.key - WHERE CASE - WHEN new_values.version >= 4294967295 THEN new_values.version >= COALESCE(existing.version, -1) - ELSE new_values.version > COALESCE(existing.version, -1) - END - ON CONFLICT (store_id, key) - DO UPDATE SET value = excluded.value, - version = excluded.version; - -END; -$$ LANGUAGE plpgsql; - From 2451e708b8bc4ff3eb5934e45d98c0fb30df5fc9 Mon Sep 17 00:00:00 2001 From: Tony Giorgio Date: Sat, 23 Sep 2023 20:02:53 -0500 Subject: [PATCH 7/7] Multiple Version asserts with drill --- README.md | 6 ++ drill.yml | 175 +++++++++++++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 174 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index a8f864d..3730de6 100644 --- a/README.md +++ b/README.md @@ -8,3 +8,9 @@ You need a postgres database and an authentication key. These can be set in the and `AUTH_KEY` respectively. This can be set in a `.env` file in the root of the project. To run the server, run `cargo run --release` in the root of the project. + +## Stress testing + +``` +AUTH_TOKEN=ey... drill --benchmark drill.yml -o 30 +``` diff --git a/drill.yml b/drill.yml index a9e55e1..8aa278a 100644 --- a/drill.yml +++ b/drill.yml @@ -1,5 +1,5 @@ base: 'https://vss-staging.fly.dev' -concurrency: 250 +concurrency: 125 iterations: 2000 plan: @@ -8,20 +8,181 @@ plan: method: GET url: /health-check - - name: Put Objects + # Generate a random UUID and assign it to the test_key variable + - name: Generate unique test key + exec: + command: "echo \"drill_test_$(uuidgen)\"" + assign: test_key + + # Get the Object to fetch the initial version + - name: Get Object Initial (1) - '{{ test_key }}' + request: + method: POST + url: /getObject + body: '{"key": "{{ test_key }}"}' + headers: + Content-Type: 'application/json' + Authorization: Bearer {{ AUTH_TOKEN }} + assign: initial_get_object_response + + - name: Extract and increment version (1) - '{{ test_key }}' + exec: + command: "input='{{ initial_get_object_response.body }}'; [ -z \"$input\" ] && echo 0 || echo \"$input\" | jq '.version? // empty | if type == \"number\" then . + 1 else 0 end'" + assign: version + + - name: Put Objects (1) - '{{ test_key }}' request: method: PUT url: /putObjects - body: '{"transaction_items": [{"key": "key", "value": [0, 1, 2], "version": 0}]}' + body: '{"transaction_items": [{"key": "{{ test_key }}", "value": [0, 1, 2], "version": {{ version }} }]}' headers: Content-Type: 'application/json' - Authorization: Bearer eyJhbGciOiJFUzI1NksiLCJraWQiOiIwMzU0N2Q5MmI2MTg4NTZmNGVkYTg0YTY0ZWMzMmYxNjk0Yzk2MDhhM2Y5ZGM3M2U5MWYwOGI1ZGFhMDg3MjYwMTYifQ.eyJleHAiOjE2OTU0NTc4MjgsIm5iZiI6MTY5NTQ1MDYyOCwiaWF0IjoxNjk1NDU0MjI4LCJzdWIiOiIwMjlhOGYwY2MxZWNmZDU5NWNjNzQyYWU5OGU4NDZlNDRjODdmMWJjYWVjNzcxOTZhOTc3MzFjNzllMmJmZDI1ODUifQ.llJlMub-FWU9tgmQRMchziyg6jLGOgKYPq5DOm4dOGUqMCtDyRQX--ILBLhgkHTVhTy0EFyYu0x4clVcb7kV0A + Authorization: Bearer {{ AUTH_TOKEN }} - - name: Get Object + - name: Get Object after Put (1) - '{{ test_key }}' request: method: POST url: /getObject - body: '{"key": "key"}' + body: '{"key": "{{ test_key }}"}' headers: Content-Type: 'application/json' - Authorization: Bearer eyJhbGciOiJFUzI1NksiLCJraWQiOiIwMzU0N2Q5MmI2MTg4NTZmNGVkYTg0YTY0ZWMzMmYxNjk0Yzk2MDhhM2Y5ZGM3M2U5MWYwOGI1ZGFhMDg3MjYwMTYifQ.eyJleHAiOjE2OTU0NTc4MjgsIm5iZiI6MTY5NTQ1MDYyOCwiaWF0IjoxNjk1NDU0MjI4LCJzdWIiOiIwMjlhOGYwY2MxZWNmZDU5NWNjNzQyYWU5OGU4NDZlNDRjODdmMWJjYWVjNzcxOTZhOTc3MzFjNzllMmJmZDI1ODUifQ.llJlMub-FWU9tgmQRMchziyg6jLGOgKYPq5DOm4dOGUqMCtDyRQX--ILBLhgkHTVhTy0EFyYu0x4clVcb7kV0A + Authorization: Bearer {{ AUTH_TOKEN }} + assign: get_object_response + + # Basic assertion to make sure we got a good response with the right key + - name: Extract key from response - '{{ test_key }}' + exec: + command: "echo '{{ get_object_response.body }}' | jq '.key'" + assign: retrieved_key + + - name: Compare test_key and retrieved_key - '{{ test_key }}' + exec: + command: "if [ \"{{ test_key }}\" = \"{{ retrieved_key }}\" ]; then echo 'true'; else echo 'false'; fi" + assign: key_comparison_result + + - name: Assert keys match - '{{ test_key }}' + assert: + key: key_comparison_result + value: "true" + + # Compare the version from "Get Object response" with the assigned version + - name: Compare versions with external command (1) - '{{ test_key }}' + exec: + command: "echo '{{ get_object_response.body }}' | jq --arg version '{{ version }}' '.version == ($version | tonumber)'" + assign: version_match_result + + # Assert that the result from the comparison is true + - name: Assert versions match (1) - '{{ test_key }}' + assert: + key: version_match_result + value: "true" + + # + ## Do this a 2nd time with a bigger version + # + - name: Extract and increment version (2) - '{{ test_key }}' + exec: + command: "echo $(({{ version }} + 1))" + assign: version + + - name: Put Objects (2) - '{{ test_key }}' + request: + method: PUT + url: /putObjects + body: '{"transaction_items": [{"key": "{{ test_key }}", "value": [0, 1, 2], "version": {{ version }} }]}' + headers: + Content-Type: 'application/json' + Authorization: Bearer {{ AUTH_TOKEN }} + + - name: Get Object after Put (2) - '{{ test_key }}' + request: + method: POST + url: /getObject + body: '{"key": "{{ test_key }}"}' + headers: + Content-Type: 'application/json' + Authorization: Bearer {{ AUTH_TOKEN }} + assign: get_object_response + + # Compare the version from "Get Object response" with the assigned version + - name: Compare versions with external command (2) - '{{ test_key }}' + exec: + command: "echo '{{ get_object_response.body }}' | jq --arg version '{{ version }}' '.version == ($version | tonumber)'" + assign: version_match_result + + # Basic assertion to make sure we got a good response with the right key + - name: Extract key from response - '{{ test_key }}' + exec: + command: "echo '{{ get_object_response.body }}' | jq '.key'" + assign: retrieved_key + + - name: Compare test_key and retrieved_key - '{{ test_key }}' + exec: + command: "if [ \"{{ test_key }}\" = \"{{ retrieved_key }}\" ]; then echo 'true'; else echo 'false'; fi" + assign: key_comparison_result + + - name: Assert keys match - '{{ test_key }}' + assert: + key: key_comparison_result + value: "true" + + # Assert that the result from the comparison is true + - name: Assert versions match (2) - '{{ test_key }}' + assert: + key: version_match_result + value: "true" + + # + ## Do this a third time with a bigger version + # + - name: Extract and increment version (3) - '{{ test_key }}' + exec: + command: "echo $(({{ version }} + 1))" + assign: version + + - name: Put Objects (3) - '{{ test_key }}' + request: + method: PUT + url: /putObjects + body: '{"transaction_items": [{"key": "{{ test_key }}", "value": [0, 1, 2], "version": {{ version }} }]}' + headers: + Content-Type: 'application/json' + Authorization: Bearer {{ AUTH_TOKEN }} + + - name: Get Object after Put (3) - '{{ test_key }}' + request: + method: POST + url: /getObject + body: '{"key": "{{ test_key }}"}' + headers: + Content-Type: 'application/json' + Authorization: Bearer {{ AUTH_TOKEN }} + assign: get_object_response + + # Basic assertion to make sure we got a good response with the right key + - name: Extract key from response - '{{ test_key }}' + exec: + command: "echo '{{ get_object_response.body }}' | jq '.key'" + assign: retrieved_key + + - name: Compare test_key and retrieved_key - '{{ test_key }}' + exec: + command: "if [ \"{{ test_key }}\" = \"{{ retrieved_key }}\" ]; then echo 'true'; else echo 'false'; fi" + assign: key_comparison_result + + - name: Assert keys match - '{{ test_key }}' + assert: + key: key_comparison_result + value: "true" + + # Compare the version from "Get Object response" with the assigned version + - name: Compare versions with external command (3) - '{{ test_key }}' + exec: + command: "echo '{{ get_object_response.body }}' | jq --arg version '{{ version }}' '.version == ($version | tonumber)'" + assign: version_match_result + + # Assert that the result from the comparison is true + - name: Assert versions match (3) - '{{ test_key }}' + assert: + key: version_match_result + value: "true"