From b59eb532f092b866a0c18b0532fbcb0338653f7e Mon Sep 17 00:00:00 2001 From: Jawad Laraqui Date: Mon, 6 May 2024 15:40:54 -0400 Subject: [PATCH] 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 <51471651+LukaFontanilla@users.noreply.github.com> --- .gitignore | 4 +- CHANGELOG.md | 6 ++ explore-assistant-backend/README.md | 9 ++ .../terraform/cloud_function/main.tf | 85 +++++++++++++------ explore-assistant-backend/terraform/main.tf | 3 +- explore-assistant-cloud-function/README.md | 6 +- explore-assistant-cloud-function/main.py | 55 ++++++++---- explore-assistant-cloud-function/test.py | 43 ++++++++++ explore-assistant-extension/.env_example | 4 +- explore-assistant-extension/README.md | 3 +- explore-assistant-extension/package-lock.json | 12 +++ explore-assistant-extension/package.json | 2 + explore-assistant-extension/src/App.tsx | 26 +++--- .../src/pages/ExploreAssistantPage/index.tsx | 22 +++-- 14 files changed, 212 insertions(+), 68 deletions(-) create mode 100644 explore-assistant-cloud-function/test.py diff --git a/.gitignore b/.gitignore index cb834524..5f818319 100644 --- 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 index e4459613..a8be92bd 100644 --- 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 index 41f5beea..01536ce0 100644 --- 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 index ff658554..45ebb063 100644 --- 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 index 7888b297..563f2e9a 100644 --- 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 index bcf3193d..4ab33a5b 100644 --- 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 index 712dc7f9..4d5b98be 100644 --- 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 new file mode 100644 index 00000000..731af4da --- /dev/null +++ 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 index 01c1e5ce..ff48427e 100644 --- a/explore-assistant-extension/.env_example +++ b/explore-assistant-extension/.env_example @@ -1,6 +1,8 @@ -VERTEX_AI_ENDPOINT= LOOKER_MODEL= LOOKER_EXPLORE= +VERTEX_AI_ENDPOINT= +VERTEX_CF_AUTH_TOKEN= + VERTEX_BIGQUERY_LOOKER_CONNECTION_NAME= VERTEX_BIGQUERY_MODEL_ID= \ No newline at end of file diff --git a/explore-assistant-extension/README.md b/explore-assistant-extension/README.md index b640dda8..92b12ee3 100644 --- a/explore-assistant-extension/README.md +++ b/explore-assistant-extension/README.md @@ -78,10 +78,11 @@ jsonPayload.component="explore-assistant-metadata" VERTEX_BIGQUERY_LOOKER_CONNECTION_NAME= ``` - If you're using the Cloud Function backend, replace the default: + If you're using the Cloud Function backend, replace the defaults: ``` VERTEX_AI_ENDPOINT= + VERTEX_CF_AUTH_TOKEN= ``` If you're using the BigQuery Backend replace the default: diff --git a/explore-assistant-extension/package-lock.json b/explore-assistant-extension/package-lock.json index 7bfe4220..7741dead 100644 --- a/explore-assistant-extension/package-lock.json +++ b/explore-assistant-extension/package-lock.json @@ -16,6 +16,8 @@ "@looker/extension-sdk-react": "^24.2.0", "@looker/sdk": "^24.2.0", "@reduxjs/toolkit": "^2.2.2", + "@types/crypto-js": "^4.2.2", + "crypto-js": "^4.2.0", "react": "^17.0.2", "react-dom": "^17.0.2", "react-is": "^16.13.1", @@ -2775,6 +2777,11 @@ "@types/node": "*" } }, + "node_modules/@types/crypto-js": { + "version": "4.2.2", + "resolved": "https://registry.npmjs.org/@types/crypto-js/-/crypto-js-4.2.2.tgz", + "integrity": "sha512-sDOLlVbHhXpAUAL0YHDUUwDZf3iN4Bwi4W6a0W0b+QcAezUbRtH4FVb+9J4h+XFPW7l/gQ9F8qC7P+Ec4k8QVQ==" + }, "node_modules/@types/eslint": { "version": "8.56.10", "resolved": "https://registry.npmjs.org/@types/eslint/-/eslint-8.56.10.tgz", @@ -4348,6 +4355,11 @@ "node": ">= 8" } }, + "node_modules/crypto-js": { + "version": "4.2.0", + "resolved": "https://registry.npmjs.org/crypto-js/-/crypto-js-4.2.0.tgz", + "integrity": "sha512-KALDyEYgpY+Rlob/iriUtjV6d5Eq+Y191A5g4UqLAi8CyGP9N1+FdVbkc1SxKc2r4YAYqG8JzO2KGL+AizD70Q==" + }, "node_modules/css-color-keywords": { "version": "1.0.0", "resolved": "https://registry.npmjs.org/css-color-keywords/-/css-color-keywords-1.0.0.tgz", diff --git a/explore-assistant-extension/package.json b/explore-assistant-extension/package.json index 244da355..95825f14 100644 --- a/explore-assistant-extension/package.json +++ b/explore-assistant-extension/package.json @@ -24,6 +24,8 @@ "@looker/extension-sdk-react": "^24.2.0", "@looker/sdk": "^24.2.0", "@reduxjs/toolkit": "^2.2.2", + "@types/crypto-js": "^4.2.2", + "crypto-js": "^4.2.0", "react": "^17.0.2", "react-dom": "^17.0.2", "react-is": "^16.13.1", diff --git a/explore-assistant-extension/src/App.tsx b/explore-assistant-extension/src/App.tsx index 147daa3d..edd16015 100644 --- a/explore-assistant-extension/src/App.tsx +++ b/explore-assistant-extension/src/App.tsx @@ -12,19 +12,19 @@ import ExploreAssistantPage from './pages/ExploreAssistantPage' const ExploreApp = () => { return ( - - - - - - - - - - - - - + + + + + + + + + + + + + ) } diff --git a/explore-assistant-extension/src/pages/ExploreAssistantPage/index.tsx b/explore-assistant-extension/src/pages/ExploreAssistantPage/index.tsx index c299c85f..e10fe933 100644 --- a/explore-assistant-extension/src/pages/ExploreAssistantPage/index.tsx +++ b/explore-assistant-extension/src/pages/ExploreAssistantPage/index.tsx @@ -31,6 +31,7 @@ import { RootState } from '../../store' import process from 'process' import { UtilsHelper } from '../../utils/Helper' import { useExampleData } from '../../hooks/useExampleData' +import CryptoJS from 'crypto-js' interface ModelParameters { max_output_tokens?: number @@ -42,8 +43,8 @@ const generateSQL = ( parameters: ModelParameters, ) => { const escapedPrompt = UtilsHelper.escapeQueryAll(prompt); - var subselect = `SELECT '` + escapedPrompt + `' AS prompt`; - + const subselect = `SELECT '` + escapedPrompt + `' AS prompt`; + return ` SELECT ml_generate_text_llm_result AS generated_content @@ -68,7 +69,11 @@ const ExploreAssistantPage = () => { const LOOKER_MODEL = process.env.LOOKER_MODEL || '' const LOOKER_EXPLORE = process.env.LOOKER_EXPLORE || '' + // cloud function const VERTEX_AI_ENDPOINT = process.env.VERTEX_AI_ENDPOINT || '' + const VERTEX_CF_AUTH_TOKEN = process.env.VERTEX_CF_AUTH_TOKEN || '' + + // bigquery const VERTEX_BIGQUERY_LOOKER_CONNECTION_NAME = process.env.VERTEX_BIGQUERY_LOOKER_CONNECTION_NAME || '' const VERTEX_BIGQUERY_MODEL_ID = process.env.VERTEX_BIGQUERY_MODEL_ID || '' @@ -274,16 +279,21 @@ Output contents: string, parameters: ModelParameters, ) => { + const body = JSON.stringify({ + contents: contents, + parameters: parameters, + }) + + const signature = CryptoJS.HmacSHA256(body, VERTEX_CF_AUTH_TOKEN).toString() + const responseData = await fetch(VERTEX_AI_ENDPOINT, { method: 'POST', headers: { 'Content-Type': 'application/json', + 'X-Signature': signature, }, - body: JSON.stringify({ - contents: contents, - parameters: parameters, - }), + body: body, }) const response = await responseData.text() return response.trim()