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.
This commit is contained in:
Akarshan Biswas 2025-05-29 12:02:55 +05:30 committed by Louis
parent d6edb1e944
commit da23673a44
No known key found for this signature in database
GPG Key ID: 44FA9F4D33C37DE2
7 changed files with 97 additions and 71 deletions

View File

@ -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

View File

@ -101,6 +101,7 @@ export type listResult = modelInfo[]
// 3. /load
export interface loadOptions {
modelId: string
modelPath: string
port?: number
}

View File

@ -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<string> {
const hash = await invoke<string>('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) {

View File

@ -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"

View File

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

View File

@ -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<Sha256>;
// 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<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
@ -35,18 +36,19 @@ impl serde::Serialize for ServerError {
}
}
type ServerResult<T> = Result<T, ServerError>;
type ServerResult<T> = Result<T, serverError>;
#[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<String>,
}
@ -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<String>, // Arguments from the frontend
) -> ServerResult<SessionInfo> {
) -> ServerResult<sessionInfo> {
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<UnloadResult> {
pub async fn unload_llama_model(session_id: String, state: State<'_, AppState>) -> ServerResult<unloadResult> {
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<String, String> {
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)
}

View File

@ -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()),