vecna
/
lox
mirror of https://gitlab.torproject.org/vecna/lox.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Make roll back date an arg, reduce fn complexity

This commit is contained in:
onyinyang 2023-08-21 14:17:04 -04:00
parent e1057143fa
commit b3ddc03ba6
No known key found for this signature in database
GPG Key ID: 156A6435430C2036
3 changed files with 90 additions and 130 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
crates/lox-distributor/config.json
View File

@ -1,10 +1,7 @@
{
"db": {
"db_path": "lox_db",
"roll_back_date": {
"date_range_start": "2023-08-04_16:35:32",
"date_range_end": "2023-08-04_15:35:32"
}
"db_path": "lox_db"
},
"rtype": {
"endpoint": "http://127.0.0.1:7100/resource-stream",

166
crates/lox-distributor/src/db_handler.rs
View File

@ -1,111 +1,38 @@
use std::sync::{Arc, Mutex};
use crate::{lox_context, DbConfig};
use chrono::prelude::*;
use chrono::{naive::Days, DateTime, Local, Utc};
use lox_library::{BridgeAuth, BridgeDb};
use sled::IVec;
pub fn write_context_to_db(db: sled::Db, context: lox_context::LoxServerContext) {
let date = Local::now().format("%Y-%m-%d_%H:%M:%S").to_string();
let json_date = serde_json::to_vec(&date).unwrap();
let date = Local::now().format("context_%Y-%m-%d_%H:%M:%S").to_string();
let json_result = serde_json::to_vec(&context).unwrap();
println!("Date: {:?}", date);
let new_ivec = db.insert(
IVec::from(json_date.clone()),
IVec::from(date.as_bytes().to_vec()),
IVec::from(json_result.clone()),
);
assert_eq!(
db.get(IVec::from(json_date.clone())).unwrap().unwrap(),
IVec::from(json_result.clone())
db.get(IVec::from(date.as_bytes().to_vec()))
.unwrap()
.unwrap(),
IVec::from(json_result)
);
println!("What is the new key? {:?}", new_ivec);
println!("New entry key: {:?}", new_ivec);
}
pub fn find_existing_db(
pub fn open_new_or_existing_db(
db_config: DbConfig,
roll_back_date: Option<String>,
) -> Result<(sled::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() {
// 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 db_config.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 exact date is specified and it's in the database, use that to populate the context
match roll_back_date.exact_date {
Some(exact_date) => {
match lox_db.contains_key(exact_date.clone()) {
// Find exact date/time in db and use the context from that date.
Ok(_) => {
let json_exact_date =
serde_json::to_vec(&exact_date).unwrap();
let ivec_context = lox_db
.get(IVec::from(json_exact_date))
.unwrap()
.unwrap();
context = serde_json::from_slice(&ivec_context).unwrap();
println!(
"Successfully used exact key {:?}",
exact_date
);
}
// If the entered date is not found, use the last saved context
Err(_) => {
println!("UNEXPECTED DATE: Specified exact date not found in lox db! Using last saved context");
context = use_last_context(lox_db.clone());
}
}
}
// Otherwise check that the start of a range has been specified
None => {
match roll_back_date.date_range_start {
Some(start) => {
// If so, ensure the end of the range has also been specified
if error_check_input_string(
start.clone(),
roll_back_date.date_range_end.clone(),
) {
let start_json: Vec<u8> =
serde_json::to_vec(&start).unwrap();
let end_json: Vec<u8> =
serde_json::to_vec(&roll_back_date.date_range_end)
.unwrap();
let r: sled::Iter = lox_db.range(
IVec::from(start_json)..IVec::from(end_json),
);
let ivec_context = r.last().unwrap().unwrap();
context =
serde_json::from_slice(&ivec_context.1).unwrap();
let json_key: String =
serde_json::from_slice(&ivec_context.0).unwrap();
println!(
"Successfully used range and key {:?}",
json_key
);
} else {
println!("UNEXPECTED DATE: Start of range was specified but End of range was not! Using last saved context");
context = use_last_context(lox_db.clone());
}
}
// If not, roll_back_date should not have been indicated but it was! So just use the last
// context and print that this is happening.
None => {
// If exact date and rollback range are blank, use the last saved context
println!("UNEXPECTED DATE: No date specified despite indicating rollback! Using last saved context");
context = use_last_context(lox_db.clone());
}
}
}
}
}
// 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.clone());
}
}
context = read_lox_context_from_db(lox_db.clone(), roll_back_date);
//Otherwise, create a new Lox context
} else {
let new_db = BridgeDb::new();
@ -115,7 +42,7 @@ pub fn find_existing_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())),
}
};
}
(lox_db, context)
}
@ -126,22 +53,65 @@ pub fn find_existing_db(
Ok((lox_db, context))
}
fn error_check_input_string(start: String, date_range_end: Option<String>) -> bool {
if let Some(end) = date_range_end {
let parsed_start = Utc.datetime_from_str(&start, "%Y-%m-%d_%H:%M:%S").unwrap();
let parsed_end = Utc.datetime_from_str(&end, "%Y-%m-%d_%H:%M:%S").unwrap();
if parsed_start <= parsed_end {
return true;
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 json_exact_date = serde_json::to_vec(&roll_back_date).unwrap();
let ivec_context = lox_db.get(IVec::from(json_exact_date)).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 sd_string = format!("{}", start_date.format("%Y-%m-%d_%H:%M:%S"));
let start_json: Vec<u8> = serde_json::to_vec(&sd_string).unwrap();
let end_json: Vec<u8> = serde_json::to_vec(&roll_back_date).unwrap();
let r: sled::Iter = lox_db.range(IVec::from(start_json)..IVec::from(end_json), );
match r.last(){
Some(entry) => {
let ivec_context = entry.unwrap();
let json_key: String = serde_json::from_slice(&ivec_context.0).unwrap();
println!("Successfully used range and key {:?}", json_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);
}
println!("date_range_start must be <= date_range_end! Check config file");
return false;
}
println!("date_range_start specified but no date_range_end! Check config file");
false
context
}
fn use_last_context(cloned_db: sled::Db) -> lox_context::LoxServerContext {
let ivec_context = cloned_db.last().unwrap().unwrap();
fn compute_startdate_string(date_range_end: String) -> Option<DateTime<Utc>> {
let parsed_end =
DateTime::parse_from_str(&date_range_end, "context_%Y-%m-%d_%H:%M:%S").unwrap();
parsed_end
.with_timezone(&Utc)
.checked_sub_days(Days::new(1))
}
fn use_last_context(db: sled::Db) -> lox_context::LoxServerContext {
let ivec_context = db.last().unwrap().unwrap();
let ivec_date: String = serde_json::from_slice(&ivec_context.0).unwrap();
println!("Using last context with date: {:?}", ivec_date);
serde_json::from_slice(&ivec_context.1).unwrap()

47
crates/lox-distributor/src/main.rs
View File

@ -17,7 +17,7 @@ use std::{
};
mod db_handler;
use db_handler::{find_existing_db, write_context_to_db};
use db_handler::{open_new_or_existing_db, write_context_to_db};
mod lox_context;
mod request_handler;
use request_handler::handle;
@ -44,16 +44,20 @@ struct Args {
#[arg(short, long, default_value = "config.json")]
config: PathBuf,
/// Optional name/path of the lox context json backup file
/// Used to populate the Lox context. If none is provided, an empty
/// Lox context will be created
// Optional Date/time to roll back to as a %Y-%m-%d_%H:%M:%S string
// This argument should be passed if the lox_context should be rolled back to a
// previous state due to, for example, a mass blocking event that is likely not
// due to Lox user behaviour. If the exact roll back date/time is not known, the
// last db entry within 24 hours from the passed roll_back_date will be used or else
// the program will fail gracefully.
#[arg(short, long, verbatim_doc_comment)]
backup_context: Option<PathBuf>,
roll_back_date: Option<String>,
}
#[derive(Debug, Deserialize)]
struct Config {
db: DbConfig,
rtype: ResourceInfo,
}
@ -61,32 +65,15 @@ struct Config {
pub struct DbConfig {
// The path for the lox_context database, default is "lox_db"
db_path: String,
// Optional Date/time to roll back to
// This should be blank for normal operation and only changed if the lox_context
// should be rolled back to a previous state
roll_back_date: Option<RollBackDate>,
}
impl Default for DbConfig {
fn default() -> DbConfig {
DbConfig {
db_path: "lox_db".to_owned(),
roll_back_date: None,
}
}
}
#[derive(Debug, Deserialize)]
pub struct RollBackDate {
// Optional exact Date/time to roll back to as a %Y-%m-%d_%H:%M:%S string
// This should only be changed if the lox_context should be rolled back to a
// previous state and the exact roll back date/time is known
exact_date: Option<String>,
// Since the exact time the database last saved the context may not be obvious,
// this can be presented as a range from which the most recent date in the database will be
// selected
date_range_start: Option<String>,
date_range_end: Option<String>,
}
#[derive(Debug, Deserialize)]
struct ResourceInfo {
@ -146,11 +133,12 @@ async fn parse_bridges(rdsys_tx: mpsc::Sender<Command>, mut rx: mpsc::Receiver<R
async fn create_context_manager(
db_config: DbConfig,
roll_back_date: Option<String>,
context_rx: mpsc::Receiver<Command>,
mut kill: broadcast::Receiver<()>,
) {
tokio::select! {
create_context = context_manager(db_config, context_rx) => create_context,
create_context = context_manager(db_config, roll_back_date, context_rx) => create_context,
_ = kill.recv() => {println!("Shut down context_manager");},
}
}
@ -158,8 +146,12 @@ async fn create_context_manager(
// Context Manager handles the Lox BridgeDB and Bridge Authority, ensuring
// that the DB can be updated from the rdsys stream and client requests
// can be responded to with an updated BridgeDB state
async fn context_manager(db_config: DbConfig, mut context_rx: mpsc::Receiver<Command>) {
let (lox_db, context) = match find_existing_db(db_config) {
async fn context_manager(
db_config: DbConfig,
roll_back_date: Option<String>,
mut context_rx: mpsc::Receiver<Command>,
) {
let (lox_db, context) = match open_new_or_existing_db(db_config, roll_back_date) {
Ok((lox_db, context)) => (lox_db, context),
Err(e) => {
panic!("Error: {:?}", e);
@ -379,8 +371,9 @@ async fn main() {
}
});
let context_manager =
spawn(async move { create_context_manager(config.db, context_rx, kill_context).await });
let context_manager = spawn(async move {
create_context_manager(config.db, args.roll_back_date, context_rx, kill_context).await
});
let (tx, rx) = mpsc::channel(32);
let rdsys_stream_handler = spawn(async { rdsys_stream(config.rtype, tx, kill_stream).await });