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

use hyper::{body::Bytes, header::HeaderValue, Body, Response};

use lox_library::{
    bridge_table::{BridgeLine, EncryptedBucket, MAX_BRIDGES_PER_BUCKET},
    proto::{
        blockage_migration, check_blockage, issue_invite, level_up, migration, open_invite,
        redeem_invite, trust_promotion,
    },
    BridgeAuth, BridgeDb, IssuerPubKey,
};
use serde::{Deserialize, Serialize};

use std::{
    collections::HashMap,
    sync::{Arc, Mutex},
};
use zkp::ProofError;

#[derive(Clone, Serialize, Deserialize)]
pub struct LoxServerContext {
    pub db: Arc<Mutex<BridgeDb>>,
    pub ba: Arc<Mutex<BridgeAuth>>,
    pub extra_bridges: Arc<Mutex<Vec<BridgeLine>>>,
    pub to_be_replaced_bridges: Arc<Mutex<Vec<BridgeLine>>>,
}

impl LoxServerContext {
    pub fn append_extra_bridges(&self, bridge: BridgeLine) {
        let mut extra_bridges = self.extra_bridges.lock().unwrap();
        extra_bridges.push(bridge);
    }

    pub fn remove_extra_bridges(&self) -> [BridgeLine; MAX_BRIDGES_PER_BUCKET] {
        let mut extra_bridges = self.extra_bridges.lock().unwrap();
        let mut return_bridges = [BridgeLine::default(); MAX_BRIDGES_PER_BUCKET];
        for bridge in return_bridges.iter_mut() {
            if let Some(extra) = extra_bridges.pop() {
                *bridge = extra
            }
        }
        return_bridges
    }

    pub fn remove_single_bridge(&self) {
        let mut extra_bridges = self.extra_bridges.lock().unwrap();
        let length = extra_bridges.len();
        _ = extra_bridges.remove(length - 1)
    }

    pub fn new_to_be_replaced_bridge(&self, bridge: BridgeLine) {
        let mut to_be_replaced_bridges = self.to_be_replaced_bridges.lock().unwrap();
        to_be_replaced_bridges.push(bridge);
    }

    pub fn allocate_leftover_bridges(&self) {
        let mut ba_obj = self.ba.lock().unwrap();
        let mut db_obj = self.db.lock().unwrap();
        let mut extra_bridges = self.extra_bridges.lock().unwrap();
        ba_obj.allocate_bridges(&mut extra_bridges, &mut db_obj);
    }

    pub fn add_openinv_bucket(&self, bucket: [BridgeLine; 3]) {
        let mut ba_obj = self.ba.lock().unwrap();
        let mut db_obj = self.db.lock().unwrap();
        match ba_obj.add_openinv_bridges(bucket, &mut db_obj) {
            Ok(_) => (),
            Err(e) => {
                println!("Error: {:?}", e);
                for bridge in bucket {
                    self.append_extra_bridges(bridge);
                }
            }
        }
    }

    pub fn add_spare_bucket(&self, bucket: [BridgeLine; 3]) {
        let mut ba_obj = self.ba.lock().unwrap();
        let mut db_obj = self.db.lock().unwrap();
        match ba_obj.add_spare_bucket(bucket, &mut db_obj) {
            Ok(_) => (),
            Err(e) => {
                println!("Error: {:?}", e);
                for bridge in bucket {
                    self.append_extra_bridges(bridge);
                }
            }
        }
    }

    pub fn replace_with_new(&self, bridgeline: BridgeLine) -> lox_library::ReplaceSuccess {
        let mut ba_obj = self.ba.lock().unwrap();
        let eb_obj = self.extra_bridges.lock().unwrap();
        let available_bridge = eb_obj.last();
        let result = ba_obj.bridge_replace(&bridgeline, available_bridge);
        // .last() doesn't actually remove the object so we still have to do that if the bridge was
        // replaced with an available bridge
        if result == lox_library::ReplaceSuccess::Replaced && eb_obj.len() > 0 {
            self.remove_single_bridge();
        }
        result
    }

    /* Uncomment when bridge blocking is finalized
        pub fn add_unreachable(&self, bridgeline: BridgeLine) -> bool {
            let mut ba_obj = self.ba.lock().unwrap();
            let mut db_obj = self.db.lock().unwrap();
            ba_obj.bridge_unreachable(&bridgeline, &mut db_obj)
        }
    */

    pub fn update_bridge(&self, bridgeline: BridgeLine) -> bool {
        let mut ba_obj = self.ba.lock().unwrap();
        ba_obj.bridge_update(&bridgeline)
    }

    #[cfg(test)]
    /// For testing only: manually advance the day by the given number
    /// of days.
    pub fn advance_days_test(&self, num: u16) {
        let mut ba_obj = self.ba.lock().unwrap();
        ba_obj.advance_days(num); // FOR TESTING ONLY
        println!("Today's date according to server: {}", ba_obj.today());
    }

    pub fn encrypt_table(&self) -> HashMap<u32, EncryptedBucket> {
        let mut ba_obj = self.ba.lock().unwrap();
        ba_obj.enc_bridge_table().clone()
    }

    fn pubkeys(&self) -> Vec<IssuerPubKey> {
        let ba_obj = self.ba.lock().unwrap();
        // vector of public keys (to serialize)
        vec![
            ba_obj.lox_pub.clone(),
            ba_obj.migration_pub.clone(),
            ba_obj.migrationkey_pub.clone(),
            ba_obj.reachability_pub.clone(),
            ba_obj.invitation_pub.clone(),
        ]
    }

    fn gen_invite(&self) -> lox_utils::Invite {
        let mut obj = self.db.lock().unwrap();
        lox_utils::Invite {
            invite: obj.invite(),
        }
    }

    fn open_inv(&self, req: open_invite::Request) -> Result<open_invite::Response, ProofError> {
        let mut ba_obj = self.ba.lock().unwrap();
        ba_obj.handle_open_invite(req)
    }

    fn trust_promo(
        &self,
        req: trust_promotion::Request,
    ) -> Result<trust_promotion::Response, ProofError> {
        let mut ba_obj = self.ba.lock().unwrap();
        ba_obj.handle_trust_promotion(req)
    }

    fn trust_migration(&self, req: migration::Request) -> Result<migration::Response, ProofError> {
        let mut ba_obj = self.ba.lock().unwrap();
        ba_obj.handle_migration(req)
    }

    fn level_up(&self, req: level_up::Request) -> Result<level_up::Response, ProofError> {
        let mut ba_obj = self.ba.lock().unwrap();
        ba_obj.handle_level_up(req)
    }

    fn issue_invite(
        &self,
        req: issue_invite::Request,
    ) -> Result<issue_invite::Response, ProofError> {
        let mut ba_obj = self.ba.lock().unwrap();
        ba_obj.handle_issue_invite(req)
    }

    fn redeem_invite(
        &self,
        req: redeem_invite::Request,
    ) -> Result<redeem_invite::Response, ProofError> {
        let mut ba_obj = self.ba.lock().unwrap();
        ba_obj.handle_redeem_invite(req)
    }

    fn check_blockage(
        &self,
        req: check_blockage::Request,
    ) -> Result<check_blockage::Response, ProofError> {
        let mut ba_obj = self.ba.lock().unwrap();
        ba_obj.handle_check_blockage(req)
    }

    fn blockage_migration(
        &self,
        req: blockage_migration::Request,
    ) -> Result<blockage_migration::Response, ProofError> {
        let mut ba_obj = self.ba.lock().unwrap();
        ba_obj.handle_blockage_migration(req)
    }

    // Generate and return an open invitation token
    pub fn generate_invite(self) -> Response<Body> {
        let invite = self.gen_invite();
        match serde_json::to_string(&invite) {
            Ok(resp) => prepare_header(resp),
            Err(e) => {
                println!("Error parsing Invite to JSON");
                prepare_error_header(e.to_string())
            }
        }
    }

    // Return the serialized encrypted bridge table
    pub fn send_reachability_cred(self) -> Response<Body> {
        let enc_table = self.encrypt_table();
        let etable = lox_utils::EncBridgeTable { etable: enc_table };
        match serde_json::to_string(&etable) {
            Ok(resp) => prepare_header(resp),
            Err(e) => {
                println!("Error parsing encrypted bridgetable to JSON");
                prepare_error_header(e.to_string())
            }
        }
    }

    // Return the serialized pubkeys for the Bridge Authority
    pub fn send_keys(self) -> Response<Body> {
        let pubkeys = self.pubkeys();
        match serde_json::to_string(&pubkeys) {
            Ok(resp) => prepare_header(resp),
            Err(e) => {
                println!("Error parsing Pubkeys to JSON");
                prepare_error_header(e.to_string())
            }
        }
    }

    pub fn verify_and_send_open_cred(self, request: Bytes) -> Response<Body> {
        let req = match serde_json::from_slice(&request) {
            Ok(req) => req,
            Err(e) => return prepare_error_header(e.to_string()),
        };
        match self.open_inv(req) {
            Ok(resp) => {
                let response = serde_json::to_string(&resp).unwrap();
                prepare_header(response)
            }
            Err(e) => {
                println!("Invalid Open Invitation request, Proof Error");
                prepare_error_header(e.to_string())
            }
        }
    }

    pub fn verify_and_send_trust_promo(self, request: Bytes) -> Response<Body> {
        let req: trust_promotion::Request = match serde_json::from_slice(&request) {
            Ok(req) => req,
            Err(e) => return prepare_error_header(e.to_string()),
        };
        match self.trust_promo(req) {
            Ok(resp) => {
                let response = serde_json::to_string(&resp).unwrap();
                prepare_header(response)
            }
            Err(e) => {
                println!("Invalid Trust Promotion request, Proof Error");
                prepare_error_header(e.to_string())
            }
        }
    }

    pub fn verify_and_send_trust_migration(self, request: Bytes) -> Response<Body> {
        let req: migration::Request = match serde_json::from_slice(&request) {
            Ok(req) => req,
            Err(e) => return prepare_error_header(e.to_string()),
        };
        match self.trust_migration(req) {
            Ok(resp) => {
                let response = serde_json::to_string(&resp).unwrap();
                prepare_header(response)
            }
            Err(e) => {
                println!("Invalid Trust Migration request, Proof Error");
                prepare_error_header(e.to_string())
            }
        }
    }

    pub fn verify_and_send_level_up(self, request: Bytes) -> Response<Body> {
        let req: level_up::Request = match serde_json::from_slice(&request) {
            Ok(req) => req,
            Err(e) => return prepare_error_header(e.to_string()),
        };
        match self.level_up(req) {
            Ok(resp) => {
                let response = serde_json::to_string(&resp).unwrap();
                prepare_header(response)
            }
            Err(e) => {
                println!("Invalid Level up request, Proof Error");
                prepare_error_header(e.to_string())
            }
        }
    }

    pub fn verify_and_send_issue_invite(self, request: Bytes) -> Response<Body> {
        let req: issue_invite::Request = match serde_json::from_slice(&request) {
            Ok(req) => req,
            Err(e) => return prepare_error_header(e.to_string()),
        };
        match self.issue_invite(req) {
            Ok(resp) => {
                let response = serde_json::to_string(&resp).unwrap();
                prepare_header(response)
            }
            Err(e) => {
                println!("Invalid Issue invite request, Proof Error");
                prepare_error_header(e.to_string())
            }
        }
    }

    pub fn verify_and_send_redeem_invite(self, request: Bytes) -> Response<Body> {
        let req: redeem_invite::Request = match serde_json::from_slice(&request) {
            Ok(req) => req,
            Err(e) => return prepare_error_header(e.to_string()),
        };
        match self.redeem_invite(req) {
            Ok(resp) => {
                let response = serde_json::to_string(&resp).unwrap();
                prepare_header(response)
            }
            Err(e) => {
                println!("Invalid Redeem Invite request, Proof Error");
                prepare_error_header(e.to_string())
            }
        }
    }

    pub fn verify_and_send_check_blockage(self, request: Bytes) -> Response<Body> {
        let req: check_blockage::Request = match serde_json::from_slice(&request) {
            Ok(req) => req,
            Err(e) => return prepare_error_header(e.to_string()),
        };
        match self.check_blockage(req) {
            Ok(resp) => {
                let response = serde_json::to_string(&resp).unwrap();
                prepare_header(response)
            }
            Err(e) => {
                println!("Invalid Check blockage request, Proof Error");
                prepare_error_header(e.to_string())
            }
        }
    }

    pub fn verify_and_send_blockage_migration(self, request: Bytes) -> Response<Body> {
        let req: blockage_migration::Request = match serde_json::from_slice(&request) {
            Ok(req) => req,
            Err(e) => return prepare_error_header(e.to_string()),
        };
        match self.blockage_migration(req) {
            Ok(resp) => {
                let response = serde_json::to_string(&resp).unwrap();
                prepare_header(response)
            }
            Err(e) => {
                println!("Invalid Blockage Migration request, Proof Error");
                prepare_error_header(e.to_string())
            }
        }
    }
}

fn prepare_header(response: String) -> Response<Body> {
    let mut resp = Response::new(Body::from(response));
    resp.headers_mut()
        .insert("Access-Control-Allow-Origin", HeaderValue::from_static("*"));
    resp
}

fn prepare_error_header(error: String) -> Response<Body> {
    Response::builder()
        .status(hyper::StatusCode::BAD_REQUEST)
        .body(Body::from(error))
        .unwrap()
}