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

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

use lox::{
    bridge_table::{BridgeLine, ENC_BUCKET_BYTES, MAX_BRIDGES_PER_BUCKET},
    proto::{
        blockage_migration, check_blockage, issue_invite, level_up, migration, open_invite,
        redeem_invite, trust_promotion,
    },
    BridgeAuth, BridgeDb, IssuerPubKey,
};
use lox_utils;
use rand::RngCore;
use std::sync::{Arc, Mutex};

/// Create a random BridgeLine for testing ONLY. Do not use in production!
/// This was copied directly from lox/src/bridge_table.rs in order
/// to easily initialize a bridgedb/lox_auth with structurally
/// correct buckets to be used for Lox requests/verifications/responses.
/// In production, existing bridges should be translated into this format
/// in a private function and sorted into buckets (3 bridges/bucket is suggested
/// but experience may suggest something else) in some intelligent way.

pub fn random() -> BridgeLine {
    let mut rng = rand::thread_rng();
    let mut res: BridgeLine = BridgeLine::default();
    // Pick a random 4-byte address
    let mut addr: [u8; 4] = [0; 4];
    rng.fill_bytes(&mut addr);
    // If the leading byte is 224 or more, that's not a valid IPv4
    // address.  Choose an IPv6 address instead (but don't worry too
    // much about it being well formed).
    if addr[0] >= 224 {
        rng.fill_bytes(&mut res.addr);
    } else {
        // Store an IPv4 address as a v4-mapped IPv6 address
        res.addr[10] = 255;
        res.addr[11] = 255;
        res.addr[12..16].copy_from_slice(&addr);
    };
    let ports: [u16; 4] = [443, 4433, 8080, 43079];
    let portidx = (rng.next_u32() % 4) as usize;
    res.port = ports[portidx];
    res.uid_fingerprint = rng.next_u64();
    let mut cert: [u8; 52] = [0; 52];
    rng.fill_bytes(&mut cert);
    let infostr: String = format!(
        "obfs4 cert={}, iat-mode=0",
        base64::encode_config(cert, base64::STANDARD_NO_PAD)
    );
    res.info[..infostr.len()].copy_from_slice(infostr.as_bytes());
    res
}

#[derive(Clone)]
pub struct LoxServerContext {
    pub db: Arc<Mutex<BridgeDb>>,
    pub ba: Arc<Mutex<BridgeAuth>>,
    pub extra_bridges: Arc<Mutex<Vec<BridgeLine>>>,
    pub unreplaced_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 i in 0..MAX_BRIDGES_PER_BUCKET {
            return_bridges[i] = extra_bridges.remove(i);
        }

        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_unreplaced_bridge(&self, bridge: BridgeLine) {
        let mut unreplaced_bridges = self.unreplaced_bridges.lock().unwrap();
        unreplaced_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();
        ba_obj.add_openinv_bridges(bucket, &mut db_obj);
    }

    pub fn add_spare_bucket(&self, bucket: [BridgeLine; 3]) {
        let mut ba_obj = self.ba.lock().unwrap();
        ba_obj.add_spare_bucket(bucket);
    }

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

    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) -> Vec<[u8; ENC_BUCKET_BYTES]> {
        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 obj = self.db.lock().unwrap();
        lox_utils::Invite {
            invite: obj.invite(),
        }
    }

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

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

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

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

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

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

    fn check_blockage(&self, req: check_blockage::Request) -> check_blockage::Response {
        let mut ba_obj = self.ba.lock().unwrap();
        // Created 5 buckets initially, so we will add 5  hot spares (for migration) and
        // block all of the existing buckets to trigger migration table propagation
        // FOR TESTING ONLY, ADD 5 NEW Buckets
        for _ in 0..5 {
            let bucket = [random(), random(), random()];
            ba_obj.add_spare_bucket(bucket);
        }
        ba_obj.enc_bridge_table();

        // FOR TESTING ONLY, BLOCK ALL BRIDGES
        let mut db_obj = self.db.lock().unwrap();
        for index in 0..5 {
            let b0 = ba_obj.bridge_table.buckets[index][0];
            let b1 = ba_obj.bridge_table.buckets[index][1];
            let b2 = ba_obj.bridge_table.buckets[index][2];
            ba_obj.bridge_unreachable(&b0, &mut db_obj);
            ba_obj.bridge_unreachable(&b1, &mut db_obj);
            ba_obj.bridge_unreachable(&b2, &mut db_obj);
        }
        ba_obj.enc_bridge_table();
        ba_obj.handle_check_blockage(req).unwrap()
    }

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

// Generate and return an open invitation token
pub fn generate_invite(context: LoxServerContext) -> Response<Body> {
    let invite = context.gen_invite();
    let token = serde_json::to_string(&invite).unwrap();
    prepare_header(token)
}

// Return the serialized encrypted bridge table
pub fn send_reachability_cred(context: LoxServerContext) -> Response<Body> {
    let enc_table = context.encrypt_table();
    let etable = lox_utils::EncBridgeTable { etable: enc_table };
    prepare_header(serde_json::to_string(&etable).unwrap())
}

// Return the serialized pubkeys for the Bridge Authority
pub fn send_keys(context: LoxServerContext) -> Response<Body> {
    let pubkeys = context.pubkeys();
    prepare_header(serde_json::to_string(&pubkeys).unwrap())
}

pub fn verify_and_send_open_cred(request: Bytes, context: LoxServerContext) -> Response<Body> {
    let req: open_invite::Request = serde_json::from_slice(&request).unwrap();
    let response = context.open_inv(req);
    let open_invite_resp_str = serde_json::to_string(&response).unwrap();
    prepare_header(open_invite_resp_str)
}

pub fn verify_and_send_trust_promo(request: Bytes, context: LoxServerContext) -> Response<Body> {
    let req: trust_promotion::Request = serde_json::from_slice(&request).unwrap();
    let response = context.trust_promo(req);
    let trust_promo_resp_str = serde_json::to_string(&response).unwrap();
    prepare_header(trust_promo_resp_str)
}

pub fn verify_and_send_trust_migration(
    request: Bytes,
    context: LoxServerContext,
) -> Response<Body> {
    let req: migration::Request = serde_json::from_slice(&request).unwrap();
    let response = context.trust_migration(req);
    let resp_str = serde_json::to_string(&response).unwrap();
    prepare_header(resp_str)
}

pub fn verify_and_send_level_up(request: Bytes, context: LoxServerContext) -> Response<Body> {
    let req: level_up::Request = serde_json::from_slice(&request).unwrap();
    let response = context.level_up(req);
    let level_up_resp_str = serde_json::to_string(&response).unwrap();
    prepare_header(level_up_resp_str)
}

pub fn verify_and_send_issue_invite(request: Bytes, context: LoxServerContext) -> Response<Body> {
    let req: issue_invite::Request = serde_json::from_slice(&request).unwrap();
    let response = context.issue_invite(req);
    let issue_invite_resp_str = serde_json::to_string(&response).unwrap();
    prepare_header(issue_invite_resp_str)
}

pub fn verify_and_send_redeem_invite(request: Bytes, context: LoxServerContext) -> Response<Body> {
    let req: redeem_invite::Request = serde_json::from_slice(&request).unwrap();
    let response = context.redeem_invite(req);
    let redeem_invite_resp_str = serde_json::to_string(&response).unwrap();
    prepare_header(redeem_invite_resp_str)
}

pub fn verify_and_send_check_blockage(request: Bytes, context: LoxServerContext) -> Response<Body> {
    let req: check_blockage::Request = serde_json::from_slice(&request).unwrap();

    let response = context.check_blockage(req);
    let check_blockage_resp_str = serde_json::to_string(&response).unwrap();
    prepare_header(check_blockage_resp_str)
}

pub fn verify_and_send_blockage_migration(
    request: Bytes,
    context: LoxServerContext,
) -> Response<Body> {
    let req: blockage_migration::Request = serde_json::from_slice(&request).unwrap();
    let response = context.blockage_migration(req);
    let resp_str = serde_json::to_string(&response).unwrap();
    prepare_header(resp_str)
}

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
}