ppc/packages/marts/auth/client.rs

use chrono::{DateTime, TimeDelta, Utc};
use mize::{
    item::ItemData, mize_err, mize_part, Mize, MizeError, MizePart, MizePartCreate, MizeResult,
};
use openidconnect::{
    ClientId, DeviceAuthorizationResponse, EmptyExtraDeviceAuthorizationFields, IssuerUrl,
    OAuth2TokenResponse, RefreshToken, Scope, TokenResponse,
};
use serde::{Deserialize, Serialize};
use std::{path::PathBuf, thread};

// console_log macro
// that can be copied into other files for debugging purposes
#[cfg(feature = "target-obsidian")]
use wasm_bindgen::prelude::*;

#[cfg(feature = "target-obsidian")]
#[wasm_bindgen]
extern "C" {
    #[wasm_bindgen(js_namespace = console)]
    fn log(s: &str);
}

#[cfg(feature = "target-obsidian")]
macro_rules! console_log {
    // Note that this is using the `log` function imported above during
    // `bare_bones`
    ($($t:tt)*) => (unsafe { log(&format_args!($($t)*).to_string())})
}

#[cfg(feature = "target-os")]
use clap::Command;

#[derive(Default)]
#[mize_part]
pub struct ClientAuth {
    mize: Mize,
    secrets: Option<ClientSecrets>,
}

#[derive(Default, Serialize, Deserialize, Clone, Debug)]
pub struct ClientSecrets {
    id_token: String,
    access_token: String,
    refresh_token: String,
    expiry: DateTime<Utc>,
}

impl MizePartCreate for ClientAuth {}

impl MizePart for ClientAuth {
    fn name(&self) -> &'static str {
        "client_auth"
    }
}

impl ClientAuth {
    pub fn id_token(&self) -> MizeResult<&str> {
        if self.secrets.is_none() {
            return Err(mize_err!("Client not logged in"));
        }
        Ok(&self.secrets.as_ref().unwrap().id_token)
    }

    pub fn access_token(&self) -> MizeResult<&str> {
        if self.secrets.is_none() {
            return Err(mize_err!("Client not logged in"));
        }
        Ok(&self.secrets.as_ref().unwrap().access_token)
    }

    pub fn refresh_token(&self) -> MizeResult<&str> {
        if self.secrets.is_none() {
            return Err(mize_err!("Client not logged in"));
        }
        Ok(&self.secrets.as_ref().unwrap().refresh_token)
    }

    pub fn refresh(&mut self) -> MizeResult<()> {
        let secrets = self.secrets.clone().ok_or(mize_err!(
            "no client secrets, meaning the refresh thread can't use the refresh token"
        ))?;
        self.mize.spawn_async_blocking(
            "auth.refresh_flow",
            refresh_flow(self.mize.clone(), secrets),
        )?;
        Ok(())
    }

    pub fn start_refresh_thread(&mut self) -> MizeResult<()> {
        let mut mize_clone = self.mize.clone();
        self.mize.spawn_background("auth.refresh", move || loop {
            let auth = mize_clone.get_part_native::<ClientAuth>("auth")?;
            let to_sleep =
                Utc::now() - auth.secrets.as_ref().unwrap().expiry.clone() - TimeDelta::minutes(20);
            drop(auth);
            println!("refresh_thread: going to sleep for: {}", to_sleep);
            thread::sleep(to_sleep.to_std()?);
            println!("refresh_thread: sleeping done");
            let mut auth = mize_clone.get_part_native::<ClientAuth>("auth")?;
            auth.refresh()?;
            drop(auth);
        })?;
        Ok(())
    }
}

pub fn client_auth(mize: &mut Mize) -> MizeResult<()> {
    let mize_clone = mize.clone();

    let mut secrets = match read_secrets(mize) {
        Ok(secrets) => Some(secrets),
        Err(err) => {
            println!("err reading secrets: {:?}", err);
            None
        }
    };

    // do a refresh if the tokens have expired
    if secrets.is_some() && secrets.clone().unwrap().expiry < Utc::now() {
        secrets = Some(
            mize.spawn_async_blocking("refresh_low", refresh_flow(mize.clone(), secrets.unwrap()))?,
        );
    }

    #[cfg(feature = "target-os")]
    {
        let mut cli = mize.get_part_native::<crate::CliPart>("cli")?;
        cli.subcommand(Command::new("login"), |_, mut mize| {
            let mut auth = mize.get_part_native::<ClientAuth>("client_auth")?;
            println!("secrets: {:?}", auth.secrets);
            let secrets = if auth.secrets.is_some() {
                match mize.spawn_async_blocking(
                    "client_auth",
                    refresh_flow(mize.clone(), auth.secrets.clone().unwrap()),
                ) {
                    Ok(secrets) => secrets,
                    Err(err) => {
                        println!("err refreshing secrets: {:?}", err);
                        mize.spawn_async_blocking("client_auth", auth_flow(mize_clone))?
                    }
                }
            } else {
                mize.spawn_async_blocking("client_auth", auth_flow(mize_clone))?
            };
            write_secrets(&mut mize, &secrets)?;
            auth.secrets = Some(secrets);
            Ok(())
        });
    }

    #[cfg(feature = "target-obsidian")]
    {}

    // check if login token is stil valid refresh if not
    /*
    if secrets.is_some() {
        secrets = Some(mize.spawn_async_blocking(
            "client_auth",
            refresh_flow(mize.clone(), secrets.clone().unwrap()),
        )?);
    }
    */

    mize.register_part(Box::new(ClientAuth {
        secrets,
        mize: mize.clone(),
    }))
}

fn read_secrets(mize: &mut Mize) -> MizeResult<ClientSecrets> {
    let mize_dir = mize.get_config("data_dir")?.value_string()?;
    let secrets_file_path = PathBuf::from(mize_dir).join("client_secrets.json");

    let secrets_string = std::fs::read_to_string(secrets_file_path)?;
    let secrets: ClientSecrets = serde_json::from_str(&secrets_string)?;
    Ok(secrets)
}

fn write_secrets(mize: &mut Mize, secrets: &ClientSecrets) -> MizeResult<()> {
    let mize_dir = mize.get_config("data_dir")?.value_string()?;
    let secrets_file_path = PathBuf::from(mize_dir).join("client_secrets.json");

    let secrets_string = serde_json::to_string(secrets)?;
    std::fs::write(secrets_file_path, secrets_string)?;

    // update the spacetimedb_token in XDG_CONFIG_HOME/spacetime/cli.toml config file if configured to do so
    if mize
        .get_config("auth.client.update_spacetime_token")
        .unwrap_or(ItemData::empty())
        .as_bool()
        .unwrap_or(false)
    {
        update_spacetimedb_token(&secrets.id_token.as_str())?;
    }

    Ok(())
}

use openidconnect::reqwest::Client as HttpClient;
use openidconnect::DeviceAuthorizationUrl;
use std::time::Duration;

async fn refresh_flow(mut mize: Mize, old_secrets: ClientSecrets) -> MizeResult<ClientSecrets> {
    println!("Refresh flow...");
    let issuer = mize.get_config("auth.client.issuer")?.value_string()?;
    let client_id = mize.get_config("auth.client.client_id")?.value_string()?;

    let issuer = IssuerUrl::new(issuer)?;
    let client_id = ClientId::new(client_id);

    let http_client = HttpClient::new();

    // Teach openidconnect-rs about the device authorization url.
    #[derive(Clone, Debug, Deserialize, Serialize)]
    struct MyExtraProviderMetadata {
        device_authorization_endpoint: String,
    }
    use openidconnect::core::{
        CoreAuthDisplay, CoreClaimName, CoreClaimType, CoreClient, CoreClientAuthMethod,
        CoreGrantType, CoreJsonWebKey, CoreJweContentEncryptionAlgorithm,
        CoreJweKeyManagementAlgorithm, CoreResponseMode, CoreResponseType,
        CoreSubjectIdentifierType,
    };
    use openidconnect::AdditionalProviderMetadata;
    use openidconnect::ProviderMetadata;
    impl AdditionalProviderMetadata for MyExtraProviderMetadata {}
    type MyProviderMetadata = ProviderMetadata<
        MyExtraProviderMetadata,
        CoreAuthDisplay,
        CoreClientAuthMethod,
        CoreClaimName,
        CoreClaimType,
        CoreGrantType,
        CoreJweContentEncryptionAlgorithm,
        CoreJweKeyManagementAlgorithm,
        CoreJsonWebKey,
        CoreResponseMode,
        CoreResponseType,
        CoreSubjectIdentifierType,
    >;

    // Discover provider metadata
    let provider_metadata: MyProviderMetadata =
        MyProviderMetadata::discover_async(issuer, &http_client).await?;

    let device_auth_url = provider_metadata
        .additional_metadata()
        .device_authorization_endpoint
        .clone();

    let client = CoreClient::from_provider_metadata(
        provider_metadata,
        client_id,
        None, // public client
    )
    .set_device_authorization_url(DeviceAuthorizationUrl::new(device_auth_url)?);

    let refresh_token = RefreshToken::new(old_secrets.refresh_token);
    let res = client
        .exchange_refresh_token(&refresh_token)?
        .request_async(&http_client)
        .await?;
    let expires_in = res.expires_in().ok_or(mize_err!(
        "no expiry information was in token response of the refresh flow"
    ))?;

    return Ok(ClientSecrets {
        id_token: res.id_token().unwrap().to_string(),
        access_token: res.access_token().secret().to_owned(),
        refresh_token: res.refresh_token().unwrap().secret().to_owned(),
        expiry: Utc::now() + expires_in,
    });
}

async fn auth_flow(mut mize: Mize) -> MizeResult<ClientSecrets> {
    let issuer = mize.get_config("auth.client.issuer")?.value_string()?;
    let client_id = mize.get_config("auth.client.client_id")?.value_string()?;

    let issuer = IssuerUrl::new(issuer)?;
    let client_id = ClientId::new(client_id);

    let http_client = HttpClient::new();

    // Teach openidconnect-rs about the device authorization url.
    #[derive(Clone, Debug, Deserialize, Serialize)]
    struct MyExtraProviderMetadata {
        device_authorization_endpoint: String,
    }
    use openidconnect::core::{
        CoreAuthDisplay, CoreClaimName, CoreClaimType, CoreClient, CoreClientAuthMethod,
        CoreGrantType, CoreJsonWebKey, CoreJweContentEncryptionAlgorithm,
        CoreJweKeyManagementAlgorithm, CoreResponseMode, CoreResponseType,
        CoreSubjectIdentifierType,
    };
    use openidconnect::AdditionalProviderMetadata;
    use openidconnect::ProviderMetadata;
    impl AdditionalProviderMetadata for MyExtraProviderMetadata {}
    type MyProviderMetadata = ProviderMetadata<
        MyExtraProviderMetadata,
        CoreAuthDisplay,
        CoreClientAuthMethod,
        CoreClaimName,
        CoreClaimType,
        CoreGrantType,
        CoreJweContentEncryptionAlgorithm,
        CoreJweKeyManagementAlgorithm,
        CoreJsonWebKey,
        CoreResponseMode,
        CoreResponseType,
        CoreSubjectIdentifierType,
    >;

    // Discover provider metadata
    let provider_metadata: MyProviderMetadata =
        MyProviderMetadata::discover_async(issuer, &http_client).await?;

    let device_auth_url = provider_metadata
        .additional_metadata()
        .device_authorization_endpoint
        .clone();

    let client = CoreClient::from_provider_metadata(
        provider_metadata,
        client_id,
        None, // public client
    )
    .set_device_authorization_url(DeviceAuthorizationUrl::new(device_auth_url)?);

    // ===== DEVICE AUTH REQUEST =====
    let details: Result<DeviceAuthorizationResponse<EmptyExtraDeviceAuthorizationFields>, _> =
        client
            .exchange_device_code()
            .add_scope(Scope::new("openid".into()))
            .add_scope(Scope::new("profile".into()))
            .add_scope(Scope::new("offline_access".into()))
            .request_async(&http_client)
            .await;

    let details = match details {
        Ok(val) => val,
        Err(err) => {
            println!("err: {:?}", err);
            println!("err: {}", err);
            panic!();
        }
    };

    println!(
        "If your browser does not open automatically, open this URL in it:\n{}\n",
        details.verification_uri_complete().unwrap().secret()
    );
    println!(
        "Or enter this PPC login code:\n{}\n",
        details.user_code().secret()
    );

    #[cfg(feature = "target-os")]
    {
        if let Some(_) = details.verification_uri_complete() {
            let _ = webbrowser::open(details.verification_uri_complete().unwrap().secret());
        }
    }

    let token_response = client
        .exchange_device_access_token(&details)?
        .request_async(
            &http_client,
            async |dur: Duration| {
                tokio::time::sleep(dur).await;
            },
            None,
        )
        .await?;

    // ===== TOKENS =====
    let id_token = token_response
        .extra_fields()
        .id_token()
        .ok_or(mize_err!("no id_token"))?;

    let refresh = token_response
        .refresh_token()
        .ok_or(mize_err!("no refresh token present"))?;

    let expires_in = token_response.expires_in().ok_or(mize_err!(
        "no expiry information was in token response of the refresh flow"
    ))?;

    Ok(ClientSecrets {
        id_token: id_token.to_string(),
        access_token: token_response.access_token().secret().to_string(),
        refresh_token: refresh.secret().to_owned(),
        expiry: Utc::now() + expires_in,
    })
}

fn update_spacetimedb_token(new_token: &str) -> MizeResult<()> {
    // Resolve config directory
    let config_dir = match std::env::var("XDG_CONFIG_HOME") {
        Ok(val) => PathBuf::from(val),
        Err(_) => {
            let home = std::env::var("HOME")?;
            std::path::PathBuf::from(home).join(".config")
        }
    };

    let file_path = config_dir.join("spacetime").join("cli.toml");

    // Read existing file (or create empty document if it doesn't exist)
    let content = std::fs::read_to_string(&file_path).unwrap_or_default();
    let mut data = ItemData::from_toml(content.as_str())?;

    // Update the field
    data.set_path("spacetimedb_token", new_token)?;

    // Write back
    std::fs::write(&file_path, data.to_toml()?)?;

    Ok(())
}