use std::sync::{Arc, Mutex}; use crate::metrics::Metrics; use crate::{lox_context, DbConfig}; use chrono::{naive::Days, DateTime, Local, NaiveDateTime, Utc}; use lox_library::{BridgeAuth, BridgeDb}; use sled::IVec; use std::collections::HashMap; use thiserror::Error; use troll_patrol::bridge_verification_info::BridgeVerificationInfo; #[derive(Error, Debug)] pub enum LoxDBError { //failed to get last db entry #[error("Failed to read last database entry")] ReadFailure(#[from] sled::Error), //no last db entries #[error("No database entries stored")] DatabaseEmpty, } // Database of Lox Distributor State pub struct DB { db: sled::Db, } impl DB { // Writes the Lox context to the lox database with "context_%Y-%m-%d_%H:%M:%S" as the // database key pub fn write_context(&mut self, context: lox_context::LoxServerContext) { let date = Local::now().format("context_%Y-%m-%d_%H:%M:%S").to_string(); /* Uncomment to generate test file for this function after making changes to lox library let file = OpenOptions::new() .create(true) .write(true) .truncate(true) .open("db_test_file.json").unwrap(); serde_json::to_writer(&file, &context).unwrap(); */ let json_result = serde_json::to_vec(&context).unwrap(); println!("Writing context to the db with key: {:?}", date); let _new_ivec = self.db.insert( IVec::from(date.as_bytes().to_vec()), IVec::from(json_result.clone()), ); assert_eq!( self.db .get(IVec::from(date.as_bytes().to_vec())) .unwrap() .unwrap(), IVec::from(json_result) ); } // If roll_back_date is empty, opens the most recent entry in the lox database or if none exists, creates a // new database. If roll_back_date is not empty, use the specified date to roll back to a previous lox-context // either exactly the entry at the roll_back_date or within 24 hours from the roll_back_date. pub fn open_new_or_existing_db( db_config: DbConfig, roll_back_date: Option, metrics: Metrics, ) -> Result<(DB, lox_context::LoxServerContext), sled::Error> { let mut context: lox_context::LoxServerContext; let (lox_db, context) = match sled::open(db_config.db_path) { Ok(lox_db) => { // Check if the lox_db already exists if lox_db.was_recovered() && !lox_db.is_empty() { context = match read_lox_context_from_db(lox_db.clone(), roll_back_date) { Ok(ctx) => ctx, Err(e) => panic!("Unable to read lox database {:?}", e), }; context.metrics = metrics; //Otherwise, create a new Lox context } else { let new_db = BridgeDb::new(); let new_ba = BridgeAuth::new(new_db.pubkey); context = lox_context::LoxServerContext { db: Arc::new(Mutex::new(new_db)), ba: Arc::new(Mutex::new(new_ba)), extra_bridges: Arc::new(Mutex::new(Vec::new())), to_be_replaced_bridges: Arc::new(Mutex::new(Vec::new())), tp_bridge_infos: Arc::new(Mutex::new(HashMap::< String, BridgeVerificationInfo, >::new())), metrics, }; } (DB { db: lox_db }, context) } Err(e) => { panic!("Unable to read or create lox database! {:?}", e); } }; Ok((lox_db, context)) } } // Logic for finding the correct context to open from the database fn read_lox_context_from_db( lox_db: sled::Db, roll_back_date: Option, ) -> Result { let context: lox_context::LoxServerContext; // Check if there is a roll back date and try to choose the appropriate context // to rollback to, otherwise, take the last saved context match roll_back_date { // If roll back date has been specified, either the exact date or range should be set Some(roll_back_date) => { // If the date is specified and it's in the database, use that to populate the context if lox_db.contains_key(roll_back_date.clone()).unwrap() { // Find date/time in db and use the context from that date. let ivec_context = lox_db .get(IVec::from(roll_back_date.as_bytes().to_vec())) .unwrap() .unwrap(); context = serde_json::from_slice(&ivec_context).unwrap(); println!("Successfully used exact key {:?}", roll_back_date); } else { // If the exact date is not found, compute the start of the range as 24 hours prior to the specified date match compute_startdate_string(roll_back_date.clone()){ Some(start_date) => { let start_date = start_date.format("context_%Y-%m-%d_%H:%M:%S").to_string(); let r: sled::Iter = lox_db.range(IVec::from(start_date.as_bytes().to_vec())..IVec::from(roll_back_date.as_bytes().to_vec()), ); match r.last(){ Some(entry) => { let ivec_context = entry.unwrap(); let key: String = String::from_utf8(ivec_context.0.to_vec()).unwrap(); println!("Successfully used range and key {:?}", key); context = serde_json::from_slice(&ivec_context.1).unwrap(); }, //If no entries are found within the 24 hours prior to the specified date, Exit None => panic!("UNEXPECTED DATE: No entries found in the 24 hours prior to the input roll_back_date"), } } None => panic!("UNEXPECTED DATE: Tried to use input date to calculate 24 hour range but failed. Exiting."), } } } // Use the last entry to populate the Lox context if no rollback date is set (which should be most common) None => context = use_last_context(lox_db)?, } Ok(context) } fn compute_startdate_string(date_range_end: String) -> Option> { let parsed_end = NaiveDateTime::parse_from_str(&date_range_end, "context_%Y-%m-%d_%H:%M:%S").unwrap(); let dt = DateTime::::from_naive_utc_and_offset(parsed_end, Utc); dt.with_timezone(&Utc).checked_sub_days(Days::new(1)) } // Use the last context that was entered into the database fn use_last_context(lox_db: sled::Db) -> Result { match lox_db.last()? { Some(ivec_context) => { let ivec_date: String = String::from_utf8(ivec_context.0.to_vec()).unwrap(); println!("Using last context with date: {:?}", ivec_date); Ok(serde_json::from_slice(&ivec_context.1).unwrap()) } None => Err(LoxDBError::DatabaseEmpty), } } #[cfg(test)] mod tests { use std::env; use std::fs; use super::lox_context::LoxServerContext; use super::DbConfig; use super::Metrics; use super::DB; #[test] fn test_write_context() { // TODO: Fix db_test_file.json env::set_var("TEST_FILE_PATH", "db_test_file.json"); let (mut lox_db, _context) = DB::open_new_or_existing_db(DbConfig::default(), None, Metrics::default()).unwrap(); assert!( lox_db.db.is_empty(), "db read from context that shouldn't exist" ); let path = env::var("TEST_FILE_PATH").unwrap(); let contents = fs::File::open(&path).unwrap(); //let test_string = std::str::from_utf8(&contents).unwrap(); let test_context: LoxServerContext = serde_json::from_reader(contents).unwrap(); lox_db.write_context(test_context); assert!( lox_db.db.len() == 1, "db should have written only one context" ); } }