feat: limit pro plans (#2794)

windmill-labs · Dec 6, 2023 · c58190e · c58190e
1 parent e4da819
commit c58190e
Show file tree

Hide file tree

Showing 10 changed files with 82 additions and 22 deletions.
diff --git a/backend/windmill-api/src/lib.rs b/backend/windmill-api/src/lib.rs
@@ -36,6 +36,8 @@ use tower_http::{
     trace::TraceLayer,
 };
 use windmill_common::db::UserDB;
+#[cfg(feature = "enterprise")]
+use windmill_common::ee::{get_license_plan, LicensePlan};
 use windmill_common::utils::rd_string;
 use windmill_common::worker::ALL_TAGS;
 use windmill_common::BASE_URL;
@@ -159,7 +161,10 @@ pub async fn run_server(
         .allow_origin(Any);
 
     #[cfg(feature = "enterprise")]
-    let sp_extension: (ServiceProviderExt, SamlSsoLogin) = saml::build_sp_extension().await?;
+    let sp_extension: (ServiceProviderExt, SamlSsoLogin) = match get_license_plan().await {
+        LicensePlan::Enterprise => saml::build_sp_extension().await?,
+        LicensePlan::Pro => (ServiceProviderExt(None), SamlSsoLogin(None)),
+    };
 
     #[cfg(not(feature = "enterprise"))]
     let sp_extension = (ServiceProviderExt(), SamlSsoLogin(None));

diff --git a/backend/windmill-api/src/saml.rs b/backend/windmill-api/src/saml.rs
@@ -40,6 +40,9 @@ pub struct ServiceProviderExt(pub Option<ServiceProvider>);
 #[cfg(not(feature = "enterprise"))]
 pub struct ServiceProviderExt();
 
+#[cfg(feature = "enterprise")]
+use windmill_common::ee::{get_license_plan, LicensePlan};
+
 pub struct SamlSsoLogin(pub Option<String>);
 
 #[cfg(feature = "enterprise")]
@@ -97,6 +100,11 @@ pub async fn acs(
     Extension(se): Extension<Arc<ServiceProviderExt>>,
     Form(s): Form<SamlForm>,
 ) -> Result<Redirect> {
+    if matches!(get_license_plan().await, LicensePlan::Pro) {
+        return Err(Error::BadRequest(
+            "SAML not available in the pro plan".to_string(),
+        ));
+    };
     if let Some(sp_m) = &se.0 {
         let sp = sp_m.clone();
         if let Some(encoded_resp) = s.SAMLResponse {

diff --git a/backend/windmill-api/src/users.rs b/backend/windmill-api/src/users.rs
@@ -42,6 +42,8 @@ use time::OffsetDateTime;
 use tower_cookies::{Cookie, Cookies};
 use tracing::{Instrument, Span};
 use windmill_audit::{audit_log, ActionKind};
+#[cfg(feature = "enterprise")]
+use windmill_common::ee::{get_license_plan, LicensePlan};
 use windmill_common::oauth2::REQUIRE_PREEXISTING_USER_FOR_OAUTH;
 use windmill_common::users::truncate_token;
 use windmill_common::worker::{CLOUD_HOSTED, SERVER_CONFIG};
@@ -1602,7 +1604,6 @@ async fn create_user(
         ));
     }
 
-    #[cfg(not(feature = "enterprise"))]
     _check_nb_of_user(&db).await?;
 
     sqlx::query!(
@@ -2344,6 +2345,10 @@ pub struct LoginUserInfo {
 }
 
 async fn _check_nb_of_user(db: &DB) -> Result<()> {
+    #[cfg(feature = "enterprise")]
+    if matches!(get_license_plan().await, LicensePlan::Enterprise) {
+        return Ok(());
+    }
     let nb_users_sso =
         sqlx::query_scalar!("SELECT COUNT(*) FROM password WHERE login_type != 'password'",)
             .fetch_one(db)
@@ -2430,7 +2435,6 @@ pub async fn login_externally(
             name = user.clone().unwrap().displayName;
         }
 
-        #[cfg(not(feature = "enterprise"))]
         _check_nb_of_user(&db).await?;
 
         sqlx::query(&format!(

diff --git a/backend/windmill-audit/src/lib.rs b/backend/windmill-audit/src/lib.rs
@@ -15,6 +15,9 @@ use windmill_common::{
     utils::Pagination,
 };
 
+#[cfg(feature = "enterprise")]
+use windmill_common::ee::{get_license_plan, LicensePlan};
+
 use serde::{Deserialize, Serialize};
 use sql_builder::SqlBuilder;
 use sqlx::{FromRow, Postgres, Transaction};
@@ -44,23 +47,32 @@ pub struct AuditLog {
 pub async fn audit_log<'c, E: sqlx::Executor<'c, Database = Postgres>>(
     db: E,
     username: &str,
-    _operation: &str,
+    mut _operation: &str,
     action_kind: ActionKind,
     w_id: &str,
-    _resource: Option<&str>,
+    mut _resource: Option<&str>,
     _parameters: Option<HashMap<&str, &str>>,
 ) -> Result<()> {
     #[cfg(feature = "enterprise")]
-    let p_json: serde_json::Value = serde_json::to_value(&_parameters).unwrap();
+    let p_json = match get_license_plan().await {
+        LicensePlan::Enterprise => serde_json::to_value(&_parameters).unwrap(),
+        LicensePlan::Pro => serde_json::json!({"redacted": "-"}),
+    };
 
     #[cfg(not(feature = "enterprise"))]
     let p_json: serde_json::Value = serde_json::json!({"redacted": "-"});
 
-    #[cfg(not(feature = "enterprise"))]
-    let _resource: Option<&str> = Some("EE only");
+    #[cfg(feature = "enterprise")]
+    if matches!(get_license_plan().await, LicensePlan::Pro) {
+        _resource = Some("EE only");
+        _operation = "redacted";
+    }
 
     #[cfg(not(feature = "enterprise"))]
-    let _operation: &str = "redacted";
+    {
+        _resource = Some("EE only");
+        _operation = "redacted";
+    }
 
     tracing::info!(
         operation = _operation,

diff --git a/backend/windmill-common/src/ee.rs b/backend/windmill-common/src/ee.rs
@@ -6,3 +6,18 @@ lazy_static::lazy_static! {
   pub static ref LICENSE_KEY_ID: Arc<RwLock<String>> = Arc::new(RwLock::new("".to_string()));
   pub static ref LICENSE_KEY: Arc<RwLock<String>> = Arc::new(RwLock::new("".to_string()));
 }
+
+pub enum LicensePlan {
+    Pro,
+    Enterprise,
+}
+
+pub async fn get_license_plan() -> LicensePlan {
+    let id = LICENSE_KEY_ID.read().await.clone();
+
+    if id.ends_with("_pro") {
+        LicensePlan::Pro
+    } else {
+        LicensePlan::Enterprise
+    }
+}
diff --git a/backend/windmill-worker/src/python_executor.rs b/backend/windmill-worker/src/python_executor.rs
@@ -10,6 +10,8 @@ use tokio::{
     process::Command,
 };
 use uuid::Uuid;
+#[cfg(feature = "enterprise")]
+use windmill_common::ee::{get_license_plan, LicensePlan};
 use windmill_common::{
     error::{self, Error},
     jobs::QueuedJob,
@@ -649,12 +651,16 @@ pub async fn handle_python_reqs(
 
         #[cfg(feature = "enterprise")]
         if let Some(ref bucket) = *S3_CACHE_BUCKET {
-            sqlx::query_scalar!("UPDATE queue SET last_ping = now() WHERE id = $1", job_id)
-                .execute(db)
-                .await?;
-            if pull_from_tar(bucket, venv_p.clone()).await.is_ok() {
-                req_paths.push(venv_p.clone());
-                continue;
+            if matches!(get_license_plan().await, LicensePlan::Pro) {
+                tracing::warn!("S3 cache not available in the pro plan");
+            } else {
+                sqlx::query_scalar!("UPDATE queue SET last_ping = now() WHERE id = $1", job_id)
+                    .execute(db)
+                    .await?;
+                if pull_from_tar(bucket, venv_p.clone()).await.is_ok() {
+                    req_paths.push(venv_p.clone());
+                    continue;
+                }
             }
         }
 
@@ -770,8 +776,12 @@ pub async fn handle_python_reqs(
 
         #[cfg(feature = "enterprise")]
         if let Some(ref bucket) = *S3_CACHE_BUCKET {
-            let venv_p = venv_p.clone();
-            tokio::spawn(build_tar_and_push(bucket, venv_p));
+            if matches!(get_license_plan().await, LicensePlan::Pro) {
+                tracing::warn!("S3 cache not available in the pro plan");
+            } else {
+                let venv_p = venv_p.clone();
+                tokio::spawn(build_tar_and_push(bucket, venv_p));
+            }
         }
         req_paths.push(venv_p);
     }

diff --git a/backend/windmill-worker/src/worker.rs b/backend/windmill-worker/src/worker.rs
@@ -27,6 +27,8 @@ use std::{
 };
 
 use uuid::Uuid;
+#[cfg(feature = "enterprise")]
+use windmill_common::ee::{get_license_plan, LicensePlan};
 use windmill_common::{
     error::{self, to_anyhow, Error},
     flows::{FlowModule, FlowModuleValue, FlowValue},
@@ -855,7 +857,9 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
     #[cfg(feature = "enterprise")]
     if i_worker == 1 {
         if let Some(ref s) = S3_CACHE_BUCKET.clone() {
-            if crate::global_cache::worker_s3_bucket_sync_enabled(&db).await {
+            if matches!(get_license_plan().await, LicensePlan::Pro) {
+                tracing::warn!("S3 cache not available in the pro plan");
+            } else if crate::global_cache::worker_s3_bucket_sync_enabled(&db).await {
                 let bucket = s.to_string();
                 let worker_name2 = worker_name.clone();
 
@@ -1254,7 +1258,9 @@ pub async fn run_worker<R: rsmq_async::RsmqConnection + Send + Sync + Clone + 's
 
         #[cfg(feature = "enterprise")]
         if i_worker == 1 && S3_CACHE_BUCKET.is_some() {
-            if last_sync.elapsed().as_secs() > *GLOBAL_CACHE_INTERVAL
+            if matches!(get_license_plan().await, LicensePlan::Pro) {
+                tracing::warn!("S3 cache not available in the pro plan");
+            } else if last_sync.elapsed().as_secs() > *GLOBAL_CACHE_INTERVAL
                 && (copy_cache_from_bucket_handle.is_none()
                     || copy_cache_from_bucket_handle
                         .as_ref()

diff --git a/frontend/src/lib/components/InstanceSettings.svelte b/frontend/src/lib/components/InstanceSettings.svelte
@@ -197,7 +197,7 @@
 					{#if category == 'SSO/OAuth'}
 						<div class="mb-6">
 							<h4 class="pb-4">SSO</h4>
-							{#if !$enterpriseLicense}
+							{#if !$enterpriseLicense || $enterpriseLicense.endsWith('_pro')}
 								<Alert type="warning" title="Limited to 10 SSO users">
 									Without EE, the number of SSO users is limited to 10. SCIM/SAML is available on EE
 								</Alert>

diff --git a/frontend/src/routes/(root)/(logged)/audit_logs/+page.svelte b/frontend/src/routes/(root)/(logged)/audit_logs/+page.svelte
@@ -71,7 +71,7 @@
 			</div>
 		</div>
 
-		{#if !$enterpriseLicense}
+		{#if !$enterpriseLicense || $enterpriseLicense.endsWith('_pro')}
 			<Alert title="Redacted audit logs" type="warning">
 				You need an enterprise license to see unredacted audit logs.
 			</Alert>

diff --git a/frontend/src/routes/(root)/(logged)/workers/+page.svelte b/frontend/src/routes/(root)/(logged)/workers/+page.svelte
@@ -135,7 +135,7 @@
 							}
 						}}
 						options={{ right: 'global cache to s3' }}
-						disabled={!$enterpriseLicense}
+						disabled={!$enterpriseLicense || $enterpriseLicense.endsWith('_pro')}
 					/>
 					<Tooltip
 						><p