tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

refactor: use the verify_secret_key util for all api handlers (#2284)

Browse Source
This commit is contained in:
Yingjie He
2025-04-21 13:27:38 -07:00
committed by GitHub
parent 1f70c34315
commit 8247fd6c1d
7 changed files with 40 additions and 142 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose-server/src/routes/agent.rs
+4 -26
View File
@@ -1,3 +1,4 @@
use super::utils::verify_secret_key;
use crate::state::AppState;
use axum::{
extract::{Query, State},
@@ -85,15 +86,7 @@ async fn extend_prompt(
headers: HeaderMap,
Json(payload): Json<ExtendPromptRequest>,
) -> Result<Json<ExtendPromptResponse>, StatusCode> {
// Verify secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
verify_secret_key(&headers, &state)?;
let mut agent = state.agent.write().await;
if let Some(ref mut agent) = *agent {
@@ -110,15 +103,7 @@ async fn create_agent(
headers: HeaderMap,
Json(payload): Json<CreateAgentRequest>,
) -> Result<Json<CreateAgentResponse>, StatusCode> {
// Verify secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
verify_secret_key(&headers, &state)?;
// Set the environment variable for the model if provided
if let Some(model) = &payload.model {
@@ -187,14 +172,7 @@ async fn get_tools(
headers: HeaderMap,
Query(query): Query<GetToolsQuery>,
) -> Result<Json<Vec<ToolInfo>>, StatusCode> {
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
verify_secret_key(&headers, &state)?;
let config = Config::global();
let goose_mode = config.get_param("GOOSE_MODE").unwrap_or("auto".to_string());
crates/goose-server/src/routes/config_management.rs
+3 -16
View File
@@ -1,3 +1,4 @@
use super::utils::verify_secret_key;
use crate::routes::utils::check_provider_configured;
use crate::state::AppState;
use axum::{
@@ -5,6 +6,7 @@ use axum::{
routing::{delete, get, post},
Json, Router,
};
use etcetera::{choose_app_strategy, AppStrategy, AppStrategyArgs};
use goose::config::Config;
use goose::config::{extensions::name_to_key, PermissionManager};
use goose::config::{ExtensionConfigManager, ExtensionEntry};
@@ -12,26 +14,13 @@ use goose::providers::base::ProviderMetadata;
use goose::providers::providers as get_providers;
use goose::{agents::ExtensionConfig, config::permission::PermissionLevel};
use http::{HeaderMap, StatusCode};
use once_cell::sync::Lazy;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use serde_yaml;
use std::collections::HashMap;
use utoipa::ToSchema;
fn verify_secret_key(headers: &HeaderMap, state: &AppState) -> Result<StatusCode, StatusCode> {
// Verify secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
Err(StatusCode::UNAUTHORIZED)
} else {
Ok(StatusCode::OK)
}
}
#[derive(Serialize, ToSchema)]
pub struct ExtensionResponse {
pub extensions: Vec<ExtensionEntry>,
@@ -431,8 +420,6 @@ pub async fn upsert_permissions(
Ok(Json("Permissions updated successfully".to_string()))
}
use etcetera::{choose_app_strategy, AppStrategy, AppStrategyArgs};
use once_cell::sync::Lazy;
pub static APP_STRATEGY: Lazy<AppStrategyArgs> = Lazy::new(|| AppStrategyArgs {
top_level_domain: "Block".to_string(),
author: "Block".to_string(),
crates/goose-server/src/routes/configs.rs
+4 -27
View File
@@ -1,3 +1,4 @@
use super::utils::verify_secret_key;
use crate::state::AppState;
use axum::{
extract::{Query, State},
@@ -29,15 +30,7 @@ async fn store_config(
headers: HeaderMap,
Json(request): Json<ConfigRequest>,
) -> Result<Json<ConfigResponse>, StatusCode> {
// Verify secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
verify_secret_key(&headers, &state)?;
let config = Config::global();
let result = if request.is_secret {
@@ -159,15 +152,7 @@ pub async fn get_config(
headers: HeaderMap,
Query(query): Query<GetConfigQuery>,
) -> Result<Json<GetConfigResponse>, StatusCode> {
// Verify secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
verify_secret_key(&headers, &state)?;
// Fetch the configuration value. Right now we don't allow get a secret.
let config = Config::global();
@@ -193,15 +178,7 @@ async fn delete_config(
headers: HeaderMap,
Json(request): Json<DeleteConfigRequest>,
) -> Result<StatusCode, StatusCode> {
// Verify secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
verify_secret_key(&headers, &state)?;
// Attempt to delete the key
let config = Config::global();
crates/goose-server/src/routes/extension.rs
+4 -18
View File
@@ -2,6 +2,7 @@ use std::env;
use std::path::Path;
use std::sync::OnceLock;
use super::utils::verify_secret_key;
use crate::state::AppState;
use axum::{extract::State, routing::post, Json, Router};
use goose::agents::{extension::Envs, ExtensionConfig};
@@ -82,6 +83,8 @@ async fn add_extension(
headers: HeaderMap,
raw: axum::extract::Json<serde_json::Value>,
) -> Result<Json<ExtensionResponse>, StatusCode> {
verify_secret_key(&headers, &state)?;
// Log the raw request for debugging
tracing::info!(
"Received extension request: {}",
@@ -100,15 +103,6 @@ async fn add_extension(
return Err(StatusCode::UNPROCESSABLE_ENTITY);
}
};
// Verify the presence and validity of the secret key.
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
// If this is a Stdio extension that uses npx, check for Node.js installation
#[cfg(target_os = "windows")]
@@ -264,15 +258,7 @@ async fn remove_extension(
headers: HeaderMap,
Json(name): Json<String>,
) -> Result<Json<ExtensionResponse>, StatusCode> {
// Verify the presence and validity of the secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
verify_secret_key(&headers, &state)?;
// Acquire a lock on the agent and attempt to remove the extension
let mut agent = state.agent.write().await;
crates/goose-server/src/routes/reply.rs
+6 -37
View File
@@ -1,3 +1,4 @@
use super::utils::verify_secret_key;
use crate::state::AppState;
use axum::{
extract::State,
@@ -104,15 +105,7 @@ async fn handler(
headers: HeaderMap,
Json(request): Json<ChatRequest>,
) -> Result<SseResponse, StatusCode> {
// Verify secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
verify_secret_key(&headers, &state)?;
// Create channel for streaming
let (tx, rx) = mpsc::channel(100);
@@ -273,15 +266,7 @@ async fn ask_handler(
headers: HeaderMap,
Json(request): Json<AskRequest>,
) -> Result<Json<AskResponse>, StatusCode> {
// Verify secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
verify_secret_key(&headers, &state)?;
let session_working_dir = request.session_working_dir;
@@ -393,15 +378,7 @@ pub async fn confirm_permission(
headers: HeaderMap,
Json(request): Json<PermissionConfirmationRequest>,
) -> Result<Json<Value>, StatusCode> {
// Verify secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
verify_secret_key(&headers, &state)?;
let agent = state.agent.clone();
let agent = agent.read().await;
@@ -437,6 +414,8 @@ async fn submit_tool_result(
headers: HeaderMap,
raw: axum::extract::Json<serde_json::Value>,
) -> Result<Json<Value>, StatusCode> {
verify_secret_key(&headers, &state)?;
// Log the raw request for debugging
tracing::info!(
"Received tool result request: {}",
@@ -456,16 +435,6 @@ async fn submit_tool_result(
}
};
// Verify secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
let agent = state.agent.read().await;
let agent = agent.as_ref().ok_or(StatusCode::NOT_FOUND)?;
agent.handle_tool_result(payload.id, payload.result).await;
crates/goose-server/src/routes/session.rs
+3 -18
View File
@@ -1,3 +1,4 @@
use super::utils::verify_secret_key;
use crate::state::AppState;
use axum::{
extract::{Path, State},
@@ -27,15 +28,7 @@ async fn list_sessions(
State(state): State<AppState>,
headers: HeaderMap,
) -> Result<Json<SessionListResponse>, StatusCode> {
// Verify secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
verify_secret_key(&headers, &state)?;
let sessions = get_session_info().map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?;
@@ -48,15 +41,7 @@ async fn get_session_history(
headers: HeaderMap,
Path(session_id): Path<String>,
) -> Result<Json<SessionHistoryResponse>, StatusCode> {
// Verify secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
return Err(StatusCode::UNAUTHORIZED);
}
verify_secret_key(&headers, &state)?;
let session_path = session::get_path(session::Identifier::Name(session_id.clone()));
crates/goose-server/src/routes/utils.rs
+16
View File
@@ -1,5 +1,7 @@
use crate::state::AppState;
use goose::config::Config;
use goose::providers::base::{ConfigKey, ProviderMetadata};
use http::{HeaderMap, StatusCode};
use serde::{Deserialize, Serialize};
use std::env;
use std::error::Error;
@@ -21,6 +23,20 @@ pub struct KeyInfo {
pub value: Option<String>, // Only populated for non-secret keys that are set
}
pub fn verify_secret_key(headers: &HeaderMap, state: &AppState) -> Result<StatusCode, StatusCode> {
// Verify secret key
let secret_key = headers
.get("X-Secret-Key")
.and_then(|value| value.to_str().ok())
.ok_or(StatusCode::UNAUTHORIZED)?;
if secret_key != state.secret_key {
Err(StatusCode::UNAUTHORIZED)
} else {
Ok(StatusCode::OK)
}
}
/// Inspects a configuration key to determine if it's set, its location, and value (for non-secret keys)
#[allow(dead_code)]
pub fn inspect_key(key_name: &str, is_secret: bool) -> Result<KeyInfo, Box<dyn Error>> {