feat: Enabled workload identity access for kms and gcs

README.md

@@ -112,6 +112,7 @@ resources that lack official modules.
 | <a name="input_force_ssl"></a> [force\_ssl](#input\_force\_ssl) | Enforce SSL through the usage of the Cloud SQL Proxy (cloudsql://) in the DB connection string | `bool` | `false` | no |
 | <a name="input_gke_machine_type"></a> [gke\_machine\_type](#input\_gke\_machine\_type) | Specifies the machine type to be allocated for the database | `string` | `"n1-standard-4"` | no |
 | <a name="input_gke_node_count"></a> [gke\_node\_count](#input\_gke\_node\_count) | n/a | `number` | `2` | no |
+| <a name="input_kms_gcs_sa_name"></a> [kms\_gcs\_sa\_name](#input\_kms\_gcs\_sa\_name) | n/a | `string` | `"wandb-serviceaccount"` | no |
 | <a name="input_labels"></a> [labels](#input\_labels) | Labels to apply to resources | `map(string)` | `{}` | no |
 | <a name="input_license"></a> [license](#input\_license) | Your wandb/local license | `string` | n/a | yes |
 | <a name="input_local_restore"></a> [local\_restore](#input\_local\_restore) | Restores W&B to a stable state if needed | `bool` | `false` | no |
@@ -157,6 +158,7 @@ resources that lack official modules.
 | <a name="output_fqdn"></a> [fqdn](#output\_fqdn) | The FQDN to the W&B application |
 | <a name="output_gke_node_count"></a> [gke\_node\_count](#output\_gke\_node\_count) | n/a |
 | <a name="output_gke_node_instance_type"></a> [gke\_node\_instance\_type](#output\_gke\_node\_instance\_type) | n/a |
+| <a name="output_sa_account_email"></a> [sa\_account\_email](#output\_sa\_account\_email) | This output provides the email address of the service account created for workload identity, if workload identity is enabled. Otherwise, it returns null |
 | <a name="output_service_account"></a> [service\_account](#output\_service\_account) | Weights & Biases service account used to manage resources. |
 | <a name="output_standardized_size"></a> [standardized\_size](#output\_standardized\_size) | n/a |
 | <a name="output_url"></a> [url](#output\_url) | The URL to the W&B application |

main.tf

@@ -29,12 +29,14 @@ locals {
 }
 
 module "service_accounts" {
- source = "./modules/service_accounts"
- namespace = var.namespace
- bucket_name = var.bucket_name
- stackdriver_sa_name = var.stackdriver_sa_name
- enable_stackdriver = var.enable_stackdriver
- depends_on = [module.project_factory_project_services]
+ source = "./modules/service_accounts"
+ namespace = var.namespace
+ bucket_name = var.bucket_name
+ kms_gcs_sa_name = var.kms_gcs_sa_name
+ create_workload_identity = var.create_workload_identity
+ stackdriver_sa_name = var.stackdriver_sa_name
+ enable_stackdriver = var.enable_stackdriver
+ depends_on = [module.project_factory_project_services]
 }
 
 module "kms" {
@@ -150,7 +152,10 @@ module "gke_app" {
  database_connection_string = module.database.connection_string
  redis_connection_string = local.redis_connection_string
  redis_ca_cert = local.redis_certificate
-
+ service_account_name = var.kms_gcs_sa_name
+ service_account_annotations = var.create_workload_identity ? {
+ "iam.gke.io/gcp-service-account" : module.service_accounts.sa_account_email
+ } : {}
  oidc_client_id = var.oidc_client_id
  oidc_issuer = var.oidc_issuer
  oidc_auth_method = var.oidc_auth_method
@@ -246,6 +251,15 @@ module "wandb" {
 
  app = {
  extraEnvs = var.app_wandb_env
+ serviceaccount = var.create_workload_identity ? {
+ create = false
+ name = var.kms_gcs_sa_name
+ annotations = { "iam.gke.io/gcp-service-account" = module.service_accounts.sa_account_email }
+ } : {
+ create = true
+ name = ""
+ annotations = {}
+ }
  }
 
  ingress = {

modules/app_gke/main.tf

@@ -7,12 +7,12 @@ locals {
 resource "google_container_cluster" "default" {
  name = "${var.namespace}-cluster"
 
- network = var.network.self_link
- subnetwork = var.subnetwork.self_link
- networking_mode = "VPC_NATIVE"
-
+ network = var.network.self_link
+ subnetwork = var.subnetwork.self_link
+ networking_mode = "VPC_NATIVE"
  enable_intranode_visibility = true
 
+
  binary_authorization {
  evaluation_mode = "PROJECT_SINGLETON_POLICY_ENFORCE"
  }
@@ -24,7 +24,7 @@ resource "google_container_cluster" "default" {
  workload_pool = "${local.project_id}.svc.id.goog"
  }
  }
- 
+
  ip_allocation_policy {
  cluster_ipv4_cidr_block = "/14"
  services_ipv4_cidr_block = "/19"
@@ -78,12 +78,13 @@ resource "google_container_node_pool" "default" {
  "https://www.googleapis.com/auth/sqlservice.admin",
  ]
 
- dynamic "workload_metadata_config" {
- for_each = var.create_workload_identity == true ? [1] : []
- content {
- mode = "GKE_METADATA"
+ dynamic "workload_metadata_config" {
+ for_each = var.create_workload_identity == true ? [1] : []
+ content {
+ mode = "GKE_METADATA"
+ }
  }
- }
+
  shielded_instance_config {
  enable_secure_boot = true
  }

modules/service_accounts/main.tf

@@ -62,7 +62,43 @@ resource "google_project_iam_member" "secretmanager_admin" {
  role = "roles/secretmanager.admin"
 }
 
+####### service account for kms and gcs cross project access
+resource "google_service_account" "kms_gcs_sa" {
+ count = var.create_workload_identity == true ? 1 : 0
+ account_id = var.kms_gcs_sa_name
+ display_name = "Service Account For Workload Identity"
+}
+
+resource "google_project_iam_member" "storage" {
+ count = var.create_workload_identity == true ? 1 : 0
+ project = local.project_id
+ role = "roles/storage.admin"
+ member = "serviceAccount:${google_service_account.kms_gcs_sa[count.index].email}"
+}
+
+resource "google_project_iam_member" "kms" {
+ count = var.create_workload_identity == true ? 1 : 0
+ project = local.project_id
+ role = "roles/cloudkms.admin"
+ member = "serviceAccount:${google_service_account.kms_gcs_sa[count.index].email}"
+}
+
+resource "google_service_account_iam_member" "token_creator_binding" {
+ count = var.create_workload_identity == true ? 1 : 0
+ service_account_id = google_service_account.kms_gcs_sa[count.index].id
+ role = "roles/iam.serviceAccountTokenCreator"
+ member = "serviceAccount:${google_service_account.kms_gcs_sa[count.index].email}"
+}
+
+resource "google_service_account_iam_member" "workload_binding" {
+ count = var.create_workload_identity == true ? 1 : 0
+ service_account_id = google_service_account.kms_gcs_sa[count.index].id
+ role = "roles/iam.workloadIdentityUser"
+ member = "serviceAccount:${local.project_id}.svc.id.goog[default/${var.kms_gcs_sa_name}]"
+}
+
 
+### service account for stackdriver
 resource "google_service_account" "stackdriver" {
  count = var.enable_stackdriver == true ? 1 : 0
  account_id = var.stackdriver_sa_name

modules/service_accounts/outputs.tf

@@ -1,9 +1,12 @@
 output "service_account" {
  value = google_service_account.main
-
  description = "The service account."
 }
 
+output "sa_account_email" {
+ value = var.create_workload_identity == true ? google_service_account.kms_gcs_sa[0].email : null
+}
+
 output "stackdriver_email" {
  value = var.enable_stackdriver == true ? google_service_account.stackdriver[0].email : null
 }

modules/service_accounts/variables.tf

@@ -10,9 +10,20 @@ variable "bucket_name" {
 }
 
 
+variable "create_workload_identity" {
+ description = "Flag to indicate whether to create a workload identity for the service account."
+ type = bool
+}
+
+variable "kms_gcs_sa_name" {
+ type = string
+ default = "wandb-serviceaccount"
+}
+
 variable "stackdriver_sa_name" {
  description = "The name of the service account."
  type = string
+ default = "wandb-stackdriver"
 }
 
 variable "enable_stackdriver" {

outputs.tf

@@ -85,5 +85,7 @@ output "database_instance_type" {
  value = coalesce(try(local.deployment_size[var.size].db, null), var.database_machine_type)
 }
 
-
-
+output "sa_account_email" {
+ description = "This output provides the email address of the service account created for workload identity, if workload identity is enabled. Otherwise, it returns null"
+ value = var.create_workload_identity == true ? module.service_accounts.sa_account_email : null
+}

variables.tf

@@ -254,18 +254,24 @@ variable "parquet_wandb_env" {
  default = {}
 }
 
-variable "enable_stackdriver" {
- type = bool
- default = true
-}
-
 variable "create_workload_identity" {
  description = "Flag to indicate whether to create a workload identity for the service account."
  type = bool
  default = true
 }
 
+variable "kms_gcs_sa_name" {
+ type = string
+ default = "wandb-serviceaccount"
+}
+
+variable "enable_stackdriver" {
+ type = bool
+ default = true
+}
+
 variable "stackdriver_sa_name" {
  type = string
  default = "wandb-stackdriver"
 }
+