lox/crates/lox-distributor/src/db_handler.rs

use std::sync::{Arc, Mutex};

use crate::{lox_context, DbConfig};
use chrono::{naive::Days, DateTime, Local, NaiveDateTime, Utc};
use lox_library::{BridgeAuth, BridgeDb};
use sled::IVec;

pub struct DB {
    db: sled::Db,
}

impl DB {
    pub fn write_context(&mut self, context: lox_context::LoxServerContext) {
        let date = Local::now().format("context_%Y-%m-%d_%H:%M:%S").to_string();
        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)
        );
    }

    pub fn open_new_or_existing_db(
        db_config: DbConfig,
        roll_back_date: Option<String>,
    ) -> Result<(DB, lox_context::LoxServerContext), sled::Error> {
        let 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() {
                    context = read_lox_context_from_db(lox_db.clone(), roll_back_date);
                //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())),
                    };
                }
                (DB { db: lox_db }, context)
            }
            Err(e) => {
                panic!("Unable to read or create lox database! {:?}", e);
            }
        };
        Ok((lox_db, context))
    }
}

fn read_lox_context_from_db(
    lox_db: sled::Db,
    roll_back_date: Option<String>,
) -> lox_context::LoxServerContext {
    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);
        }
    }
    context
}

fn compute_startdate_string(date_range_end: String) -> Option<DateTime<Utc>> {
    let parsed_end =
        NaiveDateTime::parse_from_str(&date_range_end, "context_%Y-%m-%d_%H:%M:%S").unwrap();
    let dt = DateTime::<Utc>::from_naive_utc_and_offset(parsed_end, Utc);
    dt.with_timezone(&Utc).checked_sub_days(Days::new(1))
}

fn use_last_context(lox_db: sled::Db) -> lox_context::LoxServerContext {
    let ivec_context = lox_db.last().unwrap().unwrap();
    let ivec_date: String = String::from_utf8(ivec_context.0.to_vec()).unwrap();
    println!("Using last context with date: {:?}", ivec_date);
    serde_json::from_slice(&ivec_context.1).unwrap()
}

#[cfg(test)]
mod tests {
    use super::lox_context::LoxServerContext;
    use super::DB;
    use super::DbConfig;

    #[test]
    fn test_write_context() {
        let (mut lox_db, _context) = DB::open_new_or_existing_db(DbConfig::default(), None).unwrap();
        assert!(
            lox_db.db.is_empty(),
            "db read from context that shouldn't exist"
        );
        let test_string = "{\"db\":{\"keypair\":[201,84,79,188,183,80,124,38,24,57,74,101,168,234,97,113,45,2,27,227,236,170,128,86,236,16,167,242,10,166,252,90,4,176,236,14,247,73,94,98,246,65,69,32,32,82,12,127,18,85,244,194,243,248,95,15,35,105,203,156,11,111,19,29],\"pubkey\":[4,176,236,14,247,73,94,98,246,65,69,32,32,82,12,127,18,85,244,194,243,248,95,15,35,105,203,156,11,111,19,29],\"openinv_buckets\":[4,5,9,7,8],\"distributed_buckets\":[3],\"current_k\":1},\"ba\":{\"lox_priv\":{\"x0tilde\":[85,41,72,109,86,46,81,9,210,97,10,49,74,248,36,92,65,58,213,176,216,158,218,104,12,74,2,195,125,21,181,2],\"x\":[[123,196,110,59,160,128,12,49,184,235,21,191,5,211,238,241,140,17,45,124,43,0,132,200,245,60,39,231,171,103,1,1],[128,1,189,179,132,58,53,239,39,140,129,8,248,214,198,30,108,62,192,126,180,213,186,2,122,251,125,209,115,123,62,15],[23,138,233,141,142,126,154,226,142,12,251,164,251,235,56,66,235,183,113,251,33,16,26,201,119,40,206,152,127,72,25,11],[61,94,54,105,111,7,90,98,138,16,180,84,221,166,224,139,241,221,247,49,103,91,79,78,3,241,79,175,62,181,179,7],[120,251,196,223,77,58,205,154,26,78,186,178,178,26,70,203,123,55,136,128,106,95,213,141,119,124,16,40,24,66,193,10],[107,64,82,11,205,92,33,21,141,85,186,19,166,69,198,15,69,18,244,229,212,214,243,124,239,128,153,97,174,221,195,1],[89,185,222,176,72,23,67,104,75,194,113,15,48,236,71,181,122,80,144,167,218,253,41,59,59,248,243,4,191,74,203,9]]},\"lox_pub\":{\"X\":[[124,86,226,73,90,65,51,142,228,106,134,73,243,140,30,66,95,245,111,108,13,103,241,221,222,231,4,202,168,86,75,120],[38,178,246,142,27,106,61,213,8,212,201,114,158,48,91,31,213,139,0,152,157,165,45,220,75,77,34,177,15,158,84,84],[152,36,2,34,71,93,87,173,89,7,250,213,65,242,146,194,137,188,157,123,228,102,85,130,78,34,117,55,113,100,200,101],[182,145,71,159,110,42,203,61,239,214,178,8,120,9,203,90,222,118,207,135,239,125,55,173,99,157,41,246,87,178,198,29],[134,145,250,125,207,137,246,47,213,61,137,28,254,57,36,198,112,81,117,8,126,177,35,72,5,2,175,116,117,7,246,0],[210,51,80,8,26,219,12,108,164,137,122,142,190,146,151,136,8,129,78,2,170,79,174,177,249,235,235,203,94,118,155,53],[12,214,157,115,201,64,215,118,158,108,77,219,253,222,33,248,17,233,225,150,231,193,160,51,21,1,21,109,157,36,28,54]]},\"migration_priv\":{\"x0tilde\":[181,29,215,96,118,231,180,229,98,80,101,106,72,144,19,165,214,253,0,50,245,213,91,27,185,48,159,234,38,205,147,14],\"x\":[[77,128,28,95,123,122,78,104,230,37,172,252,211,32,27,30,215,3,168,49,219,142,72,86,20,108,175,30,155,183,89,15],[131,213,100,130,177,195,196,157,155,94,160,42,146,56,244,21,21,225,217,124,223,129,208,238,72,179,222,1,22,197,80,10],[66,103,172,113,175,41,44,227,83,37,146,206,86,235,5,140,106,166,38,211,84,53,106,135,147,230,251,52,21,251,65,13],[108,44,131,107,94,121,95,238,253,225,133,200,239,190,181,29,95,141,65,201,188,140,162,41,20,49,240,249,160,188,86,8],[42,223,124,126,107,165,159,133,60,253,47,144,95,206,172,211,23,79,188,46,59,236,39,82,217,226,210,0,111,250,79,7]]},\"migration_pub\":{\"X\":[[118,3,198,185,75,234,119,161,104,244,73,220,137,14,215,106,71,172,9,105,43,196,212,8,175,36,231,13,244,150,99,5],[246,175,188,123,108,192,70,27,203,226,226,125,186,253,150,191,86,185,16,217,227,134,197,187,32,181,56,49,59,82,223,85],[106,236,9,177,125,238,13,89,118,20,190,146,180,51,246,198,153,111,55,73,233,53,16,171,145,122,200,97,173,95,49,0],[112,81,204,172,90,70,92,192,153,105,192,40,224,188,26,58,153,9,246,255,76,80,28,245,127,66,123,84,20,150,82,92],[4,67,164,55,135,180,11,122,35,219,138,108,243,185,19,6,147,165,191,82,17,143,254,177,218,101,126,58,80,116,16,54]]},\"migrationkey_priv\":{\"x0tilde\":[241,119,133,223,169,255,180,151,120,139,119,244,42,225,75,147,157,144,176,193,154,66,205,167,190,151,245,25,111,100,175,4],\"x\":[[13,164,109,8,131,170,158,121,184,95,70,236,244,54,248,30,190,152,40,147,143,174,215,76,62,125,59,140,56,46,191,11],[136,13,238,110,70,92,169,61,90,145,38,115,60,100,174,89,25,241,40,155,39,27,25,191,255,101,128,202,54,173,112,3],[138,178,205,1,214,251,19,220,100,169,146,87,169,247,92,135,238,214,176,115,68,248,98,151,249,223,250,127,
        let test_context: LoxServerContext = serde_json::from_str(&test_string).unwrap();
        lox_db.write_context(test_context);
        assert!(
            lox_db.db.len() == 1,
            "db should have written only one context"
        );
    }
}