From ad6f21f0ed5e89d7a24661f3d5c534ad0e1413f2 Mon Sep 17 00:00:00 2001 From: gngpp Date: Wed, 22 Nov 2023 19:41:04 +0800 Subject: [PATCH] refactor(error): Refactor error handling --- mitm/src/lib.rs | 2 +- openai/src/arkose/error.rs | 15 +++++++++++++++ openai/src/arkose/mod.rs | 14 ++++++++------ openai/src/{ => auth}/error.rs | 20 -------------------- openai/src/auth/mod.rs | 3 ++- openai/src/auth/provide/apple.rs | 10 ++++------ openai/src/auth/provide/mod.rs | 8 ++++---- openai/src/auth/provide/platform.rs | 10 ++++------ openai/src/auth/provide/web.rs | 21 +++++++++------------ openai/src/lib.rs | 1 - openai/src/serve/error.rs | 16 ++++++++++++++++ openai/src/serve/mod.rs | 27 +++++++++++++-------------- openai/src/serve/proxy/req.rs | 22 +++++++--------------- openai/src/serve/proxy/resp.rs | 3 ++- 14 files changed, 85 insertions(+), 87 deletions(-) create mode 100644 openai/src/arkose/error.rs rename openai/src/{ => auth}/error.rs (75%) diff --git a/mitm/src/lib.rs b/mitm/src/lib.rs index 5014da921..5b06f198f 100644 --- a/mitm/src/lib.rs +++ b/mitm/src/lib.rs @@ -21,7 +21,7 @@ pub struct Builder { } impl Builder { - pub async fn mitm_proxy(self) -> anyhow::Result<()> { + pub async fn proxy(self) -> anyhow::Result<()> { info!("PreAuth CA Private key use: {}", self.key.display()); let private_key_bytes = fs::read(self.key).context("ca private key file path not valid!")?; diff --git a/openai/src/arkose/error.rs b/openai/src/arkose/error.rs new file mode 100644 index 000000000..adab83eee --- /dev/null +++ b/openai/src/arkose/error.rs @@ -0,0 +1,15 @@ +#[derive(thiserror::Error, Debug)] +pub enum ArkoseError { + #[error("submit funcaptcha answer error {0:?}")] + SubmitAnswerError(anyhow::Error), + #[error("Invalid arkose platform type: {0:?}")] + InvalidPlatformType(String), + #[error("Invalid GPT model: {0:?}")] + InvalidGptModel(String), + #[error("No solver available or solver is invalid")] + NoSolverAvailable, + #[error("Error creating arkose session error {0:?}")] + CreateSessionError(anyhow::Error), + #[error("invalid funcaptcha error")] + InvalidFunCaptcha, +} diff --git a/openai/src/arkose/mod.rs b/openai/src/arkose/mod.rs index 8ae6fb4eb..90336c7aa 100644 --- a/openai/src/arkose/mod.rs +++ b/openai/src/arkose/mod.rs @@ -1,4 +1,5 @@ pub mod crypto; +mod error; pub mod funcaptcha; pub mod har; pub mod murmur; @@ -23,6 +24,7 @@ use crate::generate_random_string; use crate::warn; use crate::with_context; use crate::HEADER_UA; +use error::ArkoseError; use self::funcaptcha::solver::SubmitSolver; use self::funcaptcha::ArkoseSolver; @@ -65,7 +67,7 @@ impl std::str::FromStr for Type { "gpt4" => Ok(Type::GPT4), "auth" => Ok(Type::Auth), "platform" => Ok(Type::Platform), - _ => anyhow::bail!("Invalid type"), + _ => anyhow::bail!(ArkoseError::InvalidPlatformType(s.to_owned())), } } } @@ -114,7 +116,7 @@ impl std::str::FromStr for GPTModel { s if s.starts_with("gpt-3.5") || s.starts_with("text-davinci") => { Ok(GPTModel::Gpt35Other) } - _ => anyhow::bail!("Invalid GPT model"), + _ => anyhow::bail!(ArkoseError::InvalidGptModel(value.to_owned())), } } } @@ -391,7 +393,7 @@ async fn get_from_context(t: Type) -> anyhow::Result { } } - anyhow::bail!("No solver available") + anyhow::bail!(ArkoseError::NoSolverAvailable) } #[inline] @@ -434,9 +436,9 @@ async fn submit_captcha( ) -> anyhow::Result { let session = funcaptcha::start_challenge(arkose_token.value()) .await - .map_err(|error| anyhow::anyhow!("Error creating session: {error}"))?; + .map_err(ArkoseError::CreateSessionError)?; - let funs = anyhow::Context::context(session.funcaptcha(), "Valid funcaptcha error")?; + let funs = session.funcaptcha().ok_or(ArkoseError::InvalidFunCaptcha)?; let mut rx = match solver { Solver::Yescaptcha => { let (tx, rx) = tokio::sync::mpsc::channel(funs.len()); @@ -528,6 +530,6 @@ async fn submit_captcha( let new_token = arkose_token.value().replace("at=40", "at=40|sup=1"); Ok(ArkoseToken::from(new_token)) } - Err(err) => anyhow::bail!("submit funcaptcha answer error: {err}"), + Err(err) => anyhow::bail!(ArkoseError::SubmitAnswerError(err)), }; } diff --git a/openai/src/error.rs b/openai/src/auth/error.rs similarity index 75% rename from openai/src/error.rs rename to openai/src/auth/error.rs index c00006ad1..70889e493 100644 --- a/openai/src/error.rs +++ b/openai/src/auth/error.rs @@ -8,8 +8,6 @@ pub enum AuthError { Unauthorized(String), #[error("Server error ({0:?})")] ServerError(String), - #[error("failed to get public key")] - FailedPubKeyRequest, #[error("failed login")] FailedLogin, #[error(transparent)] @@ -42,14 +40,10 @@ pub enum AuthError { InvalidEmail, #[error("invalid Location")] InvalidLocation, - #[error("invalid access token")] - InvalidAccessToken, #[error("invalid refresh token")] InvalidRefreshToken, #[error("invalid location path")] InvalidLocationPath, - #[error("token expired")] - TokenExpired, #[error("MFA failed")] MFAFailed, #[error("MFA required")] @@ -61,17 +55,3 @@ pub enum AuthError { #[error("failed to get preauth cookie")] PreauthCookieNotFound, } - -#[derive(thiserror::Error, Debug)] -pub enum TokenStoreError { - #[error("failed to access token")] - AccessError, - #[error("token not found error")] - NotFoundError, - #[error("failed token deserialize")] - DeserializeError(#[from] serde_json::error::Error), - #[error("failed to verify access_token")] - AccessTokenVerifyError, - #[error("failed to create default token store file")] - CreateDefaultTokenFileError, -} diff --git a/openai/src/auth/mod.rs b/openai/src/auth/mod.rs index 682587e51..26d044f56 100644 --- a/openai/src/auth/mod.rs +++ b/openai/src/auth/mod.rs @@ -1,3 +1,4 @@ +mod error; pub mod model; pub mod provide; @@ -21,9 +22,9 @@ use reqwest::{Client, Proxy, StatusCode, Url}; use sha2::{Digest, Sha256}; use tokio::sync::OnceCell; -use crate::error::AuthError; use crate::URL_CHATGPT_API; use crate::{debug, random_impersonate}; +use error::AuthError; use self::model::{ApiKeyData, AuthStrategy}; #[cfg(feature = "preauth")] diff --git a/openai/src/auth/provide/apple.rs b/openai/src/auth/provide/apple.rs index 250e3c044..9d9ea25e9 100644 --- a/openai/src/auth/provide/apple.rs +++ b/openai/src/auth/provide/apple.rs @@ -1,11 +1,9 @@ +use crate::auth::error::AuthError; use crate::auth::provide::{AuthenticateData, GrantType}; use crate::auth::AuthClient; -use crate::{ - auth::{ - model::{self, AuthStrategy}, - OPENAI_OAUTH_REVOKE_URL, OPENAI_OAUTH_TOKEN_URL, OPENAI_OAUTH_URL, - }, - error::AuthError, +use crate::auth::{ + model::{self, AuthStrategy}, + OPENAI_OAUTH_REVOKE_URL, OPENAI_OAUTH_TOKEN_URL, OPENAI_OAUTH_URL, }; use crate::{warn, with_context}; use anyhow::{bail, Context}; diff --git a/openai/src/auth/provide/mod.rs b/openai/src/auth/provide/mod.rs index f733764fa..e9865fa69 100644 --- a/openai/src/auth/provide/mod.rs +++ b/openai/src/auth/provide/mod.rs @@ -5,12 +5,12 @@ pub mod web; use std::collections::HashSet; -use crate::{ - arkose::{self, ArkoseToken, Type}, +use crate::arkose::{self, ArkoseToken, Type}; + +use super::{ error::AuthError, + model::{self, AuthStrategy}, }; - -use super::model::{self, AuthStrategy}; use reqwest::header; use serde::Serialize; use typed_builder::TypedBuilder; diff --git a/openai/src/auth/provide/platform.rs b/openai/src/auth/provide/platform.rs index 5445ce5de..94a0e79e1 100644 --- a/openai/src/auth/provide/platform.rs +++ b/openai/src/auth/provide/platform.rs @@ -1,11 +1,9 @@ +use crate::auth::error::AuthError; use crate::auth::provide::{AuthenticateData, GrantType}; use crate::auth::AuthClient; -use crate::{ - auth::{ - model::{self, AuthStrategy}, - OPENAI_OAUTH_REVOKE_URL, OPENAI_OAUTH_TOKEN_URL, OPENAI_OAUTH_URL, - }, - error::AuthError, +use crate::auth::{ + model::{self, AuthStrategy}, + OPENAI_OAUTH_REVOKE_URL, OPENAI_OAUTH_TOKEN_URL, OPENAI_OAUTH_URL, }; use crate::{debug, warn}; use anyhow::{bail, Context}; diff --git a/openai/src/auth/provide/web.rs b/openai/src/auth/provide/web.rs index ceb3418b7..0fdf3086d 100644 --- a/openai/src/auth/provide/web.rs +++ b/openai/src/auth/provide/web.rs @@ -1,10 +1,12 @@ -use crate::{ - auth::{ - model::{self, AuthStrategy}, - provide::AuthenticateData, - AuthClient, API_AUTH_SESSION_COOKIE_KEY, OPENAI_OAUTH_URL, - }, - error::AuthError, +use super::{ + AuthProvider, AuthResult, AuthenticateMfaData, GetAuthorizedUrlData, IdentifierData, + RequestContext, RequestContextExt, +}; +use crate::auth::error::AuthError; +use crate::auth::{ + model::{self, AuthStrategy}, + provide::AuthenticateData, + AuthClient, API_AUTH_SESSION_COOKIE_KEY, OPENAI_OAUTH_URL, }; use crate::{debug, warn, URL_CHATGPT_API}; use anyhow::{bail, Context}; @@ -12,11 +14,6 @@ use reqwest::{Client, StatusCode}; use serde_json::Value; use url::Url; -use super::{ - AuthProvider, AuthResult, AuthenticateMfaData, GetAuthorizedUrlData, IdentifierData, - RequestContext, RequestContextExt, -}; - pub(crate) struct WebAuthProvider { inner: Client, } diff --git a/openai/src/lib.rs b/openai/src/lib.rs index 52961bb56..6777df752 100644 --- a/openai/src/lib.rs +++ b/openai/src/lib.rs @@ -4,7 +4,6 @@ pub mod auth; pub mod balancer; pub mod chatgpt; pub mod context; -pub mod error; pub mod eventsource; pub mod homedir; pub mod log; diff --git a/openai/src/serve/error.rs b/openai/src/serve/error.rs index 9fc38715d..01d82e8a4 100644 --- a/openai/src/serve/error.rs +++ b/openai/src/serve/error.rs @@ -5,6 +5,22 @@ use axum::response::Response; use axum::Json; use serde_json::json; +#[derive(thiserror::Error, Debug)] +pub enum ProxyError { + #[error("Session not found")] + SessionNotFound, + #[error("Authentication Key error")] + AuthKeyError, + #[error("AccessToken required")] + AccessTokenRequired, + #[error("Model required")] + ModelRequired, + #[error("Body required")] + BodyRequired, + #[error("Body must be a json object")] + BodyMustBeJsonObject, +} + // Make our own error that wraps `anyhow::Error`. pub struct ResponseError { msg: Option, diff --git a/openai/src/serve/mod.rs b/openai/src/serve/mod.rs index f2fb62df6..4bcbb8dd8 100644 --- a/openai/src/serve/mod.rs +++ b/openai/src/serve/mod.rs @@ -27,6 +27,7 @@ use crate::auth::model::{AccessToken, AuthAccount, RefreshToken, SessionAccessTo use crate::auth::provide::AuthProvider; use crate::auth::API_AUTH_SESSION_COOKIE_KEY; use crate::context::{self, ContextArgs}; +use crate::serve::error::ProxyError; use crate::serve::error::ResponseError; use crate::serve::middleware::tokenbucket::{Strategy, TokenBucketLimitContext}; use crate::{info, warn, with_context}; @@ -199,8 +200,8 @@ impl Serve { .mitm_filters(vec![String::from("ios.chat.openai.com")]) .handler(preauth::PreAuthHanlder) .build(); - if let Some(err) = builder.mitm_proxy().await.err() { - crate::error!("PreAuth proxy error: {}", err); + if let Some(err) = builder.proxy().await.err() { + warn!("PreAuth proxy error: {}", err); } } @@ -267,9 +268,9 @@ async fn get_session(jar: CookieJar) -> Result let resp: Response = session_token.try_into()?; Ok(resp.into_response()) } - _ => Err(ResponseError::InternalServerError(anyhow!( - "Session error!" - ))), + _ => Err(ResponseError::InternalServerError( + ProxyError::SessionNotFound, + )), } } @@ -282,9 +283,7 @@ async fn post_access_token( // check bearer token exist let bearer = bearer.ok_or(ResponseError::Unauthorized(anyhow!("Auth Key required!")))?; if auth_key.ne(bearer.token()) { - return Err(ResponseError::Unauthorized(anyhow!( - "Authentication Key error!" - ))); + return Err(ResponseError::Unauthorized(ProxyError::AuthKeyError)); } } @@ -354,15 +353,15 @@ impl TryInto> for SessionAccessToken { let session = self .session_token .clone() - .ok_or(ResponseError::InternalServerError(anyhow!( - "Session error!" - )))?; + .ok_or(ResponseError::InternalServerError( + ProxyError::SessionNotFound, + ))?; let timestamp_secs = session .expires .unwrap_or_else(|| SystemTime::now()) .duration_since(UNIX_EPOCH) - .expect("Failed to get timestamp") + .map_err(ResponseError::InternalServerError)? .as_secs_f64(); let cookie = cookie::Cookie::build(API_AUTH_SESSION_COOKIE_KEY, session.value) @@ -378,7 +377,7 @@ impl TryInto> for SessionAccessToken { Ok(Response::builder() .status(axum::http::StatusCode::OK) .header(header::SET_COOKIE, cookie.to_string()) - .header(header::CONTENT_TYPE, "application/json") + .header(header::CONTENT_TYPE, mime::APPLICATION_JSON.as_ref()) .body(Body::from(serde_json::to_string(&self)?)) .map_err(ResponseError::InternalServerError)?) } @@ -388,7 +387,7 @@ async fn check_wan_address() { match with_context!(client) .get("https://ifconfig.me") .timeout(Duration::from_secs(70)) - .header(header::ACCEPT, "application/json") + .header(header::ACCEPT, mime::APPLICATION_JSON.as_ref()) .send() .await { diff --git a/openai/src/serve/proxy/req.rs b/openai/src/serve/proxy/req.rs index 13355f190..016931478 100644 --- a/openai/src/serve/proxy/req.rs +++ b/openai/src/serve/proxy/req.rs @@ -15,7 +15,7 @@ use crate::{arkose, with_context}; use super::ext::{RequestExt, ResponseExt, SendRequestExt}; use super::resp::header_convert; use super::toapi; -use crate::serve::error::ResponseError; +use crate::serve::error::{ProxyError, ResponseError}; use crate::serve::puid::{get_or_init, reduce_key}; use crate::serve::EMPTY; @@ -81,32 +81,26 @@ async fn handle_conv_request(req: &mut RequestExt) -> Result<(), ResponseError> let body = req .body .as_ref() - .ok_or(ResponseError::BadRequest(anyhow::anyhow!( - "Body can not be empty!" - )))?; + .ok_or(ResponseError::BadRequest(ProxyError::BodyRequired))?; // Use serde_json to parse body let mut json = serde_json::from_slice::(&body).map_err(ResponseError::BadRequest)?; let body = json .as_object_mut() - .ok_or(ResponseError::BadRequest(anyhow::anyhow!("Body is empty")))?; + .ok_or(ResponseError::BadRequest(ProxyError::BodyMustBeJsonObject))?; // If model is not exist, then return error let model = body .get("model") .and_then(|m| m.as_str()) - .ok_or(ResponseError::BadRequest(anyhow::anyhow!( - "Model is not exist in body!" - )))?; + .ok_or(ResponseError::BadRequest(ProxyError::ModelRequired))?; // If puid is exist, then return if !has_puid(&req.headers)? { // extract token from Authorization header let token = req .bearer_auth() - .ok_or(ResponseError::Unauthorized(anyhow::anyhow!( - "AccessToken required!" - )))?; + .ok_or(ResponseError::Unauthorized(ProxyError::AccessTokenRequired))?; // Exstract the token from the Authorization header let cache_id = reduce_key(token)?; @@ -162,15 +156,13 @@ async fn handle_dashboard_request(req: &mut RequestExt) -> Result<(), ResponseEr let body = req .body .as_ref() - .ok_or(ResponseError::BadRequest(anyhow::anyhow!( - "Body can not be empty!" - )))?; + .ok_or(ResponseError::BadRequest(ProxyError::BodyRequired))?; // Use serde_json to parse body let mut json = serde_json::from_slice::(&body).map_err(ResponseError::BadRequest)?; let body = json .as_object_mut() - .ok_or(ResponseError::BadRequest(anyhow::anyhow!("Body is empty")))?; + .ok_or(ResponseError::BadRequest(ProxyError::BodyMustBeJsonObject))?; // If arkose_token is not exist, then add it if body.get("arkose_token").is_none() { diff --git a/openai/src/serve/proxy/resp.rs b/openai/src/serve/proxy/resp.rs index 012a1ee0d..24a854fb0 100644 --- a/openai/src/serve/proxy/resp.rs +++ b/openai/src/serve/proxy/resp.rs @@ -66,7 +66,8 @@ pub(crate) fn header_convert( if !cookies.is_empty() { headers.insert( header::COOKIE, - header::HeaderValue::from_str(&cookies.join(";")).expect("setting cookie error"), + header::HeaderValue::from_str(&cookies.join(";")) + .map_err(ResponseError::InternalServerError)?, ); } Ok(headers)