feat: self-signed HTTPS for goosed server (#7126)
This commit is contained in:
@@ -13,6 +13,7 @@ pub async fn check_token(
|
||||
if request.uri().path() == "/status"
|
||||
|| request.uri().path() == "/mcp-ui-proxy"
|
||||
|| request.uri().path() == "/mcp-app-proxy"
|
||||
|| request.uri().path() == "/mcp-app-guest"
|
||||
{
|
||||
return Ok(next.run(request).await);
|
||||
}
|
||||
|
||||
@@ -2,11 +2,12 @@ use crate::configuration;
|
||||
use crate::state;
|
||||
use anyhow::Result;
|
||||
use axum::middleware;
|
||||
use axum_server::Handle;
|
||||
use goose_server::auth::check_token;
|
||||
use goose_server::tls::self_signed_config;
|
||||
use tower_http::cors::{Any, CorsLayer};
|
||||
use tracing::info;
|
||||
|
||||
// Graceful shutdown signal
|
||||
#[cfg(unix)]
|
||||
async fn shutdown_signal() {
|
||||
use tokio::signal::unix::{signal, SignalKind};
|
||||
@@ -53,8 +54,17 @@ pub async fn run() -> Result<()> {
|
||||
))
|
||||
.layer(cors);
|
||||
|
||||
let listener = tokio::net::TcpListener::bind(settings.socket_addr()).await?;
|
||||
info!("listening on {}", listener.local_addr()?);
|
||||
let addr = settings.socket_addr();
|
||||
let tls_setup = self_signed_config().await?;
|
||||
|
||||
let handle = Handle::new();
|
||||
let shutdown_handle = handle.clone();
|
||||
tokio::spawn(async move {
|
||||
shutdown_signal().await;
|
||||
shutdown_handle.graceful_shutdown(None);
|
||||
});
|
||||
|
||||
info!("listening on https://{}", addr);
|
||||
|
||||
let tunnel_manager = app_state.tunnel_manager.clone();
|
||||
tokio::spawn(async move {
|
||||
@@ -66,8 +76,9 @@ pub async fn run() -> Result<()> {
|
||||
gateway_manager.check_auto_start().await;
|
||||
});
|
||||
|
||||
axum::serve(listener, app)
|
||||
.with_graceful_shutdown(shutdown_signal())
|
||||
axum_server::bind_rustls(addr, tls_setup.config)
|
||||
.handle(handle)
|
||||
.serve(app.into_make_service())
|
||||
.await?;
|
||||
|
||||
if goose::otel::otlp::is_otlp_initialized() {
|
||||
|
||||
@@ -4,6 +4,7 @@ pub mod error;
|
||||
pub mod openapi;
|
||||
pub mod routes;
|
||||
pub mod state;
|
||||
pub mod tls;
|
||||
pub mod tunnel;
|
||||
|
||||
// Re-export commonly used items
|
||||
|
||||
@@ -848,6 +848,16 @@ async fn update_working_dir(
|
||||
Ok(StatusCode::OK)
|
||||
}
|
||||
|
||||
async fn ensure_extensions_loaded(state: &AppState, session_id: &str) {
|
||||
if let Some(_results) = state.take_extension_loading_task(session_id).await {
|
||||
tracing::debug!(
|
||||
"Awaited background extension loading for session {} before serving request",
|
||||
session_id
|
||||
);
|
||||
state.remove_extension_loading_task(session_id).await;
|
||||
}
|
||||
}
|
||||
|
||||
#[utoipa::path(
|
||||
post,
|
||||
path = "/agent/read_resource",
|
||||
@@ -866,6 +876,8 @@ async fn read_resource(
|
||||
) -> Result<Json<ReadResourceResponse>, StatusCode> {
|
||||
use rmcp::model::ResourceContents;
|
||||
|
||||
ensure_extensions_loaded(&state, &payload.session_id).await;
|
||||
|
||||
let agent = state
|
||||
.get_agent_for_route(payload.session_id.clone())
|
||||
.await?;
|
||||
@@ -879,7 +891,16 @@ async fn read_resource(
|
||||
CancellationToken::default(),
|
||||
)
|
||||
.await
|
||||
.map_err(|_e| StatusCode::INTERNAL_SERVER_ERROR)?;
|
||||
.map_err(|e| {
|
||||
tracing::error!(
|
||||
"read_resource failed for session={}, uri={}, extension={}: {:?}",
|
||||
payload.session_id,
|
||||
payload.uri,
|
||||
payload.extension_name,
|
||||
e
|
||||
);
|
||||
StatusCode::INTERNAL_SERVER_ERROR
|
||||
})?;
|
||||
|
||||
let content = read_result
|
||||
.contents
|
||||
@@ -936,6 +957,8 @@ async fn call_tool(
|
||||
State(state): State<Arc<AppState>>,
|
||||
Json(payload): Json<CallToolRequest>,
|
||||
) -> Result<Json<CallToolResponse>, StatusCode> {
|
||||
ensure_extensions_loaded(&state, &payload.session_id).await;
|
||||
|
||||
let agent = state
|
||||
.get_agent_for_route(payload.session_id.clone())
|
||||
.await?;
|
||||
|
||||
@@ -2,10 +2,23 @@ use axum::{
|
||||
extract::Query,
|
||||
http::{header, StatusCode},
|
||||
response::{Html, IntoResponse, Response},
|
||||
routing::get,
|
||||
Router,
|
||||
routing::{get, post},
|
||||
Json, Router,
|
||||
};
|
||||
use serde::Deserialize;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::time::Instant;
|
||||
use tokio::sync::RwLock;
|
||||
use uuid::Uuid;
|
||||
|
||||
const GUEST_HTML_TTL_SECS: u64 = 300; // 5 minutes
|
||||
const GUEST_HTML_MAX_ENTRIES: usize = 64;
|
||||
|
||||
/// In-memory store for guest HTML content.
|
||||
/// Maps nonce -> (html_content, csp_string, created_at)
|
||||
/// Entries are consumed on first read and evicted after TTL.
|
||||
type GuestHtmlStore = Arc<RwLock<HashMap<String, (String, String, Instant)>>>;
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct ProxyQuery {
|
||||
@@ -18,6 +31,22 @@ struct ProxyQuery {
|
||||
frame_domains: Option<String>,
|
||||
/// Comma-separated list of allowed base URIs (base-uri)
|
||||
base_uri_domains: Option<String>,
|
||||
/// Comma-separated list of domains for script-src (external scripts like SDKs)
|
||||
script_domains: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct GuestQuery {
|
||||
secret: String,
|
||||
nonce: String,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct StoreGuestBody {
|
||||
secret: String,
|
||||
html: String,
|
||||
/// CSP string to apply to the guest page
|
||||
csp: Option<String>,
|
||||
}
|
||||
|
||||
const MCP_APP_PROXY_HTML: &str = include_str!("templates/mcp_app_proxy.html");
|
||||
@@ -35,6 +64,7 @@ fn build_outer_csp(
|
||||
resource_domains: &[String],
|
||||
frame_domains: &[String],
|
||||
base_uri_domains: &[String],
|
||||
script_domains: &[String],
|
||||
) -> String {
|
||||
let resources = if resource_domains.is_empty() {
|
||||
String::new()
|
||||
@@ -42,16 +72,23 @@ fn build_outer_csp(
|
||||
format!(" {}", resource_domains.join(" "))
|
||||
};
|
||||
|
||||
let scripts = if script_domains.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
format!(" {}", script_domains.join(" "))
|
||||
};
|
||||
|
||||
let connections = if connect_domains.is_empty() {
|
||||
String::new()
|
||||
} else {
|
||||
format!(" {}", connect_domains.join(" "))
|
||||
};
|
||||
|
||||
// frame-src needs 'self' so the proxy can load the guest iframe from /mcp-app-guest
|
||||
let frame_src = if frame_domains.is_empty() {
|
||||
"frame-src 'none'".to_string()
|
||||
"frame-src 'self'".to_string()
|
||||
} else {
|
||||
format!("frame-src {}", frame_domains.join(" "))
|
||||
format!("frame-src 'self' {}", frame_domains.join(" "))
|
||||
};
|
||||
|
||||
let base_uris = if base_uri_domains.is_empty() {
|
||||
@@ -62,8 +99,8 @@ fn build_outer_csp(
|
||||
|
||||
format!(
|
||||
"default-src 'none'; \
|
||||
script-src 'self' 'unsafe-inline'{resources}; \
|
||||
script-src-elem 'self' 'unsafe-inline'{resources}; \
|
||||
script-src 'self' 'unsafe-inline'{resources}{scripts}; \
|
||||
script-src-elem 'self' 'unsafe-inline'{resources}{scripts}; \
|
||||
style-src 'self' 'unsafe-inline'{resources}; \
|
||||
style-src-elem 'self' 'unsafe-inline'{resources}; \
|
||||
connect-src 'self'{connections}; \
|
||||
@@ -88,6 +125,12 @@ fn parse_domains(domains: Option<&String>) -> Vec<String> {
|
||||
.unwrap_or_default()
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct AppState {
|
||||
secret_key: String,
|
||||
guest_store: GuestHtmlStore,
|
||||
}
|
||||
|
||||
#[utoipa::path(
|
||||
get,
|
||||
path = "/mcp-app-proxy",
|
||||
@@ -96,7 +139,8 @@ fn parse_domains(domains: Option<&String>) -> Vec<String> {
|
||||
("connect_domains" = Option<String>, Query, description = "Comma-separated domains for connect-src"),
|
||||
("resource_domains" = Option<String>, Query, description = "Comma-separated domains for resource loading"),
|
||||
("frame_domains" = Option<String>, Query, description = "Comma-separated origins for nested iframes (frame-src)"),
|
||||
("base_uri_domains" = Option<String>, Query, description = "Comma-separated allowed base URIs (base-uri)")
|
||||
("base_uri_domains" = Option<String>, Query, description = "Comma-separated allowed base URIs (base-uri)"),
|
||||
("script_domains" = Option<String>, Query, description = "Comma-separated domains for script-src")
|
||||
),
|
||||
responses(
|
||||
(status = 200, description = "MCP App proxy HTML page", content_type = "text/html"),
|
||||
@@ -104,10 +148,10 @@ fn parse_domains(domains: Option<&String>) -> Vec<String> {
|
||||
)
|
||||
)]
|
||||
async fn mcp_app_proxy(
|
||||
axum::extract::State(secret_key): axum::extract::State<String>,
|
||||
axum::extract::State(state): axum::extract::State<AppState>,
|
||||
Query(params): Query<ProxyQuery>,
|
||||
) -> Response {
|
||||
if params.secret != secret_key {
|
||||
if params.secret != state.secret_key {
|
||||
return (StatusCode::UNAUTHORIZED, "Unauthorized").into_response();
|
||||
}
|
||||
|
||||
@@ -116,6 +160,7 @@ async fn mcp_app_proxy(
|
||||
let resource_domains = parse_domains(params.resource_domains.as_ref());
|
||||
let frame_domains = parse_domains(params.frame_domains.as_ref());
|
||||
let base_uri_domains = parse_domains(params.base_uri_domains.as_ref());
|
||||
let script_domains = parse_domains(params.script_domains.as_ref());
|
||||
|
||||
// Build the outer CSP based on declared domains
|
||||
let csp = build_outer_csp(
|
||||
@@ -123,6 +168,7 @@ async fn mcp_app_proxy(
|
||||
&resource_domains,
|
||||
&frame_domains,
|
||||
&base_uri_domains,
|
||||
&script_domains,
|
||||
);
|
||||
|
||||
// Replace the CSP placeholder in the HTML template
|
||||
@@ -141,8 +187,98 @@ async fn mcp_app_proxy(
|
||||
.into_response()
|
||||
}
|
||||
|
||||
/// Store guest HTML and return a nonce for retrieval.
|
||||
/// The proxy page calls this via fetch, then sets the guest iframe src to /mcp-app-guest?nonce=...
|
||||
async fn store_guest_html(
|
||||
axum::extract::State(state): axum::extract::State<AppState>,
|
||||
Json(body): Json<StoreGuestBody>,
|
||||
) -> Response {
|
||||
if body.secret != state.secret_key {
|
||||
return (StatusCode::UNAUTHORIZED, "Unauthorized").into_response();
|
||||
}
|
||||
|
||||
let nonce = Uuid::new_v4().to_string();
|
||||
let csp = body.csp.unwrap_or_default();
|
||||
|
||||
{
|
||||
let mut store = state.guest_store.write().await;
|
||||
|
||||
// Evict expired entries
|
||||
let cutoff = Instant::now() - std::time::Duration::from_secs(GUEST_HTML_TTL_SECS);
|
||||
store.retain(|_, (_, _, created)| *created > cutoff);
|
||||
|
||||
// If still at capacity, drop the oldest entry
|
||||
if store.len() >= GUEST_HTML_MAX_ENTRIES {
|
||||
if let Some(oldest_key) = store
|
||||
.iter()
|
||||
.min_by_key(|(_, (_, _, created))| *created)
|
||||
.map(|(k, _)| k.clone())
|
||||
{
|
||||
store.remove(&oldest_key);
|
||||
}
|
||||
}
|
||||
|
||||
store.insert(nonce.clone(), (body.html, csp, Instant::now()));
|
||||
}
|
||||
|
||||
(
|
||||
StatusCode::OK,
|
||||
[(header::CONTENT_TYPE, "application/json")],
|
||||
format!(r#"{{"nonce":"{}"}}"#, nonce),
|
||||
)
|
||||
.into_response()
|
||||
}
|
||||
|
||||
/// Serve stored guest HTML with a real HTTPS URL.
|
||||
/// This gives the guest iframe `window.location.protocol === "https:"`,
|
||||
/// which is required by SDKs like Square Web Payments that check for secure context.
|
||||
async fn serve_guest_html(
|
||||
axum::extract::State(state): axum::extract::State<AppState>,
|
||||
Query(params): Query<GuestQuery>,
|
||||
) -> Response {
|
||||
if params.secret != state.secret_key {
|
||||
return (StatusCode::UNAUTHORIZED, "Unauthorized").into_response();
|
||||
}
|
||||
|
||||
// Consume the entry (one-time use)
|
||||
let entry = {
|
||||
let mut store = state.guest_store.write().await;
|
||||
store.remove(¶ms.nonce)
|
||||
};
|
||||
|
||||
match entry {
|
||||
Some((html, csp, _created)) => {
|
||||
let mut response = Html(html).into_response();
|
||||
let headers = response.headers_mut();
|
||||
// Use strict-origin so third-party SDKs (e.g. Square Web Payments)
|
||||
// receive the origin in their requests, which they need for auth.
|
||||
// no-referrer would cause 401s from SDK servers.
|
||||
headers.insert(
|
||||
header::HeaderName::from_static("referrer-policy"),
|
||||
"strict-origin".parse().unwrap(),
|
||||
);
|
||||
if !csp.is_empty() {
|
||||
headers.insert(header::CONTENT_SECURITY_POLICY, csp.parse().unwrap());
|
||||
}
|
||||
response
|
||||
}
|
||||
None => (
|
||||
StatusCode::NOT_FOUND,
|
||||
"Guest content not found or already consumed",
|
||||
)
|
||||
.into_response(),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn routes(secret_key: String) -> Router {
|
||||
let state = AppState {
|
||||
secret_key,
|
||||
guest_store: Arc::new(RwLock::new(HashMap::new())),
|
||||
};
|
||||
|
||||
Router::new()
|
||||
.route("/mcp-app-proxy", get(mcp_app_proxy))
|
||||
.with_state(secret_key)
|
||||
.route("/mcp-app-guest", get(serve_guest_html))
|
||||
.route("/mcp-app-guest", post(store_guest_html))
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
@@ -34,7 +34,42 @@
|
||||
|
||||
let guestIframe = null;
|
||||
|
||||
function createGuestIframe(html, permissions) {
|
||||
/**
|
||||
* Extract the secret and base URL from the current page URL.
|
||||
*/
|
||||
function getProxyParams() {
|
||||
var params = new URLSearchParams(window.location.search);
|
||||
return {
|
||||
secret: params.get('secret') || '',
|
||||
baseUrl: window.location.origin
|
||||
};
|
||||
}
|
||||
|
||||
/**
|
||||
* Store guest HTML on the server and get a nonce for retrieval.
|
||||
* This allows the guest iframe to load from a real HTTPS URL
|
||||
* instead of srcdoc (which has about: protocol).
|
||||
*/
|
||||
async function storeGuestHtml(html) {
|
||||
var proxyParams = getProxyParams();
|
||||
var response = await fetch(proxyParams.baseUrl + '/mcp-app-guest', {
|
||||
method: 'POST',
|
||||
headers: { 'Content-Type': 'application/json' },
|
||||
body: JSON.stringify({
|
||||
secret: proxyParams.secret,
|
||||
html: html
|
||||
})
|
||||
});
|
||||
|
||||
if (!response.ok) {
|
||||
throw new Error('Failed to store guest HTML: ' + response.status);
|
||||
}
|
||||
|
||||
var data = await response.json();
|
||||
return data.nonce;
|
||||
}
|
||||
|
||||
async function createGuestIframe(html, permissions) {
|
||||
if (guestIframe) {
|
||||
guestIframe.remove();
|
||||
}
|
||||
@@ -43,7 +78,7 @@
|
||||
|
||||
// Sandbox permissions for the Guest UI
|
||||
// allow-scripts: needed for the app to run
|
||||
// allow-same-origin: needed for localStorage, cookies, etc.
|
||||
// allow-same-origin: needed for localStorage, cookies, and real URL context
|
||||
// allow-forms: needed if the app has forms
|
||||
guestIframe.setAttribute('sandbox', 'allow-scripts allow-same-origin allow-forms');
|
||||
|
||||
@@ -58,9 +93,22 @@
|
||||
guestIframe.setAttribute('allow', allowList.join('; '));
|
||||
}
|
||||
|
||||
guestIframe.srcdoc = html;
|
||||
guestIframe.style.cssText = 'width:100%; height:100%; border:none;';
|
||||
|
||||
// Store the HTML server-side and load via real URL.
|
||||
// This gives the guest iframe a real https:// URL instead of about:srcdoc,
|
||||
// which is required by SDKs (like Square Web Payments) that check
|
||||
// window.location.protocol for secure context verification.
|
||||
try {
|
||||
var nonce = await storeGuestHtml(html);
|
||||
var proxyParams = getProxyParams();
|
||||
guestIframe.src = proxyParams.baseUrl + '/mcp-app-guest?secret=' + encodeURIComponent(proxyParams.secret) + '&nonce=' + encodeURIComponent(nonce);
|
||||
} catch (e) {
|
||||
// Fallback to srcdoc if the store endpoint is not available
|
||||
console.warn('Failed to use /mcp-app-guest endpoint, falling back to srcdoc:', e);
|
||||
guestIframe.srcdoc = html;
|
||||
}
|
||||
|
||||
document
|
||||
.body
|
||||
.appendChild(guestIframe);
|
||||
|
||||
@@ -0,0 +1,52 @@
|
||||
use anyhow::Result;
|
||||
use aws_lc_rs::digest;
|
||||
use axum_server::tls_rustls::RustlsConfig;
|
||||
use rcgen::{CertificateParams, DnType, KeyPair, SanType};
|
||||
|
||||
pub struct TlsSetup {
|
||||
pub config: RustlsConfig,
|
||||
pub fingerprint: String,
|
||||
}
|
||||
|
||||
/// Generate a self-signed TLS certificate for localhost (127.0.0.1) and
|
||||
/// return a [`TlsSetup`] containing the rustls config and the SHA-256
|
||||
/// fingerprint of the generated certificate (colon-separated hex).
|
||||
///
|
||||
/// The fingerprint is printed to stdout so the parent process (e.g. Electron)
|
||||
/// can pin it and reject connections from any other certificate.
|
||||
pub async fn self_signed_config() -> Result<TlsSetup> {
|
||||
let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
|
||||
|
||||
let mut params = CertificateParams::default();
|
||||
params
|
||||
.distinguished_name
|
||||
.push(DnType::CommonName, "goosed localhost");
|
||||
params.subject_alt_names = vec![
|
||||
SanType::IpAddress(std::net::IpAddr::V4(std::net::Ipv4Addr::LOCALHOST)),
|
||||
SanType::DnsName("localhost".try_into()?),
|
||||
];
|
||||
|
||||
let key_pair = KeyPair::generate()?;
|
||||
let cert = params.self_signed(&key_pair)?;
|
||||
|
||||
let cert_der = cert.der();
|
||||
let sha256 = digest::digest(&digest::SHA256, cert_der);
|
||||
let fingerprint = sha256
|
||||
.as_ref()
|
||||
.iter()
|
||||
.map(|b| format!("{b:02X}"))
|
||||
.collect::<Vec<_>>()
|
||||
.join(":");
|
||||
|
||||
println!("GOOSED_CERT_FINGERPRINT={fingerprint}");
|
||||
|
||||
let cert_pem = cert.pem();
|
||||
let key_pem = key_pair.serialize_pem();
|
||||
|
||||
let config = RustlsConfig::from_pem(cert_pem.into_bytes(), key_pem.into_bytes()).await?;
|
||||
|
||||
Ok(TlsSetup {
|
||||
config,
|
||||
fingerprint,
|
||||
})
|
||||
}
|
||||
@@ -13,6 +13,15 @@ use tokio_tungstenite::{connect_async, tungstenite::Message};
|
||||
use tracing::{error, info, warn};
|
||||
use url::Url;
|
||||
|
||||
/// Shared state for proxying tunnel requests to the local goosed server.
|
||||
#[derive(Clone)]
|
||||
struct ProxyContext {
|
||||
port: u16,
|
||||
tunnel_secret: String,
|
||||
server_secret: String,
|
||||
http_client: reqwest::Client,
|
||||
}
|
||||
|
||||
/// Constant-time comparison using hash to prevent timing attacks
|
||||
fn secure_compare(a: &str, b: &str) -> bool {
|
||||
use std::collections::hash_map::DefaultHasher;
|
||||
@@ -253,38 +262,42 @@ async fn handle_chunked_response(
|
||||
|
||||
async fn handle_request(
|
||||
message: TunnelMessage,
|
||||
port: u16,
|
||||
ctx: ProxyContext,
|
||||
ws_tx: WebSocketSender,
|
||||
tunnel_secret: String,
|
||||
server_secret: String,
|
||||
scheme: &str,
|
||||
) -> Result<()> {
|
||||
let request_id = message.request_id.clone();
|
||||
let client = &ctx.http_client;
|
||||
|
||||
let client = reqwest::Client::new();
|
||||
let url = format!("http://127.0.0.1:{}{}", port, message.path);
|
||||
let url = format!("{}://127.0.0.1:{}{}", scheme, ctx.port, message.path);
|
||||
|
||||
let request_builder =
|
||||
match validate_and_build_request(&client, &url, &message, &tunnel_secret, &server_secret) {
|
||||
Ok(builder) => builder,
|
||||
Err(e) => {
|
||||
error!("✗ Authentication error [{}]: {}", request_id, e);
|
||||
let error_response = TunnelResponse {
|
||||
request_id,
|
||||
status: 401,
|
||||
headers: None,
|
||||
body: None,
|
||||
error: Some(e.to_string()),
|
||||
chunk_index: None,
|
||||
total_chunks: None,
|
||||
is_chunked: false,
|
||||
is_streaming: false,
|
||||
is_first_chunk: false,
|
||||
is_last_chunk: false,
|
||||
};
|
||||
send_response(ws_tx, error_response).await?;
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
let request_builder = match validate_and_build_request(
|
||||
client,
|
||||
&url,
|
||||
&message,
|
||||
&ctx.tunnel_secret,
|
||||
&ctx.server_secret,
|
||||
) {
|
||||
Ok(builder) => builder,
|
||||
Err(e) => {
|
||||
error!("✗ Authentication error [{}]: {}", request_id, e);
|
||||
let error_response = TunnelResponse {
|
||||
request_id,
|
||||
status: 401,
|
||||
headers: None,
|
||||
body: None,
|
||||
error: Some(e.to_string()),
|
||||
chunk_index: None,
|
||||
total_chunks: None,
|
||||
is_chunked: false,
|
||||
is_streaming: false,
|
||||
is_first_chunk: false,
|
||||
is_last_chunk: false,
|
||||
};
|
||||
send_response(ws_tx, error_response).await?;
|
||||
return Ok(());
|
||||
}
|
||||
};
|
||||
|
||||
let response = match request_builder.send().await {
|
||||
Ok(resp) => resp,
|
||||
@@ -399,11 +412,10 @@ async fn handle_websocket_messages(
|
||||
>,
|
||||
>,
|
||||
ws_tx: WebSocketSender,
|
||||
port: u16,
|
||||
tunnel_secret: String,
|
||||
server_secret: String,
|
||||
ctx: ProxyContext,
|
||||
last_activity: Arc<RwLock<Instant>>,
|
||||
active_tasks: Arc<RwLock<Vec<JoinHandle<()>>>>,
|
||||
scheme: String,
|
||||
) {
|
||||
while let Some(msg) = read.next().await {
|
||||
match msg {
|
||||
@@ -413,17 +425,12 @@ async fn handle_websocket_messages(
|
||||
match serde_json::from_str::<TunnelMessage>(&text) {
|
||||
Ok(tunnel_msg) => {
|
||||
let ws_tx_clone = ws_tx.clone();
|
||||
let tunnel_secret_clone = tunnel_secret.clone();
|
||||
let server_secret_clone = server_secret.clone();
|
||||
let ctx_clone = ctx.clone();
|
||||
let scheme_clone = scheme.clone();
|
||||
let task = tokio::spawn(async move {
|
||||
if let Err(e) = handle_request(
|
||||
tunnel_msg,
|
||||
port,
|
||||
ws_tx_clone,
|
||||
tunnel_secret_clone,
|
||||
server_secret_clone,
|
||||
)
|
||||
.await
|
||||
if let Err(e) =
|
||||
handle_request(tunnel_msg, ctx_clone, ws_tx_clone, &scheme_clone)
|
||||
.await
|
||||
{
|
||||
error!("Error handling request: {}", e);
|
||||
}
|
||||
@@ -475,9 +482,10 @@ async fn run_single_connection(
|
||||
agent_id: String,
|
||||
tunnel_secret: String,
|
||||
server_secret: String,
|
||||
scheme: String,
|
||||
restart_tx: mpsc::Sender<()>,
|
||||
) {
|
||||
let _ = rustls::crypto::ring::default_provider().install_default();
|
||||
let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
|
||||
|
||||
let worker_url = get_worker_url();
|
||||
let ws_url = worker_url
|
||||
@@ -514,10 +522,25 @@ async fn run_single_connection(
|
||||
};
|
||||
|
||||
info!("✓ Connected as agent: {}", agent_id);
|
||||
info!("✓ Proxying to: http://127.0.0.1:{}", port);
|
||||
info!("✓ Proxying to: {}://127.0.0.1:{}", scheme, port);
|
||||
let public_url = format!("{}/tunnel/{}", worker_url, agent_id);
|
||||
info!("✓ Public URL: {}", public_url);
|
||||
|
||||
let mut client_builder = reqwest::Client::builder();
|
||||
if scheme == "https" {
|
||||
client_builder = client_builder.danger_accept_invalid_certs(true);
|
||||
}
|
||||
let http_client = client_builder
|
||||
.build()
|
||||
.expect("failed to build reqwest client");
|
||||
|
||||
let ctx = ProxyContext {
|
||||
port,
|
||||
tunnel_secret,
|
||||
server_secret,
|
||||
http_client,
|
||||
};
|
||||
|
||||
let (write, read) = ws_stream.split();
|
||||
let ws_tx: WebSocketSender = Arc::new(RwLock::new(Some(write)));
|
||||
let last_activity = Arc::new(RwLock::new(Instant::now()));
|
||||
@@ -545,11 +568,10 @@ async fn run_single_connection(
|
||||
_ = handle_websocket_messages(
|
||||
read,
|
||||
ws_tx.clone(),
|
||||
port,
|
||||
tunnel_secret.clone(),
|
||||
server_secret.clone(),
|
||||
ctx,
|
||||
last_activity,
|
||||
active_tasks.clone()
|
||||
active_tasks.clone(),
|
||||
scheme,
|
||||
) => {
|
||||
info!("✗ Connection ended");
|
||||
}
|
||||
@@ -565,6 +587,7 @@ pub async fn start(
|
||||
tunnel_secret: String,
|
||||
server_secret: String,
|
||||
agent_id: String,
|
||||
scheme: &str,
|
||||
handle: Arc<RwLock<Option<tokio::task::JoinHandle<()>>>>,
|
||||
restart_tx: mpsc::Sender<()>,
|
||||
) -> Result<TunnelInfo> {
|
||||
@@ -573,6 +596,7 @@ pub async fn start(
|
||||
let agent_id_clone = agent_id.clone();
|
||||
let tunnel_secret_clone = tunnel_secret.clone();
|
||||
let server_secret_clone = server_secret;
|
||||
let scheme = scheme.to_string();
|
||||
|
||||
let task = tokio::spawn(async move {
|
||||
run_single_connection(
|
||||
@@ -580,6 +604,7 @@ pub async fn start(
|
||||
agent_id_clone,
|
||||
tunnel_secret_clone,
|
||||
server_secret_clone,
|
||||
scheme,
|
||||
restart_tx,
|
||||
)
|
||||
.await;
|
||||
|
||||
@@ -88,6 +88,7 @@ async fn test_tunnel_end_to_end() {
|
||||
tunnel_secret.clone(),
|
||||
server_secret.clone(),
|
||||
agent_id.clone(),
|
||||
"http",
|
||||
handle.clone(),
|
||||
restart_tx,
|
||||
)
|
||||
@@ -138,6 +139,7 @@ async fn test_tunnel_post_request() {
|
||||
tunnel_secret.clone(),
|
||||
server_secret.clone(),
|
||||
agent_id.clone(),
|
||||
"http",
|
||||
handle.clone(),
|
||||
restart_tx,
|
||||
)
|
||||
|
||||
@@ -229,6 +229,7 @@ impl TunnelManager {
|
||||
tunnel_secret,
|
||||
server_secret,
|
||||
agent_id,
|
||||
"https",
|
||||
self.lapstone_handle.clone(),
|
||||
restart_tx,
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user