Add support for AKS workload identity

.github/workflows/test.yml

@@ -35,7 +35,7 @@ jobs:
       - uses: actions/checkout@v4
       - uses: ./.github/actions/setup-terraform
       - name: Terraform apply
-        run: terraform apply -auto-approve -var "create_vm=true" -var "create_app_service=true" -var "ssh_key=${{ secrets.SSH_PUBLIC_KEY }}"
+        run: terraform apply -auto-approve -var "create_vm=true" -var "create_app_service=true" -var "create_aks=true" -var "ssh_key=${{ secrets.SSH_PUBLIC_KEY }}"
 
   app_service_test:
     needs: deploy-infrastructure
@@ -97,6 +97,45 @@ jobs:
           SSH_AUTH_SOCK: /tmp/ssh_agent.sock
         run: bundle exec rake test_azure_vm
 
+  aks_test:
+    needs: deploy-infrastructure
+    runs-on: ubuntu-latest
+    steps:
+      - name: Install dependencies
+        run: sudo apt-get update && sudo apt-get install -y libvips sshuttle sqlite3 libsqlite3-dev
+      - name: SSH key
+        env:
+          SSH_AUTH_SOCK: /tmp/ssh_agent.sock
+        run: |
+          mkdir -p /home/runner/.ssh
+          echo "${{ secrets.SSH_PRIVATE_KEY }}" > /home/runner/.ssh/id_rsa
+          chmod 600 /home/runner/.ssh/id_rsa
+          ssh-agent -a $SSH_AUTH_SOCK > /dev/null
+          ssh-add /home/runner/.ssh/id_rsa
+      - name: Checkout
+        uses: actions/checkout@v4
+      - uses: ./.github/actions/setup-terraform
+      - name: Azure login
+        uses: azure/login@v2
+        with:
+          client-id: ${{ secrets.AZURE_CLIENT_ID }}
+          tenant-id: ${{ secrets.AZURE_TENANT_ID }}
+          subscription-id: ${{ secrets.AZURE_SUBSCRIPTION_ID }}
+      - name: Setup kubectl
+        uses: azure/setup-kubectl@v4
+      - name: Setup ruby
+        uses: ruby/setup-ruby@v1
+        with:
+          bundler-cache: true
+      - name: Tests
+        env:
+          AZURE_ACCOUNT_NAME: ${{secrets.AZURE_ACCOUNT_NAME}}
+          AZURE_PRIVATE_CONTAINER: ${{secrets.AZURE_PRIVATE_CONTAINER}}
+          AZURE_PUBLIC_CONTAINER: ${{secrets.AZURE_PUBLIC_CONTAINER}}
+          USE_MANAGED_IDENTITIES: true
+          SSH_AUTH_SOCK: /tmp/ssh_agent.sock
+        run: bundle exec rake test_aks
+
   azurite_test:
     runs-on: ubuntu-latest
     steps:

Gemfile.lock

@@ -82,6 +82,7 @@ GEM
     base64 (0.2.0)
     bcrypt_pbkdf (1.1.1)
     bcrypt_pbkdf (1.1.1-arm64-darwin)
+    bcrypt_pbkdf (1.1.1-x86_64-darwin)
     benchmark (0.4.0)
     bigdecimal (3.1.9)
     builder (3.3.0)
@@ -97,6 +98,7 @@ GEM
     erb (5.0.1)
     erubi (1.13.1)
     ffi (1.17.2-arm64-darwin)
+    ffi (1.17.2-x86_64-darwin)
     ffi (1.17.2-x86_64-linux-gnu)
     globalid (1.2.1)
       activesupport (>= 6.1)
@@ -140,11 +142,13 @@ GEM
       timeout
     net-smtp (0.5.1)
       net-protocol
-    net-ssh (7.2.3)
+    net-ssh (7.3.0)
     nio4r (2.7.4)
-    nokogiri (1.18.8-arm64-darwin)
+    nokogiri (1.18.10-arm64-darwin)
       racc (~> 1.4)
-    nokogiri (1.18.8-x86_64-linux-gnu)
+    nokogiri (1.18.10-x86_64-darwin)
+      racc (~> 1.4)
+    nokogiri (1.18.10-x86_64-linux-gnu)
       racc (~> 1.4)
     parallel (1.24.0)
     parser (3.3.1.0)
@@ -244,6 +248,7 @@ GEM
       logger
     securerandom (0.4.1)
     sqlite3 (2.6.0-arm64-darwin)
+    sqlite3 (2.6.0-x86_64-darwin)
     sqlite3 (2.6.0-x86_64-linux-gnu)
     stringio (3.1.7)
     strscan (3.1.0)
@@ -262,6 +267,7 @@ GEM
 
 PLATFORMS
   arm64-darwin-24
+  x86_64-darwin-23
   x86_64-linux
 
 DEPENDENCIES

README.md

@@ -26,13 +26,13 @@ microsoft:
 
 ### Managed Identity (Entra ID)
 
-AzureBlob supports managed identities on :
+AzureBlob supports managed identities on:
 - Azure VM
 - App Service
+- AKS (Azure Kubernetes Service) - supports both node identity and workload identity
 - Azure Functions (Untested but should work)
 - Azure Containers (Untested but should work)
-
-AKS support will likely require more work. Contributions are welcome.
+- Azure AD Workload Identity (AKS/K8s)
 
 To authenticate through managed identities instead of a shared key, omit `storage_access_key` from your `storage.yml` file and pass in the identity `principal_id`.
 
@@ -46,6 +46,21 @@ prod:
   principal_id: 71b34410-4c50-451d-b456-95ead1b18cce
 ```
 
+#### Azure AD Workload Identity (AKS/K8s)
+
+ActiveStorage config example:
+
+```
+prod:
+  service: AzureBlob
+  container: container_name
+  storage_account_name: account_name
+  use_managed_identities: true
+```
+
+> uses `AZURE_CLIENT_ID`, `AZURE_TENANT_ID` and `AZURE_FEDERATED_TOKEN_FILE` environment variables, made available by AKS cluster when Azure AD Workload Identity is set up properly.
+
+
 ### Azurite
 
 To use Azurite, pass the `storage_blob_host` config key with the Azurite URL (`http://127.0.0.1:10000/devstoreaccount1` by default)
@@ -126,20 +141,34 @@ A dev environment is supplied through Nix with [devenv](https://devenv.sh/).
 
 To test with Entra ID, the `AZURE_ACCESS_KEY` environment variable must be unset and the code must be ran or proxied through a VPS with the proper roles.
 
-For cost saving, the terraform variable `create_vm` and `create_app_service` are false by default.
-To create the VPS and App service, Create a var file `var.tfvars` containing:
+For cost saving, the terraform variables `create_vm`, `create_app_service`, and `create_aks` are false by default.
+To create the VM, App Service, and/or AKS cluster, create a var file `var.tfvars` containing:
 
 ```
 create_vm = true
 create_app_service = true
+create_aks = true
 ```
 and re-apply terraform: `terraform apply -var-file=var.tfvars`.
 
-This will create the VPS and required managed identities.
+This will create the infrastructure and required managed identities.
+
+**Testing:**
+- `bin/rake test_azure_vm` - Establishes a VPN connection to the Azure VM and runs tests using node identity
+- `bin/rake test_app_service` - Establishes a VPN connection to the App Service container and runs tests
+- `bin/rake test_aks` - Establishes a VPN connection to the AKS cluster and runs tests. This tests both:
+  - **Node identity**: Storage access via the node pool's managed identity
+  - **Workload identity**: Storage access via Azure AD federated credentials for the pod's service account
+
+You might be prompted for a sudo password when the VPN starts (sshuttle).
 
-`bin/rake test_azure_vm` and `bin/rake test_app_service` will establish a VPN connection to the VM or App service container and run the test suite. You might be prompted for a sudo password when the VPN starts (sshuttle).
+After you are done, run terraform again without the var file (`terraform apply`) to destroy all resources.
 
-After you are done, run terraform again without the var file (`terraform apply`) to destroy the VPS and App service application.
+**Note for AKS:** The AKS setup includes:
+- A Kubernetes cluster with OIDC issuer and workload identity enabled
+- An SSH-enabled pod running the `linuxserver/openssh-server` image
+- A LoadBalancer service exposing SSH access for VPN tunneling
+- Both kubelet identity (for node-level access) and federated identity credentials (for workload identity)
 
 #### Cleanup
 

Rakefile

@@ -5,6 +5,7 @@ require "minitest/test_task"
 require "azure_blob"
 require_relative "test/support/app_service_vpn"
 require_relative "test/support/azure_vm_vpn"
+require_relative "test/support/aks_vpn"
 require_relative "test/support/azurite"
 
 Minitest::TestTask.create(:test_rails) do
@@ -49,6 +50,15 @@ ensure
   vpn.kill
 end
 
+task :test_aks do |t|
+  vpn = AksVpn.new
+  ENV["IDENTITY_ENDPOINT"] = vpn.endpoint
+  ENV["IDENTITY_HEADER"] = vpn.header
+  Rake::Task["test_entra_id"].execute
+ensure
+  vpn.kill
+end
+
 task :test_azurite do |t|
   azurite = Azurite.new
   # Azurite well-known credentials

bin/proxy-aks

@@ -0,0 +1,34 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Get Terraform outputs
+resource_group=$(terraform output --raw "resource_group")
+cluster_name=$(terraform output --raw "aks_cluster_name")
+ssh_username=$(terraform output --raw "aks_ssh_username")
+
+# Setup kubectl config
+echo "Setting up kubectl config..."
+az aks get-credentials --resource-group "$resource_group" --name "$cluster_name" --overwrite-existing --only-show-errors 2>/dev/null
+
+# Wait for pod to be ready
+echo "Waiting for SSH pod to be ready..."
+kubectl wait --for=condition=ready pod -l app=azure-blob-test --timeout=300s
+
+# Start port forward in the background
+echo "Setting up port forward..."
+kubectl port-forward service/azure-blob-test-ssh 2222:22 &
+PORT_FORWARD_PID=$!
+
+# Cleanup function
+cleanup() {
+  echo "Cleaning up..."
+  kill $PORT_FORWARD_PID 2>/dev/null || true
+}
+trap cleanup EXIT
+
+# Wait for port forward to be established
+sleep 3
+
+# Establish sshuttle VPN
+echo "Establishing VPN connection..."
+exec sshuttle -e "ssh -o CheckHostIP=no -o StrictHostKeyChecking=no -o UserKnownHostsFile=/dev/null" -r "[email protected]:2222" 0/0

devenv.nix

@@ -12,10 +12,10 @@
     libyaml
     terraform
     azure-cli
+    kubectl
     glib
     vips
     sshuttle
-    sshpass
     rsync
     azurite
     sqlite

input.tf

@@ -38,6 +38,21 @@ variable "create_app_service" {
   default = false
 }
 
+variable "create_aks" {
+  type    = bool
+  default = false
+}
+
+variable "aks_node_count" {
+  type    = number
+  default = 1
+}
+
+variable "aks_ssh_username" {
+  type    = string
+  default = "testuser"
+}
+
 variable "ssh_key" {
   type    = string
   default = ""

lib/azure_blob/identity_token.rb

@@ -1,21 +1,14 @@
+require_relative "instance_metadata_service"
+require_relative "workload_identity"
 require "json"
 
 module AzureBlob
   class IdentityToken
-    RESOURCE_URI = "https://storage.azure.com/"
     EXPIRATION_BUFFER = 600 # 10 minutes
 
-    IDENTITY_ENDPOINT = ENV["IDENTITY_ENDPOINT"] || "http://169.254.169.254/metadata/identity/oauth2/token"
-    API_VERSION = ENV["IDENTITY_ENDPOINT"] ? "2019-08-01" : "2018-02-01"
-
     def initialize(principal_id: nil)
-      @identity_uri = URI.parse(IDENTITY_ENDPOINT)
-      params = {
-        'api-version': API_VERSION,
-        resource: RESOURCE_URI,
-      }
-      params[:principal_id] = principal_id if principal_id
-      @identity_uri.query = URI.encode_www_form(params)
+      @service = AzureBlob::WorkloadIdentity.federated_token? ?
+                   AzureBlob::WorkloadIdentity.new : AzureBlob::InstanceMetadataService.new(principal_id: principal_id)
     end
 
     def to_s
@@ -31,13 +24,11 @@ def expired?
 
     def refresh
       return unless expired?
-      headers =  { "Metadata" => "true" }
-      headers["X-IDENTITY-HEADER"] = ENV["IDENTITY_HEADER"] if ENV["IDENTITY_HEADER"]
 
       attempt = 0
       begin
         attempt += 1
-        response = JSON.parse(AzureBlob::Http.new(identity_uri, headers).get)
+        response = JSON.parse(service.request)
       rescue AzureBlob::Http::Error => error
         if should_retry?(error, attempt)
           attempt = 1 if error.status == 410
@@ -48,7 +39,7 @@ def refresh
         raise
       end
       @token = response["access_token"]
-      @expiration = Time.at(response["expires_on"].to_i)
+      @expiration = Time.at((response["expires_on"] || response["expires_in"]).to_i)
     end
 
     def should_retry?(error, attempt)
@@ -61,6 +52,6 @@ def exponential_backoff(error, attempt)
     end
     EXPONENTIAL_BACKOFF = [ 2, 6, 14, 30 ]
 
-    attr_reader :identity_uri, :expiration, :token
+    attr_reader :service, :expiration, :token
   end
 end