Skip to content

Fix engine-side KCL websocket authentication #1816

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Joe-Spencer wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix engine-side KCL websocket authentication #1816

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
178 changes: 165 additions & 13 deletions src/context.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use std::{path::Path, str::FromStr, time::Duration};
use std::{collections::HashMap, path::Path, str::FromStr, time::Duration};

use anyhow::{Result, anyhow};
use camino::Utf8Path;
Expand All @@ -8,7 +8,7 @@ use kittycad_modeling_cmds::{
ModelingCmd, each_cmd as mcmd,
output::TakeSnapshot,
websocket::{
FailureWebSocketResponse, ModelingCmdReq, ModelingSessionData, OkWebSocketResponseData, RawFile,
ErrorCode, FailureWebSocketResponse, ModelingCmdReq, ModelingSessionData, OkWebSocketResponseData, RawFile,
SuccessWebSocketResponse, WebSocketRequest, WebSocketResponse,
},
};
Expand All @@ -28,6 +28,7 @@ type DirectWsWrite = futures::stream::SplitSink<DirectWs, WsMsg>;

const ENGINE_EXECUTION_ENV: &str = "ENGINE_EXECUTION";
const WS_RESPONSE_TIMEOUT_SECS: u64 = 600;
const MAX_TRANSIENT_AUTH_MISSING_RESPONSES: usize = 8;

pub struct Context<'a> {
pub config: &'a mut (dyn Config + Send + Sync + 'a),
Expand Down Expand Up @@ -142,9 +143,7 @@ impl<'a> Context<'a> {
}
}

/// This function returns an API client for Zoo that is based on the configured
/// user.
pub fn api_client(&self, hostname: &str) -> Result<kittycad::Client> {
fn api_client_and_token(&self, hostname: &str) -> Result<(kittycad::Client, String)> {
let (host, baseurl) = self.resolve_api_host_and_baseurl(hostname)?;

let http_client = self.http_client_builder();
Expand All @@ -158,13 +157,19 @@ impl<'a> Context<'a> {
let token = self.config.get(&host, "token")?;

// Create the client.
let mut client = kittycad::Client::new_from_reqwest(token, http_client, ws_client);
let mut client = kittycad::Client::new_from_reqwest(token.clone(), http_client, ws_client);

if baseurl != crate::DEFAULT_HOST {
client.set_base_url(&baseurl);
}

Ok(client)
Ok((client, token))
}

/// This function returns an API client for Zoo that is based on the configured
/// user.
pub fn api_client(&self, hostname: &str) -> Result<kittycad::Client> {
Ok(self.api_client_and_token(hostname)?.0)
}

pub fn raw_http_request(
Expand Down Expand Up @@ -282,8 +287,8 @@ impl<'a> Context<'a> {
&self,
hostname: &str,
settings: &kcl_lib::ExecutorSettings,
) -> Result<reqwest::Upgraded> {
let client = self.api_client(hostname)?;
) -> Result<(reqwest::Upgraded, String)> {
let (client, token) = self.api_client_and_token(hostname)?;
let pr = std::env::var("ZOO_ENGINE_PR").ok().and_then(|s| s.parse().ok());
let (ws, _headers) = client
.modeling()
Expand All @@ -306,7 +311,7 @@ impl<'a> Context<'a> {
webrtc: Some(false),
})
.await?;
Ok(ws)
Ok((ws, token))
}

/// Run this KCL on the server, then send some followup modeling commands
Expand All @@ -324,7 +329,7 @@ impl<'a> Context<'a> {
anyhow::bail!("Invalid filepath {} (must be unicode)", filepath.display());
};
let project = build_kcl_project(filepath, code)?;
let ws = self.engine_ws_with_settings(hostname, &settings).await?;
let (ws, token) = self.engine_ws_with_settings(hostname, &settings).await?;
let wsconfig = tokio_tungstenite::tungstenite::protocol::WebSocketConfig::default()
.max_message_size(Some(usize::MAX))
.max_frame_size(Some(usize::MAX));
Expand All @@ -333,6 +338,14 @@ impl<'a> Context<'a> {
let mut session_data = None;
let mut heartbeat =
tokio::time::interval(Duration::from_secs(settings.heartbeats.unwrap_or(cmd_kcl::HEARTBEATS)));
let mut auth_missing_grace = MAX_TRANSIENT_AUTH_MISSING_RESPONSES;

// Some Zoo credentials (including OAuth login tokens) must be forwarded
// in-band after the websocket upgrade. API tokens can authenticate the
// HTTP upgrade directly, but sending this request is valid for both
// credential types and keeps server-side execution consistent with the
// authentication already resolved by the CLI.
send_ws_request(&mut write, websocket_auth_request(&token)).await?;

let exec_request_id = uuid::Uuid::new_v4();
send_ws_request(
Expand All @@ -353,6 +366,9 @@ impl<'a> Context<'a> {
session_data = Some(session);
continue;
}
if take_transient_auth_token_missing(&resp, &mut auth_missing_grace) {
continue;
}

let success_resp = match resp {
WebSocketResponse::Success(success) => success,
Expand Down Expand Up @@ -868,14 +884,50 @@ fn check_server_compilation_issues(
}

async fn send_ws_request(write: &mut DirectWsWrite, request: WebSocketRequest) -> Result<()> {
let msg = serde_json::to_string(&request)?;
let msg = encode_ws_request(&request)?;
write
.send(WsMsg::Text(msg.into()))
.send(msg)
.await
.map_err(|err| anyhow!("could not send request to engine websocket: {err}"))?;
Ok(())
}

fn encode_ws_request(request: &WebSocketRequest) -> Result<WsMsg> {
if matches!(request, WebSocketRequest::ExecKclProject { .. }) {
Ok(WsMsg::Binary(rmp_serde::to_vec_named(request)?.into()))
} else {
Ok(WsMsg::Text(serde_json::to_string(request)?.into()))
}
}

fn websocket_auth_request(token: &str) -> WebSocketRequest {
let mut headers = HashMap::new();
headers.insert("Authorization".to_owned(), format!("Bearer {token}"));
WebSocketRequest::Headers { headers }
}

fn is_transient_auth_token_missing(response: &WebSocketResponse) -> bool {
matches!(
response,
WebSocketResponse::Failure(FailureWebSocketResponse {
request_id: None,
errors,
..
}) if !errors.is_empty()
&& errors
.iter()
.all(|error| error.error_code == ErrorCode::AuthTokenMissing)
)
}

fn take_transient_auth_token_missing(response: &WebSocketResponse, remaining: &mut usize) -> bool {
if *remaining == 0 || !is_transient_auth_token_missing(response) {
return false;
}
*remaining -= 1;
true
}

async fn read_ws_response_with_heartbeat(
read: &mut DirectWsRead,
write: &mut DirectWsWrite,
Expand Down Expand Up @@ -1183,6 +1235,106 @@ mod test {
}
}

#[test]
fn only_uncorrelated_auth_token_missing_is_treated_as_transient() {
use kittycad_modeling_cmds::websocket::ApiError;

let missing = ApiError {
error_code: ErrorCode::AuthTokenMissing,
message: "send authentication headers".to_owned(),
};
let invalid = ApiError {
error_code: ErrorCode::AuthTokenInvalid,
message: "invalid authentication token".to_owned(),
};

let missing_response = WebSocketResponse::failure(None, vec![missing.clone()]);
let mut remaining = MAX_TRANSIENT_AUTH_MISSING_RESPONSES;
for _ in 0..MAX_TRANSIENT_AUTH_MISSING_RESPONSES {
assert!(take_transient_auth_token_missing(&missing_response, &mut remaining));
}
assert!(!take_transient_auth_token_missing(&missing_response, &mut remaining));
assert!(!is_transient_auth_token_missing(&WebSocketResponse::failure(
Some(uuid::Uuid::new_v4()),
vec![missing.clone()],
)));
assert!(!is_transient_auth_token_missing(&WebSocketResponse::failure(
None,
vec![invalid.clone()],
)));
assert!(!is_transient_auth_token_missing(&WebSocketResponse::failure(
None,
vec![missing, invalid],
)));
assert!(!is_transient_auth_token_missing(&WebSocketResponse::failure(
None,
Vec::new(),
)));
}

#[test]
fn configured_token_becomes_text_websocket_auth_header_without_environment_override() {
let host = crate::cmd_auth::parse_host(crate::DEFAULT_HOST).unwrap().to_string();
let mut config = crate::config::new_blank_config().unwrap();
config.set(&host, "token", Some("configured-oauth-token")).unwrap();
config.set(&host, "default", Some("true")).unwrap();
let mut c = TestEnvConfig {
config: &mut config,
env: Arc::new(HashMap::new()),
};
let (io, _stdout_path, _stderr_path) = crate::iostreams::IoStreams::test();
let ctx = Context {
config: &mut c,
io,
debug: false,
override_host: None,
};

let (_client, token) = ctx.api_client_and_token("").unwrap();
let request = websocket_auth_request(&token);
let WsMsg::Text(encoded) = encode_ws_request(&request).unwrap() else {
panic!("websocket authentication must be sent as JSON text");
};

assert_eq!(
serde_json::from_str::<serde_json::Value>(encoded.as_ref()).unwrap(),
serde_json::json!({
"type": "headers",
"headers": {
"Authorization": "Bearer configured-oauth-token",
},
})
);
}

#[test]
fn exec_kcl_project_is_sent_as_named_messagepack() {
use kittycad_modeling_cmds::{
exec_kcl::{KclFile, KclProject},
shared::safe_filepath::SafeFilepath,
};

let entrypoint = SafeFilepath::validate("main.kcl").unwrap();
let project = KclProject::new(vec![KclFile::new(entrypoint.clone(), b"cube = 1".to_vec())], entrypoint);
let expected_project = project.clone();
let expected_request_id = uuid::Uuid::new_v4();
let request = WebSocketRequest::ExecKclProject {
request_id: expected_request_id,
project,
};

let WsMsg::Binary(encoded) = encode_ws_request(&request).unwrap() else {
panic!("ExecKclProject must be sent as MessagePack binary");
};
let decoded: WebSocketRequest = rmp_serde::from_slice(encoded.as_ref()).unwrap();

let WebSocketRequest::ExecKclProject { request_id, project } = decoded else {
panic!("decoded request was not ExecKclProject");
};
assert_eq!(request_id, expected_request_id);
assert_eq!(project, expected_project);
}

#[test]
fn reasoning_to_markdown_text_has_no_header() {
let md = super::reasoning_to_markdown(&kittycad::types::ReasoningMessage::Text {
Expand Down
Loading