rust - simple code to get access tokens via client credentials

This is my rust code to get a client credentials access token with the intention of doing performance test.  The client id, client secret and endpoint is configurable as environment variables.


use chrono::{DateTime, Utc, Duration};
use indicatif::{ProgressBar, ProgressStyle};
use serde::Deserialize;
use std::env;

#[derive(Debug, Deserialize)]
pub struct TokenResponse {
    pub access_token: String,
    pub expires_in: i64,
}

pub struct AuthClient {
    client_id: String,
    client_secret: String,
    token_url: String,
    http_client: reqwest::Client,
}

impl AuthClient {
    pub fn new(client_id: String, client_secret: String, token_url: String) -> Self {
        Self {
            client_id,
            client_secret,
            token_url,
            http_client: reqwest::Client::new(),
        }
    }

    pub async fn fetch_token(&self) -> Result<TokenResponse, String> {
        let params = [
            ("grant_type", "client_credentials"),
            ("client_id", &self.client_id),
            ("client_secret", &self.client_secret),
        ];

        let response = self.http_client
            .post(&self.token_url)
            .form(&params)
            .send()
            .await
            .map_err(|e| format!("Network Error: {}", e))?;

        // Check for 401 specifically
        if response.status() == reqwest::StatusCode::UNAUTHORIZED {
            return Err("ERROR 401: Invalid Client ID or Secret".to_string());
        }

        // Check for other non-success codes
        if !response.status().is_success() {
            return Err(format!("Server Error: {}", response.status()));
        }

        let token = response
            .json::<TokenResponse>()
            .await
            .map_err(|e| format!("JSON Parse Error: {}", e))?;

        Ok(token)
    }

    pub fn format_expiry(expires_in: i64) -> String {
        let expiry_time: DateTime<Utc> = Utc::now() + Duration::seconds(expires_in);
        expiry_time.format("%Y-%m-%d %H:%M:%S UTC").to_string()
    }
}

// --- MAIN EXECUTION LOGIC ---

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    dotenvy::dotenv().ok();

    let auth_endpoint = env::var("AUTH_ENDPOINT").expect("AUTH_ENDPOINT must be set");
    let client_id = env::var("CLIENT_ID").expect("CLIENT_ID must be set");
    let client_secret = env::var("CLIENT_SECRET").expect("CLIENT_SECRET must be set");
    let iterations: usize = env::var("ITERATIONS").unwrap_or("10".to_string()).parse()?;

    let auth_service = AuthClient::new(client_id, client_secret, auth_endpoint);

    let pb = ProgressBar::new(iterations as u64);
    pb.set_style(ProgressStyle::default_bar()
        .template("{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} {msg}")?
        .progress_chars("#>-"));

    println!("Running {} requests sequentially...\n", iterations);

    let mut last_token: Option<TokenResponse> = None;

    for i in 1..=iterations {
        match auth_service.fetch_token().await {
            Ok(token) => {
                //last_token = Some(token);
                println!("Latest Token:  {}...", &token.access_token[..10]);
                pb.inc(1);
            }
            Err(e) => {
                // Use pb.suspend to print to terminal without breaking the progress bar
                pb.suspend(|| {
                    eprintln!("[Request {} Failed] {}", i, e);
                });
                pb.inc(1);
            }
        }
    }

    pb.finish_with_message("Done!");

    if let Some(token) = last_token {
        println!("\n--- Success ---");
        println!("Latest Token:  {}...", &token.access_token[..10]);
        println!("Expiry Date:   {}", AuthClient::format_expiry(token.expires_in));
    } else {
        println!("\n❌ Failure: No tokens were successfully retrieved.");
    }

    Ok(())
}




Comments

Popular posts from this blog

vllm : Failed to infer device type

NodeJS: Error: spawn EINVAL in window for node version 20.20 and 18.20

android studio kotlin source is null error