From da23673a4485d8178b01e3ffe6898c0029f473df Mon Sep 17 00:00:00 2001 From: Akarshan Biswas Date: Thu, 29 May 2025 12:02:55 +0530 Subject: [PATCH] feat: Add API key generation for Llama.cpp This commit introduces API key generation for the Llama.cpp extension. The API key is now generated on the server side using HMAC-SHA256 and a secret key to ensure security and uniqueness. The frontend now passes the model ID and API secret to the server to generate the key. This addresses the requirement for secure model access and authorization. --- Makefile | 62 ++++++------- .../browser/extensions/engines/AIEngine.ts | 1 + extensions/llamacpp-extension/src/index.ts | 12 +-- src-tauri/Cargo.toml | 3 + src-tauri/src/core/setup.rs | 2 + .../inference_llamacpp_extension/server.rs | 87 +++++++++++-------- src-tauri/src/lib.rs | 1 + 7 files changed, 97 insertions(+), 71 deletions(-) diff --git a/Makefile b/Makefile index 6e69c602d..51f24885c 100644 --- a/Makefile +++ b/Makefile @@ -66,36 +66,36 @@ ifeq ($(OS),Windows_NT) -powershell -Command "Remove-Item -Recurse -Force ./src-tauri/target" -powershell -Command "if (Test-Path \"$($env:USERPROFILE)\jan\extensions\") { Remove-Item -Path \"$($env:USERPROFILE)\jan\extensions\" -Recurse -Force }" else ifeq ($(shell uname -s),Linux) - find . -name "node_modules" -type d -prune -exec rm -rf '{}' + - find . -name ".next" -type d -exec rm -rf '{}' + - find . -name "dist" -type d -exec rm -rf '{}' + - find . -name "build" -type d -exec rm -rf '{}' + - find . -name "out" -type d -exec rm -rf '{}' + - find . -name ".turbo" -type d -exec rm -rf '{}' + - find . -name ".yarn" -type d -exec rm -rf '{}' + - find . -name "packake-lock.json" -type f -exec rm -rf '{}' + - find . -name "package-lock.json" -type f -exec rm -rf '{}' + - rm -rf ./pre-install/*.tgz - rm -rf ./extensions/*/*.tgz - rm -rf ./electron/pre-install/*.tgz - rm -rf ./src-tauri/resources - rm -rf ./src-tauri/target - rm -rf "~/jan/extensions" - rm -rf "~/.cache/jan*" + find . -name "node_modules" -type d -prune -exec rm -rfv '{}' + + find . -name ".next" -type d -exec rm -rfv '{}' + + find . -name "dist" -type d -exec rm -rfv '{}' + + find . -name "build" -type d -exec rm -rfv '{}' + + find . -name "out" -type d -exec rm -rfv '{}' + + find . -name ".turbo" -type d -exec rm -rfv '{}' + + find . -name ".yarn" -type d -exec rm -rfv '{}' + + find . -name "packake-lock.json" -type f -exec rm -rfv '{}' + + find . -name "package-lock.json" -type f -exec rm -rfv '{}' + + rm -rfv ./pre-install/*.tgz + rm -rfv ./extensions/*/*.tgz + rm -rfv ./electron/pre-install/*.tgz + rm -rfv ./src-tauri/resources + rm -rfv ./src-tauri/target + rm -rfv "~/jan/extensions" + rm -rfv "~/.cache/jan*" else - find . -name "node_modules" -type d -prune -exec rm -rf '{}' + - find . -name ".next" -type d -exec rm -rf '{}' + - find . -name "dist" -type d -exec rm -rf '{}' + - find . -name "build" -type d -exec rm -rf '{}' + - find . -name "out" -type d -exec rm -rf '{}' + - find . -name ".turbo" -type d -exec rm -rf '{}' + - find . -name ".yarn" -type d -exec rm -rf '{}' + - find . -name "package-lock.json" -type f -exec rm -rf '{}' + - rm -rf ./pre-install/*.tgz - rm -rf ./extensions/*/*.tgz - rm -rf ./electron/pre-install/*.tgz - rm -rf ./src-tauri/resources - rm -rf ./src-tauri/target - rm -rf ~/jan/extensions - rm -rf ~/Library/Caches/jan* + find . -name "node_modules" -type d -prune -exec rm -rfv '{}' + + find . -name ".next" -type d -exec rm -rfv '{}' + + find . -name "dist" -type d -exec rm -rfv '{}' + + find . -name "build" -type d -exec rm -rfv '{}' + + find . -name "out" -type d -exec rm -rfv '{}' + + find . -name ".turbo" -type d -exec rm -rfv '{}' + + find . -name ".yarn" -type d -exec rm -rfv '{}' + + find . -name "package-lock.json" -type f -exec rm -rfv '{}' + + rm -rfv ./pre-install/*.tgz + rm -rfv ./extensions/*/*.tgz + rm -rfv ./electron/pre-install/*.tgz + rm -rfv ./src-tauri/resources + rm -rfv ./src-tauri/target + rm -rfv ~/jan/extensions + rm -rfv ~/Library/Caches/jan* endif diff --git a/core/src/browser/extensions/engines/AIEngine.ts b/core/src/browser/extensions/engines/AIEngine.ts index 9947a9604..d957e8f41 100644 --- a/core/src/browser/extensions/engines/AIEngine.ts +++ b/core/src/browser/extensions/engines/AIEngine.ts @@ -101,6 +101,7 @@ export type listResult = modelInfo[] // 3. /load export interface loadOptions { + modelId: string modelPath: string port?: number } diff --git a/extensions/llamacpp-extension/src/index.ts b/extensions/llamacpp-extension/src/index.ts index d883ac635..d12f9420d 100644 --- a/extensions/llamacpp-extension/src/index.ts +++ b/extensions/llamacpp-extension/src/index.ts @@ -23,7 +23,6 @@ import { } from '@janhq/core' import { invoke } from '@tauri-apps/api/core' -// import { createHmac } from 'crypto' type LlamacppConfig = { n_gpu_layers: number; @@ -130,10 +129,12 @@ export default class llamacpp_extension extends AIEngine { this.config[key] = value } - private generateApiKey(modelId: string): string { - return '' - // const hash = createHmac('sha256', this.apiSecret).update(modelId).digest("base64") - // return hash + private async generateApiKey(modelId: string): Promise { + const hash = await invoke('generate_api_key', { + modelId: modelId, + apiSecret: this.apiSecret + }) + return hash } // Implement the required LocalProvider interface methods @@ -304,6 +305,7 @@ export default class llamacpp_extension extends AIEngine { // model option is required // TODO: llama.cpp extension lookup model path based on modelId args.push('-m', opts.modelPath) + args.push('-a', opts.modelId) args.push('--port', String(opts.port || 8080)) // Default port if not specified if (cfg.ctx_size !== undefined) { diff --git a/src-tauri/Cargo.toml b/src-tauri/Cargo.toml index df939aa55..b6a8b4936 100644 --- a/src-tauri/Cargo.toml +++ b/src-tauri/Cargo.toml @@ -53,6 +53,9 @@ nvml-wrapper = "0.10.0" tauri-plugin-deep-link = "2" fix-path-env = { git = "https://github.com/tauri-apps/fix-path-env-rs" } serde_yaml = "0.9.34" +hmac = "0.12.1" +sha2 = "0.10.9" +base64 = "0.22.1" [target.'cfg(windows)'.dependencies] libloading = "0.8.7" diff --git a/src-tauri/src/core/setup.rs b/src-tauri/src/core/setup.rs index 44cff7aea..bbf8dc9e7 100644 --- a/src-tauri/src/core/setup.rs +++ b/src-tauri/src/core/setup.rs @@ -10,6 +10,8 @@ use tauri_plugin_store::StoreExt; use tokio::sync::Mutex; use tokio::time::{sleep, Duration}; // Using tokio::sync::Mutex // MCP + +// MCP use super::{ cmd::{get_jan_data_folder_path, get_jan_extensions_path}, mcp::run_mcp_commands, diff --git a/src-tauri/src/core/utils/extensions/inference_llamacpp_extension/server.rs b/src-tauri/src/core/utils/extensions/inference_llamacpp_extension/server.rs index 2f520b314..5de95f099 100644 --- a/src-tauri/src/core/utils/extensions/inference_llamacpp_extension/server.rs +++ b/src-tauri/src/core/utils/extensions/inference_llamacpp_extension/server.rs @@ -1,24 +1,25 @@ use std::path::PathBuf; use serde::{Serialize, Deserialize}; -use tauri::path::BaseDirectory; -use tauri::{AppHandle, Manager, State}; // Import Manager trait +use tauri::{AppHandle, State}; // Import Manager trait use tokio::process::Command; use uuid::Uuid; use thiserror; +use hmac::{Hmac, Mac}; +use sha2::Sha256; +use base64::{Engine as _, engine::general_purpose}; use crate::core::state::AppState; +type HmacSha256 = Hmac; // Error type for server commands #[derive(Debug, thiserror::Error)] -pub enum ServerError { +pub enum serverError { #[error("Server is already running")] AlreadyRunning, // #[error("Server is not running")] // NotRunning, #[error("Failed to locate server binary: {0}")] BinaryNotFound(String), - #[error("Failed to determine resource path: {0}")] - ResourcePathError(String), #[error("IO error: {0}")] Io(#[from] std::io::Error), #[error("Jan API error: {0}")] @@ -26,7 +27,7 @@ pub enum ServerError { } // impl serialization for tauri -impl serde::Serialize for ServerError { +impl serde::Serialize for serverError { fn serialize(&self, serializer: S) -> Result where S: serde::Serializer, @@ -35,18 +36,19 @@ impl serde::Serialize for ServerError { } } -type ServerResult = Result; +type ServerResult = Result; #[derive(Debug, Serialize, Deserialize)] -pub struct SessionInfo { - pub session_id: String, // opaque handle for unload/chat - pub port: u16, // llama-server output port - pub model_path: String, // path of the loaded model - pub api_key: String, +pub struct sessionInfo { + pub pid: String, // opaque handle for unload/chat + pub port: u16, // llama-server output port + pub modelId: String, + pub modelPath: String, // path of the loaded model + pub apiKey: String, } #[derive(serde::Serialize, serde::Deserialize)] -pub struct UnloadResult { +pub struct unloadResult { success: bool, error: Option, } @@ -54,40 +56,41 @@ pub struct UnloadResult { // --- Load Command --- #[tauri::command] pub async fn load_llama_model( - app_handle: AppHandle, // Get the AppHandle + _app_handle: AppHandle, // Get the AppHandle state: State<'_, AppState>, // Access the shared state - server_path: String, + engineBasePath: String, args: Vec, // Arguments from the frontend -) -> ServerResult { +) -> ServerResult { let mut process_lock = state.llama_server_process.lock().await; if process_lock.is_some() { log::warn!("Attempted to load server, but it's already running."); - return Err(ServerError::AlreadyRunning); + return Err(serverError::AlreadyRunning); } - log::info!("Attempting to launch server at path: {:?}", server_path); + log::info!("Attempting to launch server at path: {:?}", engineBasePath); log::info!("Using arguments: {:?}", args); - let server_path_buf = PathBuf::from(&server_path); + let server_path_buf = PathBuf::from(&engineBasePath); if !server_path_buf.exists() { log::error!( "Server binary not found at expected path: {:?}", - server_path + engineBasePath ); - return Err(ServerError::BinaryNotFound(format!( + return Err(serverError::BinaryNotFound(format!( "Binary not found at {:?}", - server_path + engineBasePath ))); } let port = 8080; // Default port // Configure the command to run the server - let mut command = Command::new(server_path); + let mut command = Command::new(engineBasePath); - let model_path = args[2].replace("-m", ""); - let api_key = args[1].replace("--api-key", ""); + let modelPath = args[2].replace("-m", ""); + let apiKey = args[1].replace("--api-key", ""); + let modelId = args[3].replace("-a", ""); command.args(args); // Optional: Redirect stdio if needed (e.g., for logging within Jan) @@ -95,7 +98,7 @@ pub async fn load_llama_model( // command.stderr(Stdio::piped()); // Spawn the child process - let child = command.spawn().map_err(ServerError::Io)?; + let child = command.spawn().map_err(serverError::Io)?; // Get the PID to use as session ID let pid = child.id().map(|id| id.to_string()).unwrap_or_else(|| { @@ -108,11 +111,12 @@ pub async fn load_llama_model( // Store the child process handle in the state *process_lock = Some(child); - let session_info = SessionInfo { - session_id: pid, // Use PID as session ID + let session_info = sessionInfo { + pid, port, - model_path, - api_key, + modelId, + modelPath, + apiKey, }; Ok(session_info) @@ -120,7 +124,7 @@ pub async fn load_llama_model( // --- Unload Command --- #[tauri::command] -pub async fn unload_llama_model(session_id: String, state: State<'_, AppState>) -> ServerResult { +pub async fn unload_llama_model(session_id: String, state: State<'_, AppState>) -> ServerResult { let mut process_lock = state.llama_server_process.lock().await; // Take the child process out of the Option, leaving None in its place if let Some(mut child) = process_lock.take() { @@ -138,7 +142,7 @@ pub async fn unload_llama_model(session_id: String, state: State<'_, AppState>) process_pid ); - return Ok(UnloadResult { + return Ok(unloadResult { success: false, error: Some(format!("Session ID mismatch: provided {} doesn't match process {}", session_id, process_pid)), @@ -155,7 +159,7 @@ pub async fn unload_llama_model(session_id: String, state: State<'_, AppState>) Ok(_) => { log::info!("Server process termination signal sent successfully"); - Ok(UnloadResult { + Ok(unloadResult { success: true, error: None, }) @@ -164,7 +168,7 @@ pub async fn unload_llama_model(session_id: String, state: State<'_, AppState>) log::error!("Failed to kill server process: {}", e); // Return formatted error - Ok(UnloadResult { + Ok(unloadResult { success: false, error: Some(format!("Failed to kill server process: {}", e)), }) @@ -175,10 +179,23 @@ pub async fn unload_llama_model(session_id: String, state: State<'_, AppState>) // If no process is running but client thinks there is, // still report success since the end state is what they wanted - Ok(UnloadResult { + Ok(unloadResult { success: true, error: None, }) } } +// crypto +#[allow(clippy::camel_case_variables)] +#[tauri::command] +pub fn generate_api_key(modelId: String, apiSecret: String) -> Result { + let mut mac = HmacSha256::new_from_slice(apiSecret.as_bytes()) + .map_err(|e| format!("Invalid key length: {}", e))?; + mac.update(modelId.as_bytes()); + let result = mac.finalize(); + let code_bytes = result.into_bytes(); + let hash = general_purpose::STANDARD.encode(code_bytes); + Ok(hash) +} + diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index 6f57d3cbc..a139bcd01 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -91,6 +91,7 @@ pub fn run() { // llama-cpp extension core::utils::extensions::inference_llamacpp_extension::server::load_llama_model, core::utils::extensions::inference_llamacpp_extension::server::unload_llama_model, + core::utils::extensions::inference_llamacpp_extension::server::generate_api_key, ]) .manage(AppState { app_token: Some(generate_app_token()),