feature: cloud function security improvement with shared secret (#31)

* terraform and backend changes to support the auth token fix auth and cors use the X-Signature header instaed of Authorization add type definitions for crypto-js ignore all dist folders * Add a test script for the local script * improve readme for cloud function * Update changelog * docs: minor typos * docs: updating extension source for new cloud function deployment env var --------- Co-authored-by: Luka Fontanilla <[email protected]>
looker-open-source · May 6, 2024 · b59eb53 · b59eb53
1 parent aa75432
commit b59eb53
Show file tree

Hide file tree

Showing 14 changed files with 212 additions and 68 deletions.
diff --git a/.gitignore b/.gitignore
@@ -8,4 +8,6 @@ terraform.tfstate*
 *.tfstate
 .venv
 node_modules
-dist/
+
+.vertex_cf_auth_token
+dist
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,11 @@
 # Changelog
 
+## v2.1
+
+### Added
+- Shared secret for the cloud function and explore assistant extension
+- Terraform code for managing the token in the GCP Secrets Manager
+
 ## v2.0
 
 There are many breaking changes in this version.

diff --git a/explore-assistant-backend/README.md b/explore-assistant-backend/README.md
@@ -24,12 +24,21 @@ terraform init
 
 ### Cloud Function Backend
 
+First create a file that will contain the LOOKER_AUTH_TOKEN and place it at the root. This will be used my the cloud function locally, as well as the extension framework app. The value of this token will uploaded to the GCP project as secret to be used by the Cloud Function.
+
+```bash
+openssl rand -base64 32 > .vertex_cf_auth_token
+
+```
+
 To deploy the Cloud Function backend:
 
 ```bash
 export TF_VAR_project_id=XXX
 export TF_VAR_use_bigquery_backend=0
 export TF_VAR_use_cloud_function_backend=1
+export TF_VAR_looker_auth_token=$(cat ../../.vertex_cf_auth_token)
+terraform init
 terraform plan
 terraform apply
 ```

diff --git a/explore-assistant-backend/terraform/cloud_function/main.tf b/explore-assistant-backend/terraform/cloud_function/main.tf
@@ -1,14 +1,14 @@
 
 variable "cloud_run_service_name" {
-    type = string
+  type = string
 }
 
 variable "deployment_region" {
-    type = string
+  type = string
 }
 
 variable "project_id" {
-    type = string
+  type = string
 }
 
 resource "google_service_account" "explore-assistant-sa" {
@@ -36,12 +36,37 @@ resource "google_project_iam_member" "iam_service_account_act_as" {
 }
 
 # IAM permission as Editor
-resource "google_project_iam_member" "iam_looker_service_usage" {  
+resource "google_project_iam_member" "iam_looker_service_usage" {
   project = var.project_id
   role    = "roles/serviceusage.serviceUsageConsumer"
   member  = format("serviceAccount:%s", google_service_account.explore-assistant-sa.email)
 }
 
+resource "google_secret_manager_secret" "vertex_cf_auth_token" {
+  project   = var.project_id
+  secret_id = "VERTEX_CF_AUTH_TOKEN"
+  replication {
+    user_managed {
+      replicas {
+        location = var.deployment_region
+      }
+    }
+  }
+}
+
+resource "google_secret_manager_secret_version" "vertex_cf_auth_token_version" {
+  secret      = google_secret_manager_secret.vertex_cf_auth_token.name
+  secret_data = file("${path.module}/../../../.vertex_cf_auth_token")
+}
+
+resource "google_secret_manager_secret_iam_binding" "vertex_cf_auth_token_accessor" {
+  secret_id = google_secret_manager_secret.vertex_cf_auth_token.secret_id
+  role      = "roles/secretmanager.secretAccessor"
+  members   = [
+    "serviceAccount:${google_service_account.explore-assistant-sa.email}",
+  ]
+}
+
 resource "random_id" "default" {
   byte_length = 8
 }
@@ -65,11 +90,11 @@ resource "google_storage_bucket_object" "object" {
   source = data.archive_file.default.output_path # Add path to the zipped function source code
 }
 
-resource google_artifact_registry_repository "default" {
+resource "google_artifact_registry_repository" "default" {
   repository_id = "explore-assistant-repo"
   location      = var.deployment_region
   project       = var.project_id
-  format       = "DOCKER"
+  format        = "DOCKER"
 }
 
 resource "google_cloudfunctions2_function" "default" {
@@ -78,8 +103,8 @@ resource "google_cloudfunctions2_function" "default" {
   description = "An endpoint for generating Looker queries from natural language using Generative UI"
 
   build_config {
-    runtime     = "python310"
-    entry_point = "cloud_function_entrypoint" # Set the entry point
+    runtime           = "python310"
+    entry_point       = "cloud_function_entrypoint" # Set the entry point
     docker_repository = google_artifact_registry_repository.default.id
     source {
       storage_source {
@@ -90,34 +115,42 @@ resource "google_cloudfunctions2_function" "default" {
 
     environment_variables = {
       FUNCTIONS_FRAMEWORK = 1
-      SOURCE_HASH = data.archive_file.default.output_sha
+      SOURCE_HASH         = data.archive_file.default.output_sha
     }
   }
 
   service_config {
-    max_instance_count = 10
-    min_instance_count = 1
-    available_memory   = "4Gi"
-    timeout_seconds    = 60
-    available_cpu      = "4"
+    max_instance_count               = 10
+    min_instance_count               = 1
+    available_memory                 = "4Gi"
+    timeout_seconds                  = 60
+    available_cpu                    = "4"
     max_instance_request_concurrency = 20
     environment_variables = {
-        REGION = var.deployment_region
-        PROJECT = var.project_id
+      REGION  = var.deployment_region
+      PROJECT = var.project_id
     }
+
+    secret_environment_variables {
+      key     = "VERTEX_CF_AUTH_TOKEN"
+      project_id = var.project_id
+      secret  = google_secret_manager_secret.vertex_cf_auth_token.secret_id
+      version = "latest"
+    }
+
     all_traffic_on_latest_revision = true
-    service_account_email = google_service_account.explore-assistant-sa.email
+    service_account_email          = google_service_account.explore-assistant-sa.email
   }
 }
 
 ### IAM permissions for Cloud Functions Gen2 (requires run invoker as well) for public access
 
 resource "google_cloudfunctions2_function_iam_member" "default" {
-  location = google_cloudfunctions2_function.default.location
-  project  = google_cloudfunctions2_function.default.project
-  cloud_function  = google_cloudfunctions2_function.default.name
-  role     = "roles/cloudfunctions.invoker"
-  member   = "allUsers"
+  location       = google_cloudfunctions2_function.default.location
+  project        = google_cloudfunctions2_function.default.project
+  cloud_function = google_cloudfunctions2_function.default.name
+  role           = "roles/cloudfunctions.invoker"
+  member         = "allUsers"
 }
 
 data "google_iam_policy" "noauth" {
@@ -130,9 +163,9 @@ data "google_iam_policy" "noauth" {
 }
 
 resource "google_cloud_run_service_iam_policy" "noauth" {
-  location    = google_cloudfunctions2_function.default.location
-  project     = google_cloudfunctions2_function.default.project
-  service     = google_cloudfunctions2_function.default.name
+  location = google_cloudfunctions2_function.default.location
+  project  = google_cloudfunctions2_function.default.project
+  service  = google_cloudfunctions2_function.default.name
 
   policy_data = data.google_iam_policy.noauth.policy_data
 }
@@ -143,4 +176,4 @@ output "function_uri" {
 
 output "data" {
   value = google_cloudfunctions2_function.default
-} 
+}
diff --git a/explore-assistant-backend/terraform/main.tf b/explore-assistant-backend/terraform/main.tf
@@ -21,7 +21,8 @@ module "project-services" {
     "storage-api.googleapis.com",
     "storage.googleapis.com",
     "aiplatform.googleapis.com",
-    "compute.googleapis.com"
+    "compute.googleapis.com",
+    "secretmanager.googleapis.com",
   ]
 }
 

diff --git a/explore-assistant-cloud-function/README.md b/explore-assistant-cloud-function/README.md
@@ -22,6 +22,8 @@ The cloud function integrates with Vertex AI and utilizes the `GenerativeModel`
 
 8. **Execution Environment**: When executed, the script checks if it's running in a Google Cloud Function environment and acts accordingly; otherwise, it starts a Flask web server for local development or testing.
 
+9. **Endpoint Security**: We are using a simple shared secret approach to securing the endpoint. The request body is checked against the supplied signature in the X-Signature header. We aren't yet guarding against replay attacks with nonces.
+
 ## Local Development
 
 To set up and run the function locally, follow these steps:
@@ -43,13 +45,13 @@ To set up and run the function locally, follow these steps:
 3. Run the function locally by executing the main script:
 
     ```bash
-    PROJECT=XXX REGION=us-central1 python main.py
+    PROJECT=XXXX LOCATION=us-central-1 VERTEX_CF_AUTH_TOKEN=$(cat ../.vertex_cf_auth_token) python main.py
     ```
 
 4. Test calling the endpoint locally with a custom query and parameter declaration
 
    ```bash
-     curl -X POST -H "Content-Type: application/json" -d '{"contents":"how are you doing?", "parameters":{"max_output_tokens": 1000}}' http://localhost:8000
+     python test.py
    ```
 
 This setup allows developers to test and modify the function in a local environment before deploying it to a cloud function service.
diff --git a/explore-assistant-cloud-function/main.py b/explore-assistant-cloud-function/main.py
@@ -22,8 +22,8 @@
 # SOFTWARE.
 
 import os
-import json
-from flask import Flask, request
+import hmac
+from flask import Flask, request, Response
 from flask_cors import CORS
 import functions_framework
 import vertexai
@@ -36,8 +36,30 @@
 # Initialize the Vertex AI
 project = os.environ.get("PROJECT")
 location = os.environ.get("REGION")
+vertex_cf_auth_token = os.environ.get("VERTEX_CF_AUTH_TOKEN")
 vertexai.init(project=project, location=location)
 
+def get_response_headers(request):
+    headers = {
+        "Access-Control-Allow-Origin": "*",
+        "Access-Control-Allow-Methods": "POST, OPTIONS",
+        "Access-Control-Allow-Headers": "Content-Type, X-Signature"
+    }
+    return headers
+
+
+def has_valid_signature(request):
+    signature = request.headers.get("X-Signature")
+    if signature is None:
+        return False
+
+    # Validate the signature
+    secret = vertex_cf_auth_token.encode("utf-8")
+    request_data = request.get_data()
+    hmac_obj = hmac.new(secret, request_data, "sha256")
+    expected_signature = hmac_obj.hexdigest()
+
+    return hmac.compare_digest(signature, expected_signature)
 
 def generate_looker_query(contents, parameters=None, model_name="gemini-1.0-pro-001"):
 
@@ -91,17 +113,21 @@ def create_flask_app():
     @app.route("/", methods=["POST", "OPTIONS"])
     def base():
         if request.method == "OPTIONS":
-            return handle_options_request()
+            return handle_options_request(request)
 
         incoming_request = request.get_json()
+        print(incoming_request)
         contents = incoming_request.get("contents")
         parameters = incoming_request.get("parameters")
         if contents is None:
             return "Missing 'contents' parameter", 400
-
+
+        if not has_valid_signature(request):
+            return "Invalid signature", 403
+
         response_text = generate_looker_query(contents, parameters)
-        
-        return response_text, 200, response_headers()
+
+        return response_text, 200, get_response_headers(request)
 
     return app
 
@@ -110,31 +136,26 @@ def base():
 @functions_framework.http
 def cloud_function_entrypoint(request):
     if request.method == "OPTIONS":
-        return handle_options_request()
+        return handle_options_request(request)
 
     incoming_request = request.get_json()
     contents = incoming_request.get("contents")
     parameters = incoming_request.get("parameters")
     if contents is None:
         return "Missing 'contents' parameter", 400
-        
+
     response_text = generate_looker_query(contents, parameters)
 
-    return response_text, 200, response_headers()
+    return response_text, 200, get_response_headers(request)
 
 def response_headers():
     return {
         "Access-Control-Allow-Origin": "*"
     }
 
-def handle_options_request():
-    headers = {
-        "Access-Control-Allow-Origin": "*",
-        "Access-Control-Allow-Methods": "POST, OPTIONS",
-        "Access-Control-Allow-Headers": "Content-Type",
-        "Access-Control-Max-Age": "3600"
-    }
-    return "", 204, headers
+def handle_options_request(request):
+    return "", 204, get_response_headers(request)
+
 
 # Determine the running environment and execute accordingly
 if __name__ == "__main__":

diff --git a/explore-assistant-cloud-function/test.py b/explore-assistant-cloud-function/test.py
@@ -0,0 +1,43 @@
+import hmac
+import hashlib
+import requests
+import json
+
+def generate_hmac_signature(secret_key, data):
+    """
+    Generate HMAC-SHA256 signature for the given data using the secret key.
+    """
+    hmac_obj = hmac.new(secret_key.encode(), json.dumps(data).encode(), hashlib.sha256)
+    return hmac_obj.hexdigest()
+
+def send_request(url, data, signature):
+    """
+    Send a POST request to the given URL with the provided data and HMAC signature.
+    """
+    headers = {
+        'Content-Type': 'application/json',
+        'X-Signature': signature
+    }
+    response = requests.post(url, headers=headers, json=data)
+    return response.text
+
+def main():
+    # URL of the endpoint
+    url = 'http://localhost:8000'
+
+    # Request payload
+    data = {"contents":"how are you doing?", "parameters":{"max_output_tokens": 1000}}
+
+    # Read the secret key from a file
+    with open('../.vertex_cf_auth_token', 'r') as file:
+        secret_key = file.read().strip()  # Remove any potential newline characters
+
+    # Generate HMAC signature
+    signature = generate_hmac_signature(secret_key, data)
+
+    # Send the request
+    response = send_request(url, data, signature)
+    print("Response from server:", response)
+
+if __name__ == "__main__":
+    main()
diff --git a/explore-assistant-extension/.env_example b/explore-assistant-extension/.env_example
@@ -1,6 +1,8 @@
-VERTEX_AI_ENDPOINT=<This is your Deployed Cloud Function Endpoint>
 LOOKER_MODEL=<This is your Looker model name>
 LOOKER_EXPLORE=<This is your Looker explore name>
 
+VERTEX_AI_ENDPOINT=<This is your Deployed Cloud Function Endpoint>
+VERTEX_CF_AUTH_TOKEN=<This is the token used to communicate with the cloud function>
+
 VERTEX_BIGQUERY_LOOKER_CONNECTION_NAME=<This is the connection name that has vertex ai external connector>
 VERTEX_BIGQUERY_MODEL_ID=<This is the model id that you want to use for prediction>