/* * Author: Ruben Fiszel * Copyright: Windmill Labs, Inc 2022 * This file and its contents are licensed under the AGPLv3 License. * Please see the included NOTICE for copyright information and * LICENSE-AGPL for a copy of the license. */ use crate::auth::is_devops_email; use crate::ee_oss::LICENSE_KEY_ID; #[cfg(feature = "enterprise")] use crate::ee_oss::{send_critical_alert, CriticalAlertKind}; use crate::error::{to_anyhow, Error, Result}; use crate::global_settings::UNIQUE_ID_SETTING; use crate::DB; use anyhow::Context; use gethostname::gethostname; use git_version::git_version; use chrono::Utc; use croner::Cron; use itertools::Itertools; use rand::{distr::Alphanumeric, rng, Rng}; use reqwest::Client; use semver::Version; use serde::{de::Error as SerdeDeserializerError, Deserialize, Deserializer, Serialize}; use sha2::{Digest, Sha256}; use sqlx::{Pool, Postgres}; use std::borrow::Cow; use std::fmt::Display; use std::{fs::DirBuilder as SyncDirBuilder, str::FromStr}; use tokio::fs::DirBuilder as AsyncDirBuilder; use url::Url; pub const MAX_PER_PAGE: usize = 10000; pub const DEFAULT_PER_PAGE: usize = 1000; pub const GIT_VERSION: &str = git_version!(args = ["--tag", "--always"], fallback = "unknown-version"); pub const AGENT_JWT_PREFIX: &str = "jwt_agent_"; pub const WORKER_NAME_PREFIX: &str = "wk"; pub const AGENT_WORKER_NAME_PREFIX: &str = "ag"; use crate::CRITICAL_ALERT_MUTE_UI_ENABLED; use std::panic::{self, AssertUnwindSafe, Location}; use std::sync::atomic::Ordering; use crate::worker::CLOUD_HOSTED; lazy_static::lazy_static! { pub static ref HTTP_CLIENT: Client = reqwest::ClientBuilder::new() .user_agent("windmill/beta") .timeout(std::time::Duration::from_secs(20)) .connect_timeout(std::time::Duration::from_secs(10)) .build().unwrap(); pub static ref GIT_SEM_VERSION: Version = Version::parse( if GIT_VERSION.starts_with('v') { &GIT_VERSION[1..] } else { GIT_VERSION } ).unwrap_or(Version::new(0, 1, 0)); pub static ref HOSTNAME :String = std::env::var("FORCE_HOSTNAME").unwrap_or_else(|_| { gethostname() .to_str() .map(|x| x.to_string()) .unwrap_or_else(|| rd_string(5)) }); pub static ref MODE_AND_ADDONS: ModeAndAddons = { let mut search_addon = false; let mode = std::env::var("MODE") .map(|x| x.to_lowercase()) .map(|x| { if &x == "server" { println!("Binary is in 'server' mode"); Mode::Server } else if &x == "worker" { tracing::info!("Binary is in 'worker' mode"); #[cfg(windows)] { println!("It is highly recommended to use the agent mode instead on windows (MODE=agent) and to pass a BASE_INTERNAL_URL"); } Mode::Worker } else if &x == "agent" { println!("Binary is in 'agent' mode with BASE_INTERNAL_URL={}", std::env::var("BASE_INTERNAL_URL").unwrap_or_default()); if std::env::var("BASE_INTERNAL_URL").is_err() { panic!("BASE_INTERNAL_URL is required in agent mode") } if std::env::var("AGENT_TOKEN").is_err() { println!("AGENT_TOKEN is not passed. This is required for the agent to work and contains the JWT to authenticate with the server.") } #[cfg(not(feature = "enterprise"))] { panic!("Agent mode is only available in the EE, ignoring..."); } #[cfg(feature = "enterprise")] Mode::Agent } else if &x == "indexer" { tracing::info!("Binary is in 'indexer' mode"); #[cfg(not(feature = "tantivy"))] { eprintln!("Cannot start the indexer because tantivy is not included in this binary/image. Make sure you are using the EE image if you want to access the full text search features."); panic!("Indexer mode requires compiling with the tantivy feature flag."); } #[cfg(feature = "tantivy")] Mode::Indexer } else if &x == "standalone+search"{ search_addon = true; println!("Binary is in 'standalone' mode with search enabled"); Mode::Standalone } else if &x == "mcp" { println!("Binary is in 'mcp' mode"); Mode::MCP } else { if &x != "standalone" { eprintln!("mode not recognized, defaulting to standalone: {x}"); } else { println!("Binary is in 'standalone' mode"); } Mode::Standalone } }) .unwrap_or_else(|_| { tracing::info!("Mode not specified, defaulting to standalone"); Mode::Standalone }); ModeAndAddons { indexer: search_addon, mode, } }; } lazy_static::lazy_static! { pub static ref AGENT_TOKEN: String = std::env::var("AGENT_TOKEN").unwrap_or_default(); } #[derive(Clone)] pub struct ModeAndAddons { pub indexer: bool, pub mode: Mode, } #[derive(Deserialize, Clone)] pub struct Pagination { pub page: Option, pub per_page: Option, } #[derive(Debug, Serialize, Deserialize)] pub struct StripPath(pub String); impl StripPath { pub fn to_path(&self) -> &str { if self.0.starts_with('/') { self.0.strip_prefix('/').unwrap() } else { &self.0 } } } pub fn require_admin(is_admin: bool, username: &str) -> Result<()> { if !is_admin { Err(Error::RequireAdmin(username.to_string())) } else { Ok(()) } } pub async fn require_admin_or_devops( is_admin: bool, username: &str, email: &str, db: &DB, ) -> Result<()> { if !is_admin { if !is_devops_email(db, email).await? { return Err(Error::RequireAdmin(username.to_string())); } } Ok(()) } fn instance_name(hostname: &str) -> String { hostname .replace(" ", "") .split("-") .last() .unwrap() .to_ascii_lowercase() .to_string() } const DEFAULT_WORKER_SUFFIX_LEN: usize = 5; pub const SSH_AGENT_WORKER_SUFFIX: &'static str = "/ssh"; pub fn create_worker_suffix(hostname: &str, rd_string_len: usize) -> String { let wk_suffix = format!("{}-{}", instance_name(hostname), rd_string(rd_string_len)); wk_suffix } pub fn create_default_worker_suffix(hostname: &str) -> String { create_worker_suffix(hostname, DEFAULT_WORKER_SUFFIX_LEN) } pub fn worker_name_with_suffix(is_agent: bool, worker_group: &str, suffix: &str) -> String { if is_agent { format!("{}-{}-{}", AGENT_WORKER_NAME_PREFIX, worker_group, suffix) } else { format!("{}-{}-{}", WORKER_NAME_PREFIX, worker_group, suffix) } } pub fn retrieve_common_worker_prefix(worker_name: &str) -> String { let (prefix, _) = worker_name.rsplit_once('-').unzip(); prefix .expect("Invalid worker_name: expected at least one '-' in the name") .to_owned() } pub fn paginate(pagination: Pagination) -> (usize, usize) { let per_page = pagination .per_page .unwrap_or(DEFAULT_PER_PAGE) .max(1) .min(MAX_PER_PAGE); let offset = (pagination.page.unwrap_or(1).max(1) - 1) * per_page; (per_page, offset) } pub fn paginate_without_limits(pagination: Pagination) -> (usize, usize) { let per_page = pagination.per_page.unwrap_or(MAX_PER_PAGE); let offset = (pagination.page.unwrap_or(1).max(1) - 1) * per_page; (per_page, offset) } pub async fn now_from_db<'c, E: sqlx::PgExecutor<'c>>( db: E, ) -> Result> { Ok(sqlx::query_scalar!("SELECT now()") .fetch_one(db) .await? .unwrap()) } pub async fn create_directory_async(directory_path: &str) { AsyncDirBuilder::new() .recursive(true) .create(directory_path) .await .expect("could not create dir"); } pub fn create_directory_sync(directory_path: &str) { SyncDirBuilder::new() .recursive(true) .create(directory_path) .expect("could not create dir"); } #[track_caller] pub fn not_found_if_none>(opt: Option, kind: &str, name: U) -> Result { if let Some(o) = opt { Ok(o) } else { let loc = Location::caller(); Err(Error::NotFound(format!( "{} not found at name {} ({}:{})", kind, name.as_ref(), loc.file().split("/").last().unwrap_or_default(), loc.line() ))) } } pub async fn query_elems_from_hub( http_client: &reqwest::Client, url: &str, query_params: Option>, db: &DB, ) -> Result<( reqwest::StatusCode, reqwest::header::HeaderMap, axum::body::Body, )> { let response = http_get_from_hub(http_client, url, false, query_params, Some(db)).await?; let status = response.status(); Ok(( status, response.headers().clone(), axum::body::Body::from_stream(response.bytes_stream()), )) } pub async fn http_get_from_hub( http_client: &reqwest::Client, url: &str, plain: bool, query_params: Option>, db: Option<&Pool>, ) -> Result { let uid = match db { Some(db) => match get_uid(db).await { Ok(uid) => Some(uid), Err(err) => { tracing::info!("No valid uid found: {}", err); None } }, None => None, }; let mut request = http_client.get(url).header( "Accept", if plain { "text/plain" } else { "application/json" }, ); if let Some(uid) = uid { request = request.header("X-uid", uid); } if let Some(query_params) = query_params { for (key, value) in query_params { request = request.query(&[(key, value)]); } } let response = request .send() .await .context(format!("error fetching script at {url} from hub"))?; Ok(response) } pub fn rd_string(len: usize) -> String { rng() .sample_iter(&Alphanumeric) .take(len) .map(char::from) .collect() } pub fn calculate_hash(s: &str) -> String { let mut hasher = Sha256::new(); hasher.update(s); format!("{:x}", hasher.finalize()) } pub async fn get_uid<'c, E: sqlx::Executor<'c, Database = Postgres>>(db: E) -> Result { let mut uid = LICENSE_KEY_ID.read().await.clone(); if uid == "" { let uid_value = sqlx::query_scalar!( "SELECT value FROM global_settings WHERE name = $1", UNIQUE_ID_SETTING ) .fetch_one(db) .await?; uid = serde_json::from_value::(uid_value).map_err(to_anyhow)?; } Ok(uid) } pub fn map_string_to_number(s: &str, max_number: u64) -> u64 { use std::collections::hash_map::DefaultHasher; use std::hash::{Hash, Hasher}; let mut hasher = DefaultHasher::new(); s.hash(&mut hasher); hasher.finish() % (max_number + 1) } #[derive(Serialize, Deserialize, PartialEq, Clone)] #[serde(rename_all = "lowercase")] pub enum Mode { Worker, Agent, Server, Standalone, Indexer, MCP, } impl std::fmt::Display for Mode { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { Mode::Worker => write!(f, "worker"), Mode::Agent => write!(f, "agent"), Mode::Server => write!(f, "server"), Mode::Standalone => write!(f, "standalone"), Mode::Indexer => write!(f, "indexer"), Mode::MCP => write!(f, "mcp"), } } } // inspired from rails: https://github.com/rails/rails/blob/6e49cc77ab3d16c06e12f93158eaf3e507d4120e/activerecord/lib/active_record/migration.rb#L1308 pub fn generate_lock_id(database_name: &str) -> i64 { const CRC_IEEE: crc::Crc = crc::Crc::::new(&crc::CRC_32_ISO_HDLC); // 0x3d32ad9e chosen by fair dice roll 0x3d32ad9e * (CRC_IEEE.checksum(database_name.as_bytes()) as i64) } pub async fn report_critical_error( error_message: String, _db: DB, workspace_id: Option<&str>, resource: Option<&str>, ) -> () { tracing::error!("CRITICAL ERROR: {error_message}"); let mute_global = CRITICAL_ALERT_MUTE_UI_ENABLED.load(Ordering::Relaxed); let mute_workspace = if let Some(workspace_id) = workspace_id { match fetch_mute_workspace(&_db, workspace_id).await { Ok(flag) => flag, Err(err) => { tracing::error!("Error fetching mute_workspace: {}", err); false } } } else { false }; // we ack_global if mute_global is true, or if mute_workspace is true // but we ignore global mute setting for ack_workspace let acknowledge_workspace = mute_workspace; let acknowledge_global = mute_global || mute_workspace || (workspace_id.is_some() && *CLOUD_HOSTED); if let Err(err) = sqlx::query!( "INSERT INTO alerts (alert_type, message, acknowledged, acknowledged_workspace, workspace_id, resource) VALUES ('critical_error', $1, $2, $3, $4, $5)", error_message, acknowledge_global, acknowledge_workspace, workspace_id, resource, ) .execute(&_db) .await { tracing::error!("Failed to save critical error to database: {}", err); } #[cfg(feature = "enterprise")] if *CLOUD_HOSTED && workspace_id.is_some() { tracing::error!(error_message) } else { send_critical_alert(error_message, &_db, CriticalAlertKind::CriticalError, None).await; } } pub async fn report_recovered_critical_error( message: String, _db: DB, workspace_id: Option<&str>, resource: Option<&str>, ) -> () { tracing::info!("RECOVERED CRITICAL ERROR: {message}"); if let Err(err) = sqlx::query!( "INSERT INTO alerts (alert_type, message, acknowledged, acknowledged_workspace, workspace_id, resource) VALUES ('recovered_critical_error', $1, $2, $3, $4, $5)", message, true, true, workspace_id, resource, ) .execute(&_db) .await { tracing::error!("Failed to save recovered critical error to database: {}", err); } // acknowledge all alerts with the same resource if let Some(resource) = resource { if let Err(err) = sqlx::query!( "UPDATE alerts SET acknowledged = true, acknowledged_workspace = true WHERE resource = $1 AND alert_type = 'critical_error'", resource, ) .execute(&_db) .await { tracing::error!("Failed to acknowledge critical error alerts for resource {}: {}", resource, err); } } #[cfg(feature = "enterprise")] if *CLOUD_HOSTED && workspace_id.is_some() { tracing::error!(message); } else { send_critical_alert( message, &_db, CriticalAlertKind::RecoveredCriticalError, None, ) .await; } } pub trait IsEmpty { fn is_empty(&self) -> bool; } impl IsEmpty for String { fn is_empty(&self) -> bool { self.is_empty() } } impl IsEmpty for Vec { fn is_empty(&self) -> bool { self.is_empty() } } impl IsEmpty for Option where T: IsEmpty, { fn is_empty(&self) -> bool { match self { Some(v) => v.is_empty(), None => true, } } } pub fn empty_as_none<'de, D, T>(deserializer: D) -> std::result::Result, D::Error> where D: Deserializer<'de>, T: Deserialize<'de> + IsEmpty, { let option = as serde::Deserialize>::deserialize(deserializer)?; Ok(option.filter(|s| !s.is_empty())) } pub fn is_empty(value: &T) -> bool where T: IsEmpty, { value.is_empty() } pub fn deserialize_url<'de, D: Deserializer<'de>>( de: D, ) -> std::result::Result, D::Error> { let intermediate = >>::deserialize(de)?; match intermediate.as_deref() { None | Some("") => Ok(None), Some(non_empty_string) => Url::parse(non_empty_string) .map(Some) .map_err(D::Error::custom), } } pub async fn fetch_mute_workspace(_db: &DB, workspace_id: &str) -> Result { match sqlx::query!( "SELECT mute_critical_alerts FROM workspace_settings WHERE workspace_id = $1", workspace_id ) .fetch_optional(_db) .await { Ok(Some(record)) => Ok(record.mute_critical_alerts.unwrap_or(false)), Ok(None) => { tracing::warn!( "Workspace ID {} not found in workspace_settings table", workspace_id ); Ok(false) } Err(err) => { tracing::error!( "Error querying workspace_settings for workspace_id {}: {}", workspace_id, err ); return Err(err.into()); } } } // build_arg_str(&[("name", Some("value")), ("name2", None)], " ", "=") pub fn build_arg_str(args: &[(&str, Option<&str>)], sep: &str, eq: &str) -> String { args.iter() .filter_map(|(k, v)| { if let Some(value) = v { Some(format!("{}{}{}", k, eq, value)) } else { None } }) .join(sep) } // Some errors (duckdb) leak the password in the error message pub fn sanitize_string_from_password(s: &str, passwd: &str) -> Option { if s.contains(passwd) { return Some(s.replace(passwd, "******")); } // Do NOT check substrings // In the case the user finds a string and notices that it gets substituted, // He can very easily find the next character in O(1) and thus the entire password None } pub enum ScheduleType { Croner(Cron), Cron(cron::Schedule), } impl ScheduleType { pub fn find_next( &self, starting_from: &chrono::DateTime, ) -> chrono::DateTime { match self { ScheduleType::Croner(croner_schedule) => croner_schedule .find_next_occurrence(starting_from, false) .expect("cron: a schedule should have a next event"), ScheduleType::Cron(schedule) => schedule .after(starting_from) .next() .expect("cron: a schedule should have a next event"), } } pub fn from_str( schedule_str: &str, version: Option<&str>, seconds_required: bool, ) -> Result { tracing::debug!( "Attempting to parse schedule string: {}, with version: {:?}", schedule_str, version ); match version { Some("v1") | None => { // Use Cron for v1 or if not provided cron::Schedule::from_str(schedule_str) .map(ScheduleType::Cron) .map_err(|e| { tracing::error!( "Failed to parse schedule string '{}' using Cron: {}", schedule_str, e ); Error::BadRequest(format!("cron: {}", e)) }) } Some("v2") | Some(_) => { // Use Croner for v2 let schedule_type_result = panic::catch_unwind(AssertUnwindSafe(|| { let mut croner = Cron::new(schedule_str); if seconds_required { croner.with_seconds_required(); } else { croner.with_seconds_optional(); }; croner.parse() })) .map_err(|_| { tracing::error!( "A panic occurred while parsing schedule string '{}' using Croner", schedule_str, ); Error::BadRequest(format!("cron: a panic occurred during schedule parsing")) }) .and_then(|parse_result| { parse_result.map(ScheduleType::Croner).map_err(|e| { tracing::error!( "Failed to parse schedule string '{}' using Croner: {}", schedule_str, e ); Error::BadRequest(format!("cron: {}", e)) }) }); // Additional check to make sure the provided schedule can generate a next event if let Ok(ScheduleType::Croner(croner_schedule)) = &schedule_type_result { let test_time = chrono::Utc::now().with_timezone(&chrono_tz::UTC); let result = panic::catch_unwind(AssertUnwindSafe(|| { croner_schedule .find_next_occurrence(&test_time, false) .expect("cron: a schedule should have a next event"); })); if let Err(_) = result { tracing::error!("A panic occurred while finding the next occurrence"); return Err(Error::BadRequest(format!( "cron: a panic occurred during find_next_occurrence" ))); } if let Err(e) = result { tracing::error!( "An error occurred while finding the next occurrence: {:?}", e ); return Err(Error::BadRequest(format!( "cron: error during find_next_occurrence: {:?}", e ))); } } schedule_type_result } } } pub fn upcoming( &self, tz: chrono_tz::Tz, count: usize, // Number of upcoming events to take ) -> Result>> { let start_time = Utc::now().with_timezone(&tz); let mut events: Vec> = Vec::with_capacity(count); match self { ScheduleType::Croner(croner_schedule) => { croner_schedule .iter_from(start_time) .map(|dt| dt.with_timezone(&Utc)) .take(count) .for_each(|event| events.push(event)); } ScheduleType::Cron(schedule) => { schedule .upcoming(tz) .map(|x| x.with_timezone(&Utc)) .take(count) .for_each(|event| events.push(event)); } }; // Make sure the schedule is valid and can actually generate "count" events if events.len() != count { return Err(Error::BadRequest(format!( "cron: failed to generate the requested number of events. Expected {}, got {}", count, events.len() ))); } Ok(events) } } use std::future::Future; use std::pin::Pin; use std::task::{Context as TContext, Poll}; use tokio::time::{self, Duration, Sleep}; use pin_project_lite::pin_project; pub trait WarnAfterExt: Future + Sized { /// Warns if the future takes longer than the specified number of seconds to complete. #[track_caller] fn warn_after_seconds(self, seconds: u8) -> WarnAfterFuture { let caller = Location::caller(); self.build_from_caller(seconds, caller, None) } fn build_from_caller( self, seconds: u8, caller: &Location, sql: Option, ) -> WarnAfterFuture { let location = format!("{}:{}", caller.file(), caller.line()); WarnAfterFuture { future: self, timeout: time::sleep(Duration::from_secs(seconds as u64)), warned: false, start_time: std::time::Instant::now(), location, seconds, sql, } } #[track_caller] fn warn_after_seconds_with_sql(self, seconds: u8, sql: String) -> WarnAfterFuture { let caller = Location::caller(); self.build_from_caller(seconds, caller, Some(sql)) } } // Blanket implementation for all futures. impl WarnAfterExt for F {} pin_project! { /// A future that wraps another future and prints a warning if it takes too long. pub struct WarnAfterFuture { #[pin] future: F, #[pin] timeout: Sleep, warned: bool, location: String, start_time: std::time::Instant, seconds: u8, sql: Option, } } impl Future for WarnAfterFuture { type Output = F::Output; fn poll(self: Pin<&mut Self>, cx: &mut TContext<'_>) -> Poll { let this = self.project(); fn build_query_string(location: &str, sql: Option<&str>) -> String { match sql { Some(sql) => format!("{}: {}", location, sql), None => location.to_string(), } } // Poll the timeout future to check if it has elapsed. if !*this.warned { if this.timeout.poll(cx).is_ready() { tracing::warn!( location = this.location, "SLOW_QUERY: query {} to db taking longer than expected (> {} seconds)", build_query_string(&this.location, this.sql.as_deref()), this.seconds, ); *this.warned = true; } } // Poll the wrapped future. match this.future.poll(cx) { Poll::Ready(output) => { if *this.warned { let elapsed = this.start_time.elapsed(); tracing::warn!( location = this.location, "SLOW_QUERY: completed query {} with total duration: {:.2?}", build_query_string(&this.location, this.sql.as_deref()), elapsed ); } Poll::Ready(output) } Poll::Pending => Poll::Pending, } } } #[derive(Debug, Deserialize, Clone, Copy, PartialEq, Eq, Hash)] #[serde(rename_all = "lowercase")] pub enum RunnableKind { Script, Flow, } impl Display for RunnableKind { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let runnable_kind = match self { RunnableKind::Script => "script", RunnableKind::Flow => "flow", }; write!(f, "{}", runnable_kind) } } #[cfg(test)] mod tests { use super::*; #[test] fn test_build_arg_str() { let r = build_arg_str( &[ ("host", Some("localhost")), ("port", Some("5432")), ("password", None), ("user", Some("postgres")), ("dbname", Some("test_db")), ], " ", "=", ); assert_eq!(r, "host=localhost port=5432 user=postgres dbname=test_db"); } } #[derive(Clone)] pub struct ExpiringCacheEntry { pub value: T, pub expiry: std::time::Instant, }