mirror of
				https://github.com/MindWorkAI/AI-Studio.git
				synced 2025-11-04 15:20:20 +00:00 
			
		
		
		
	Refactored api token
This commit is contained in:
		
							parent
							
								
									45f0ab234f
								
							
						
					
					
						commit
						565a3edf30
					
				
							
								
								
									
										68
									
								
								runtime/src/api_token.rs
									
									
									
									
									
										Normal file
									
								
							
							
						
						
									
										68
									
								
								runtime/src/api_token.rs
									
									
									
									
									
										Normal file
									
								
							@ -0,0 +1,68 @@
 | 
				
			|||||||
 | 
					use once_cell::sync::Lazy;
 | 
				
			||||||
 | 
					use rand::{RngCore, SeedableRng};
 | 
				
			||||||
 | 
					use rocket::http::Status;
 | 
				
			||||||
 | 
					use rocket::Request;
 | 
				
			||||||
 | 
					use rocket::request::FromRequest;
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					pub static API_TOKEN: Lazy<APIToken> = Lazy::new(|| {
 | 
				
			||||||
 | 
					    let mut token = [0u8; 32];
 | 
				
			||||||
 | 
					    let mut rng = rand_chacha::ChaChaRng::from_entropy();
 | 
				
			||||||
 | 
					    rng.fill_bytes(&mut token);
 | 
				
			||||||
 | 
					    APIToken::from_bytes(token.to_vec())
 | 
				
			||||||
 | 
					});
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					pub struct APIToken{
 | 
				
			||||||
 | 
					    hex_text: String,
 | 
				
			||||||
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					impl APIToken {
 | 
				
			||||||
 | 
					    fn from_bytes(bytes: Vec<u8>) -> Self {
 | 
				
			||||||
 | 
					        APIToken {
 | 
				
			||||||
 | 
					            hex_text: bytes.iter().fold(String::new(), |mut result, byte| {
 | 
				
			||||||
 | 
					                result.push_str(&format!("{:02x}", byte));
 | 
				
			||||||
 | 
					                result
 | 
				
			||||||
 | 
					            }),
 | 
				
			||||||
 | 
					        }
 | 
				
			||||||
 | 
					    }
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					    fn from_hex_text(hex_text: &str) -> Self {
 | 
				
			||||||
 | 
					        APIToken {
 | 
				
			||||||
 | 
					            hex_text: hex_text.to_string(),
 | 
				
			||||||
 | 
					        }
 | 
				
			||||||
 | 
					    }
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					    pub(crate) fn to_hex_text(&self) -> &str {
 | 
				
			||||||
 | 
					        self.hex_text.as_str()
 | 
				
			||||||
 | 
					    }
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					    fn validate(&self, received_token: &Self) -> bool {
 | 
				
			||||||
 | 
					        received_token.to_hex_text() == self.to_hex_text()
 | 
				
			||||||
 | 
					    }
 | 
				
			||||||
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					#[rocket::async_trait]
 | 
				
			||||||
 | 
					impl<'r> FromRequest<'r> for APIToken {
 | 
				
			||||||
 | 
					    type Error = APITokenError;
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					    async fn from_request(request: &'r Request<'_>) -> RequestOutcome<Self, Self::Error> {
 | 
				
			||||||
 | 
					        let token = request.headers().get_one("token");
 | 
				
			||||||
 | 
					        match token {
 | 
				
			||||||
 | 
					            Some(token) => {
 | 
				
			||||||
 | 
					                let received_token = APIToken::from_hex_text(token);
 | 
				
			||||||
 | 
					                if API_TOKEN.validate(&received_token) {
 | 
				
			||||||
 | 
					                    RequestOutcome::Success(received_token)
 | 
				
			||||||
 | 
					                } else {
 | 
				
			||||||
 | 
					                    RequestOutcome::Error((Status::Unauthorized, APITokenError::Invalid))
 | 
				
			||||||
 | 
					                }
 | 
				
			||||||
 | 
					            }
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					            None => RequestOutcome::Error((Status::Unauthorized, APITokenError::Missing)),
 | 
				
			||||||
 | 
					        }
 | 
				
			||||||
 | 
					    }
 | 
				
			||||||
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
 | 
					#[derive(Debug)]
 | 
				
			||||||
 | 
					enum APITokenError {
 | 
				
			||||||
 | 
					    Missing,
 | 
				
			||||||
 | 
					    Invalid,
 | 
				
			||||||
 | 
					}
 | 
				
			||||||
@ -2,3 +2,4 @@ pub mod encryption;
 | 
				
			|||||||
pub mod log;
 | 
					pub mod log;
 | 
				
			||||||
pub mod environment;
 | 
					pub mod environment;
 | 
				
			||||||
pub mod network;
 | 
					pub mod network;
 | 
				
			||||||
 | 
					pub mod api_token;
 | 
				
			||||||
@ -64,13 +64,6 @@ static MAIN_WINDOW: Lazy<Mutex<Option<Window>>> = Lazy::new(|| Mutex::new(None))
 | 
				
			|||||||
// The update response coming from the Tauri updater.
 | 
					// The update response coming from the Tauri updater.
 | 
				
			||||||
static CHECK_UPDATE_RESPONSE: Lazy<Mutex<Option<UpdateResponse<tauri::Wry>>>> = Lazy::new(|| Mutex::new(None));
 | 
					static CHECK_UPDATE_RESPONSE: Lazy<Mutex<Option<UpdateResponse<tauri::Wry>>>> = Lazy::new(|| Mutex::new(None));
 | 
				
			||||||
 | 
					
 | 
				
			||||||
static API_TOKEN: Lazy<APIToken> = Lazy::new(|| {
 | 
					 | 
				
			||||||
    let mut token = [0u8; 32];
 | 
					 | 
				
			||||||
    let mut rng = rand_chacha::ChaChaRng::from_entropy();
 | 
					 | 
				
			||||||
    rng.fill_bytes(&mut token);
 | 
					 | 
				
			||||||
    APIToken::from_bytes(token.to_vec())
 | 
					 | 
				
			||||||
});
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
static DATA_DIRECTORY: OnceLock<String> = OnceLock::new();
 | 
					static DATA_DIRECTORY: OnceLock<String> = OnceLock::new();
 | 
				
			||||||
 | 
					
 | 
				
			||||||
static CONFIG_DIRECTORY: OnceLock<String> = OnceLock::new();
 | 
					static CONFIG_DIRECTORY: OnceLock<String> = OnceLock::new();
 | 
				
			||||||
@ -382,62 +375,8 @@ async fn main() {
 | 
				
			|||||||
    }
 | 
					    }
 | 
				
			||||||
}
 | 
					}
 | 
				
			||||||
 | 
					
 | 
				
			||||||
struct APIToken{
 | 
					 | 
				
			||||||
    hex_text: String,
 | 
					 | 
				
			||||||
}
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
impl APIToken {
 | 
					 | 
				
			||||||
    fn from_bytes(bytes: Vec<u8>) -> Self {
 | 
					 | 
				
			||||||
        APIToken {
 | 
					 | 
				
			||||||
            hex_text: bytes.iter().fold(String::new(), |mut result, byte| {
 | 
					 | 
				
			||||||
                result.push_str(&format!("{:02x}", byte));
 | 
					 | 
				
			||||||
                result
 | 
					 | 
				
			||||||
            }),
 | 
					 | 
				
			||||||
        }
 | 
					 | 
				
			||||||
    }
 | 
					    }
 | 
				
			||||||
 | 
					
 | 
				
			||||||
    fn from_hex_text(hex_text: &str) -> Self {
 | 
					 | 
				
			||||||
        APIToken {
 | 
					 | 
				
			||||||
            hex_text: hex_text.to_string(),
 | 
					 | 
				
			||||||
        }
 | 
					 | 
				
			||||||
    }
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
    fn to_hex_text(&self) -> &str {
 | 
					 | 
				
			||||||
        self.hex_text.as_str()
 | 
					 | 
				
			||||||
    }
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
    fn validate(&self, received_token: &Self) -> bool {
 | 
					 | 
				
			||||||
        received_token.to_hex_text() == self.to_hex_text()
 | 
					 | 
				
			||||||
    }
 | 
					 | 
				
			||||||
}
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
#[rocket::async_trait]
 | 
					 | 
				
			||||||
impl<'r> FromRequest<'r> for APIToken {
 | 
					 | 
				
			||||||
    type Error = APITokenError;
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
    async fn from_request(request: &'r Request<'_>) -> RequestOutcome<Self, Self::Error> {
 | 
					 | 
				
			||||||
        let token = request.headers().get_one("token");
 | 
					 | 
				
			||||||
        match token {
 | 
					 | 
				
			||||||
            Some(token) => {
 | 
					 | 
				
			||||||
                let received_token = APIToken::from_hex_text(token);
 | 
					 | 
				
			||||||
                if API_TOKEN.validate(&received_token) {
 | 
					 | 
				
			||||||
                    RequestOutcome::Success(received_token)
 | 
					 | 
				
			||||||
                } else {
 | 
					 | 
				
			||||||
                    RequestOutcome::Error((Status::Unauthorized, APITokenError::Invalid))
 | 
					 | 
				
			||||||
                }
 | 
					 | 
				
			||||||
            }
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
            None => RequestOutcome::Error((Status::Unauthorized, APITokenError::Missing)),
 | 
					 | 
				
			||||||
        }
 | 
					 | 
				
			||||||
    }
 | 
					 | 
				
			||||||
}
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
#[derive(Debug)]
 | 
					 | 
				
			||||||
enum APITokenError {
 | 
					 | 
				
			||||||
    Missing,
 | 
					 | 
				
			||||||
    Invalid,
 | 
					 | 
				
			||||||
}
 | 
					 | 
				
			||||||
 | 
					 | 
				
			||||||
#[get("/system/dotnet/port")]
 | 
					#[get("/system/dotnet/port")]
 | 
				
			||||||
fn dotnet_port(_token: APIToken) -> String {
 | 
					fn dotnet_port(_token: APIToken) -> String {
 | 
				
			||||||
    let dotnet_server_port = *DOTNET_SERVER_PORT;
 | 
					    let dotnet_server_port = *DOTNET_SERVER_PORT;
 | 
				
			||||||
 | 
				
			|||||||
		Loading…
	
		Reference in New Issue
	
	Block a user