diff --git a/.gitignore b/.gitignore index ea8c4bf..355560c 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,3 @@ /target +*.swp +*.swo diff --git a/Cargo.lock b/Cargo.lock index 992a492..198c00a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1053,7 +1053,7 @@ dependencies = [ [[package]] name = "khm" -version = "0.2.2" +version = "0.4.0" dependencies = [ "actix-web", "base64 0.21.7", diff --git a/Cargo.toml b/Cargo.toml index 96f8962..4573adc 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,8 +1,8 @@ [package] name = "khm" -version = "0.3.0" +version = "0.4.0" edition = "2021" -authors = ["AB ", "ChatGPT-4o"] +authors = ["AB "] [dependencies] actix-web = "4" diff --git a/src/db.rs b/src/db.rs new file mode 100644 index 0000000..c38df94 --- /dev/null +++ b/src/db.rs @@ -0,0 +1,286 @@ +use log::{info}; +use std::collections::HashMap; +use std::collections::HashSet; +use tokio_postgres::Client; +use crate::server::SshKey; + +// Структура для хранения статистики обработки ключей +pub struct KeyInsertStats { + pub total: usize, // Общее количество полученных ключей + pub inserted: usize, // Количество новых ключей + pub updated: usize, // Количество обновленных ключей + pub unchanged: usize, // Количество неизмененных ключей + pub key_id_map: Vec<(SshKey, i32)>, // Связь ключей с их ID в базе +} + +pub async fn initialize_db_schema(client: &Client) -> Result<(), tokio_postgres::Error> { + info!("Checking and initializing database schema if needed"); + + // Check if tables exist by querying information_schema + let tables_exist = client + .query( + "SELECT EXISTS ( + SELECT FROM information_schema.tables + WHERE table_schema = 'public' + AND table_name = 'keys' + ) AND EXISTS ( + SELECT FROM information_schema.tables + WHERE table_schema = 'public' + AND table_name = 'flows' + )", + &[], + ) + .await? + .get(0) + .map(|row| row.get::<_, bool>(0)) + .unwrap_or(false); + + if !tables_exist { + info!("Database schema doesn't exist. Creating tables..."); + + // Create the keys table + client + .execute( + "CREATE TABLE IF NOT EXISTS public.keys ( + key_id SERIAL PRIMARY KEY, + host VARCHAR(255) NOT NULL, + key TEXT NOT NULL, + updated TIMESTAMP WITH TIME ZONE NOT NULL, + CONSTRAINT unique_host_key UNIQUE (host, key) + )", + &[], + ) + .await?; + + // Create the flows table + client + .execute( + "CREATE TABLE IF NOT EXISTS public.flows ( + flow_id SERIAL PRIMARY KEY, + name VARCHAR(255) NOT NULL, + key_id INTEGER NOT NULL, + CONSTRAINT fk_key + FOREIGN KEY(key_id) + REFERENCES public.keys(key_id) + ON DELETE CASCADE, + CONSTRAINT unique_flow_key UNIQUE (name, key_id) + )", + &[], + ) + .await?; + + // Create an index for faster lookups + client + .execute( + "CREATE INDEX IF NOT EXISTS idx_flows_name ON public.flows(name)", + &[], + ) + .await?; + + info!("Database schema created successfully"); + } else { + info!("Database schema already exists"); + } + + Ok(()) +} + +pub async fn batch_insert_keys( + client: &Client, + keys: &[SshKey], +) -> Result { + if keys.is_empty() { + return Ok(KeyInsertStats { + total: 0, + inserted: 0, + updated: 0, + unchanged: 0, + key_id_map: Vec::new(), + }); + } + + // Подготавливаем массивы для пакетной вставки + let mut host_values: Vec<&str> = Vec::with_capacity(keys.len()); + let mut key_values: Vec<&str> = Vec::with_capacity(keys.len()); + + for key in keys { + host_values.push(&key.server); + key_values.push(&key.public_key); + } + + // Сначала проверяем, какие ключи уже существуют в базе + let mut existing_keys = HashMap::new(); + let mut key_query = String::from("SELECT host, key, key_id FROM public.keys WHERE "); + + for i in 0..keys.len() { + if i > 0 { + key_query.push_str(" OR "); + } + key_query.push_str(&format!("(host = ${} AND key = ${})", i * 2 + 1, i * 2 + 2)); + } + + let mut params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::with_capacity(keys.len() * 2); + for i in 0..keys.len() { + params.push(&host_values[i]); + params.push(&key_values[i]); + } + + let rows = client.query(&key_query, ¶ms[..]).await?; + + for row in rows { + let host: String = row.get(0); + let key: String = row.get(1); + let key_id: i32 = row.get(2); + existing_keys.insert((host, key), key_id); + } + + // Определяем, какие ключи нужно вставить, а какие уже существуют + let mut keys_to_insert = Vec::new(); + let mut unchanged_keys = Vec::new(); + + for key in keys { + let key_tuple = (key.server.clone(), key.public_key.clone()); + if existing_keys.contains_key(&key_tuple) { + unchanged_keys.push((key.clone(), *existing_keys.get(&key_tuple).unwrap())); + } else { + keys_to_insert.push(key.clone()); + } + } + + let mut inserted_keys = Vec::new(); + + // Если есть ключи для вставки, выполняем вставку + if !keys_to_insert.is_empty() { + let mut insert_sql = String::from( + "INSERT INTO public.keys (host, key, updated) VALUES " + ); + + let mut insert_params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::new(); + let mut param_count = 1; + + for (i, key) in keys_to_insert.iter().enumerate() { + if i > 0 { + insert_sql.push_str(", "); + } + insert_sql.push_str(&format!("(${}, ${}, NOW())", param_count, param_count + 1)); + insert_params.push(&key.server); + insert_params.push(&key.public_key); + param_count += 2; + } + + insert_sql.push_str(" RETURNING key_id, host, key"); + + let inserted_rows = client.query(&insert_sql, &insert_params[..]).await?; + + for row in inserted_rows { + let host: String = row.get(1); + let key_text: String = row.get(2); + let key_id: i32 = row.get(0); + + if let Some(orig_key) = keys_to_insert.iter().find(|k| k.server == host && k.public_key == key_text) { + inserted_keys.push((orig_key.clone(), key_id)); + } + } + } + + // Сохраняем количество элементов перед объединением + let inserted_count = inserted_keys.len(); + let unchanged_count = unchanged_keys.len(); + + // Комбинируем результаты и формируем статистику + let mut key_id_map = Vec::with_capacity(unchanged_count + inserted_count); + key_id_map.extend(unchanged_keys); + key_id_map.extend(inserted_keys); + + let stats = KeyInsertStats { + total: keys.len(), + inserted: inserted_count, + updated: 0, // В этой версии мы не обновляем существующие ключи + unchanged: unchanged_count, + key_id_map, + }; + + info!("Keys stats: received={}, new={}, unchanged={}", + stats.total, stats.inserted, stats.unchanged); + + Ok(stats) +} + +pub async fn batch_insert_flow_keys( + client: &Client, + flow_name: &str, + key_ids: &[i32], +) -> Result { + if key_ids.is_empty() { + info!("No keys to associate with flow '{}'", flow_name); + return Ok(0); + } + + // Сначала проверим, какие связи уже существуют + let mut existing_query = String::from( + "SELECT key_id FROM public.flows WHERE name = $1 AND key_id IN (" + ); + + for i in 0..key_ids.len() { + if i > 0 { + existing_query.push_str(", "); + } + existing_query.push_str(&format!("${}", i + 2)); + } + existing_query.push_str(")"); + + let mut params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::with_capacity(key_ids.len() + 1); + params.push(&flow_name); + for key_id in key_ids { + params.push(key_id); + } + + let rows = client.query(&existing_query, ¶ms[..]).await?; + + let mut existing_associations = HashSet::new(); + for row in rows { + let key_id: i32 = row.get(0); + existing_associations.insert(key_id); + } + + // Фильтруем только те ключи, которые еще не связаны с потоком + let new_key_ids: Vec<&i32> = key_ids.iter() + .filter(|&id| !existing_associations.contains(id)) + .collect(); + + if new_key_ids.is_empty() { + info!("All {} keys are already associated with flow '{}'", key_ids.len(), flow_name); + return Ok(0); + } + + // Строим SQL запрос с множественными значениями только для новых связей + let mut sql = String::from( + "INSERT INTO public.flows (name, key_id) VALUES " + ); + + for i in 0..new_key_ids.len() { + if i > 0 { + sql.push_str(", "); + } + sql.push_str(&format!("($1, ${})", i + 2)); + } + + sql.push_str(" ON CONFLICT (name, key_id) DO NOTHING"); + + // Подготавливаем параметры для запроса + let mut insert_params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::with_capacity(new_key_ids.len() + 1); + insert_params.push(&flow_name); + for key_id in &new_key_ids { + insert_params.push(*key_id); + } + + // Выполняем запрос + let affected = client.execute(&sql, &insert_params[..]).await?; + + let affected_usize = affected as usize; + + info!("Added {} new key-flow associations for flow '{}' (skipped {} existing)", + affected_usize, flow_name, existing_associations.len()); + + Ok(affected_usize) +} diff --git a/src/main.rs b/src/main.rs index dc78228..dfaadab 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,4 +1,5 @@ mod client; +mod db; mod server; use clap::Parser; diff --git a/src/server.rs b/src/server.rs index 247ffac..1967049 100644 --- a/src/server.rs +++ b/src/server.rs @@ -6,6 +6,8 @@ use std::collections::HashMap; use std::sync::{Arc, Mutex}; use tokio_postgres::{Client, NoTls}; +use crate::db; + #[derive(Serialize, Deserialize, Clone, Debug)] pub struct SshKey { pub server: String, @@ -148,6 +150,8 @@ pub async fn add_keys( return HttpResponse::Forbidden().body("Flow ID not allowed"); } + // Проверяем формат SSH ключей + let mut valid_keys = Vec::new(); for new_key in new_keys.iter() { if !is_valid_ssh_key(&new_key.public_key) { error!("Invalid SSH key format for server: {}", new_key.server); @@ -156,23 +160,42 @@ pub async fn add_keys( new_key.server )); } - - match insert_key_if_not_exists(&db_client, new_key).await { - Ok(key_id) => { - if let Err(e) = insert_flow_key(&db_client, &flow_id_str, key_id).await { - error!("Failed to insert flow key into database: {}", e); - return HttpResponse::InternalServerError() - .body("Failed to insert flow key into database"); - } - } - Err(e) => { - error!("Failed to insert key into database: {}", e); - return HttpResponse::InternalServerError() - .body("Failed to insert key into database"); - } - } + valid_keys.push(new_key.clone()); } + info!("Processing batch of {} keys for flow: {}", valid_keys.len(), flow_id_str); + + // Батчевая вставка ключей с получением статистики + let key_stats = match crate::db::batch_insert_keys(&db_client, &valid_keys).await { + Ok(stats) => stats, + Err(e) => { + error!("Failed to batch insert keys into database: {}", e); + return HttpResponse::InternalServerError() + .body("Failed to batch insert keys into database"); + } + }; + + // Если нет новых ключей, нет необходимости обновлять связи с flow + if key_stats.inserted > 0 { + // Извлекаем только ID ключей из статистики + let key_ids: Vec = key_stats.key_id_map + .iter() + .map(|(_, id)| *id) + .collect(); + + // Батчевая вставка связей ключей с flow + if let Err(e) = crate::db::batch_insert_flow_keys(&db_client, &flow_id_str, &key_ids).await { + error!("Failed to batch insert flow keys into database: {}", e); + return HttpResponse::InternalServerError() + .body("Failed to batch insert flow keys into database"); + } + + info!("Added flow associations for {} keys in flow '{}'", key_ids.len(), flow_id_str); + } else { + info!("No new keys to associate with flow '{}'", flow_id_str); + } + + // Получаем обновленные данные let updated_flows = match get_keys_from_db(&db_client).await { Ok(flows) => flows, Err(e) => { @@ -188,8 +211,16 @@ pub async fn add_keys( let updated_flow = flows_guard.iter().find(|flow| flow.name == flow_id_str); if let Some(flow) = updated_flow { let servers: Vec<&SshKey> = flow.servers.iter().collect(); - info!("Updated flow: {} with {} keys", flow_id_str, servers.len()); - HttpResponse::Ok().json(servers) + info!("Keys summary for flow '{}': total received={}, new={}, unchanged={}, total in flow={}", + flow_id_str, key_stats.total, key_stats.inserted, key_stats.unchanged, servers.len()); + + // Добавляем статистику в HTTP заголовки ответа + let mut response = HttpResponse::Ok(); + response.append_header(("X-Keys-Total", key_stats.total.to_string())); + response.append_header(("X-Keys-New", key_stats.inserted.to_string())); + response.append_header(("X-Keys-Unchanged", key_stats.unchanged.to_string())); + + response.json(servers) } else { error!("Flow ID not found after update: {}", flow_id_str); HttpResponse::NotFound().body("Flow ID not found") @@ -207,7 +238,17 @@ pub async fn run_server(args: crate::Args) -> std::io::Result<()> { args.db_host, db_user, db_password, args.db_name ); - let (db_client, connection) = tokio_postgres::connect(&db_conn_str, NoTls).await.unwrap(); + info!("Connecting to database at {}", args.db_host); + let (db_client, connection) = match tokio_postgres::connect(&db_conn_str, NoTls).await { + Ok((client, conn)) => (client, conn), + Err(e) => { + error!("Failed to connect to the database: {}", e); + return Err(std::io::Error::new( + std::io::ErrorKind::ConnectionRefused, + format!("Database connection error: {}", e), + )); + } + }; let db_client = Arc::new(db_client); // Spawn a new thread to run the database connection @@ -217,6 +258,15 @@ pub async fn run_server(args: crate::Args) -> std::io::Result<()> { } }); + // Initialize database schema if needed + if let Err(e) = db::initialize_db_schema(&db_client).await { + error!("Failed to initialize database schema: {}", e); + return Err(std::io::Error::new( + std::io::ErrorKind::Other, + format!("Database schema initialization error: {}", e), + )); + } + let mut initial_flows = match get_keys_from_db(&db_client).await { Ok(flows) => flows, Err(e) => {