From a7c0e4dcd0ae6f61a84d2a36e9e40a3113cb54f0 Mon Sep 17 00:00:00 2001
From: Anastasia Lubennikova <anastasia@neon.tech>
Date: Wed, 30 Aug 2023 16:04:48 +0300
Subject: [PATCH] Check if custiom extension is enabled. This check was lost in
 the latest refactoring.

If extension is not present in 'public_extensions' or 'custom_extensions' don't download it
---
 compute_tools/src/extension_server.rs | 14 +++++++++++++-
 libs/compute_api/src/spec.rs          | 12 ++++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/compute_tools/src/extension_server.rs b/compute_tools/src/extension_server.rs
index cb54a603e0..54c22026e7 100644
--- a/compute_tools/src/extension_server.rs
+++ b/compute_tools/src/extension_server.rs
@@ -180,7 +180,19 @@ pub async fn download_extension(
 // Create extension control files from spec
 pub fn create_control_files(remote_extensions: &RemoteExtSpec, pgbin: &str) {
     let local_sharedir = Path::new(&get_pg_config("--sharedir", pgbin)).join("extension");
-    for ext_data in remote_extensions.extension_data.values() {
+    for (ext_name, ext_data) in remote_extensions.extension_data.iter() {
+        // Check if extension is present in public or custom.
+        // If not, then it is not allowed to be used by this compute.
+        if let Some(public_extensions) = &remote_extensions.public_extensions {
+            if !public_extensions.contains(ext_name) {
+                if let Some(custom_extensions) = &remote_extensions.custom_extensions {
+                    if !custom_extensions.contains(ext_name) {
+                        continue; // skip this extension, it is not allowed
+                    }
+                }
+            }
+        }
+
         for (control_name, control_content) in &ext_data.control_data {
             let control_path = local_sharedir.join(control_name);
             if !control_path.exists() {
diff --git a/libs/compute_api/src/spec.rs b/libs/compute_api/src/spec.rs
index 78361ce07e..b41ca8c9cf 100644
--- a/libs/compute_api/src/spec.rs
+++ b/libs/compute_api/src/spec.rs
@@ -106,6 +106,18 @@ impl RemoteExtSpec {
                 .ok_or(anyhow::anyhow!("library {} is not found", lib_raw_name))?;
         }
 
+        // Check if extension is present in public or custom.
+        // If not, then it is not allowed to be used by this compute.
+        if let Some(public_extensions) = &self.public_extensions {
+            if !public_extensions.contains(&real_ext_name.to_string()) {
+                if let Some(custom_extensions) = &self.custom_extensions {
+                    if !custom_extensions.contains(&real_ext_name.to_string()) {
+                        return Err(anyhow::anyhow!("extension {} is not found", real_ext_name));
+                    }
+                }
+            }
+        }
+
         match self.extension_data.get(real_ext_name) {
             Some(_ext_data) => {
                 // Construct the path to the extension archive