use crate::{Paths, is_plain, set_plain};
use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use std::cell::Cell;
use std::env;
use std::io::{Read, Write};
use std::net::TcpListener;
use std::path::Path;
use std::sync::{Mutex, MutexGuard};
use std::thread;
pub(crate) static ENV_MUTEX: Mutex<()> = Mutex::new(());
pub(crate) static PLAIN_MUTEX: Mutex<()> = Mutex::new(());
thread_local! {
static PLAIN_DEPTH: Cell<usize> = const { Cell::new(0) };
}
pub(crate) struct EnvVarGuard {
key: String,
prev: Option<String>,
}
pub(crate) struct PlainGuard {
prev: bool,
_lock: Option<MutexGuard<'static, ()>>,
}
fn set_env(key: &str, value: Option<&str>) -> Option<String> {
let prev = env::var(key).ok();
if let Some(value) = value {
unsafe {
env::set_var(key, value);
}
} else {
unsafe {
env::remove_var(key);
}
}
prev
}
pub(crate) fn set_env_guard(key: &str, value: Option<&str>) -> EnvVarGuard {
EnvVarGuard {
key: key.to_string(),
prev: set_env(key, value),
}
}
pub(crate) fn set_plain_guard(value: bool) -> PlainGuard {
let lock = PLAIN_DEPTH.with(|depth| {
let current = depth.get();
depth.set(current + 1);
if current == 0 {
Some(PLAIN_MUTEX.lock().unwrap())
} else {
None
}
});
let prev = is_plain();
set_plain(value);
PlainGuard { prev, _lock: lock }
}
fn restore_env(key: &str, prev: Option<String>) {
if let Some(value) = prev {
unsafe {
env::set_var(key, value);
}
} else {
unsafe {
env::remove_var(key);
}
}
}
impl Drop for EnvVarGuard {
fn drop(&mut self) {
let prev = self.prev.take();
restore_env(&self.key, prev);
}
}
impl Drop for PlainGuard {
fn drop(&mut self) {
set_plain(self.prev);
PLAIN_DEPTH.with(|depth| {
let current = depth.get();
depth.set(current.saturating_sub(1));
});
}
}
pub(crate) fn spawn_server(response: String) -> String {
let listener = TcpListener::bind("127.0.0.1:0").expect("bind");
let addr = listener.local_addr().unwrap();
thread::spawn(move || {
if let Ok((mut stream, _)) = listener.accept() {
let mut buf = [0u8; 4096];
let _ = stream.read(&mut buf);
let _ = stream.write_all(response.as_bytes());
}
});
format!("http://{}", addr)
}
pub(crate) fn build_id_token(email: &str, plan: &str) -> String {
let header = serde_json::json!({
"alg": "none",
"typ": "JWT",
});
let auth = serde_json::json!({
"chatgpt_plan_type": plan,
});
let payload = serde_json::json!({
"email": email,
"https://api.openai.com/auth": auth,
});
let header = URL_SAFE_NO_PAD.encode(serde_json::to_string(&header).unwrap());
let payload = URL_SAFE_NO_PAD.encode(serde_json::to_string(&payload).unwrap());
format!("{header}.{payload}.")
}
pub(crate) fn http_ok_response(body: &str, content_type: &str) -> String {
format!(
"HTTP/1.1 200 OK\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\n\r\n{body}",
body.len(),
)
}
pub(crate) fn make_paths(root: &Path) -> Paths {
let codex = root.to_path_buf();
let auth = codex.join("auth.json");
let profiles = codex.join("profiles");
let profiles_index = profiles.join("profiles.json");
let update_cache = profiles.join("update.json");
let profiles_lock = profiles.join("profiles.lock");
Paths {
codex,
auth,
profiles,
profiles_index,
update_cache,
profiles_lock,
}
}