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

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;

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>,
        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() {
                    context = read_lox_context_from_db(lox_db.clone(), roll_back_date);
                    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())),
                        metrics,
                    };
                }
                (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::DbConfig;
    use super::Metrics;
    use super::DB;

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