refactor(error): Refactor error handling

mitm/src/lib.rs

@@ -21,7 +21,7 @@ pub struct Builder<T: HttpHandler + Clone> {
 }
 
 impl<T: HttpHandler + Clone> Builder<T> {
-    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!")?;

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,
+}

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<ArkoseToken> {
         }
     }
 
-    anyhow::bail!("No solver available")
+    anyhow::bail!(ArkoseError::NoSolverAvailable)
 }
 
 #[inline]
@@ -434,9 +436,9 @@ async fn submit_captcha(
 ) -> anyhow::Result<ArkoseToken> {
     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)),
     };
 }

openai/src/error.rs → 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,
-}

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")]

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};

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;

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};

openai/src/auth/provide/web.rs

@@ -1,22 +1,19 @@
-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};
 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,
 }

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;

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<String>,

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<impl IntoResponse, ResponseError>
             let resp: Response<Body> = 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<Response<Body>> 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<Response<Body>> 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
     {