Add support for AKS workload identity

Author: arashnd

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)
@@ -142,9 +144,11 @@ GEM
       net-protocol
     net-ssh (7.2.3)
     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

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

lib/azure_blob/instance_metadata_service.rb

@@ -0,0 +1,24 @@
+module AzureBlob
+  class InstanceMetadataService # Azure Instance Metadata Service (IMDS)
+    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"
+    RESOURCE_URI = "https://storage.azure.com/"
+
+    def initialize(principal_id: nil)
+      @identity_uri = URI.parse(IDENTITY_ENDPOINT)
+      params = {
+        'api-version': API_VERSION,
+        resource: AzureBlob::IdentityToken::RESOURCE_URI
+      }
+      params[:principal_id] = principal_id if principal_id
+      @identity_uri.query = URI.encode_www_form(params)
+    end
+
+    def request
+      headers =  { "Metadata" => "true" }
+      headers["X-IDENTITY-HEADER"] = ENV["IDENTITY_HEADER"] if ENV["IDENTITY_HEADER"]
+
+      AzureBlob::Http.new(@identity_uri, headers).get
+    end
+  end
+end

lib/azure_blob/workload_identity.rb

@@ -0,0 +1,33 @@
+module AzureBlob
+  class WorkloadIdentity # Azure AD Workload Identity
+    IDENTITY_ENDPOINT = "https://login.microsoftonline.com/#{ENV['AZURE_TENANT_ID']}/oauth2/v2.0/token"
+    CLIENT_ID = ENV['AZURE_CLIENT_ID']
+    SCOPE = "https://storage.azure.com/.default"
+    GRANT_TYPE = "client_credentials"
+    CLIENT_ASSERTION_TYPE = "urn:ietf:params:oauth:client-assertion-type:jwt-bearer"
+
+    FEDERATED_TOKEN_FILE = ENV["AZURE_FEDERATED_TOKEN_FILE"].to_s
+
+    def self.federated_token?
+      !FEDERATED_TOKEN_FILE.empty?
+    end
+
+    def request
+      AzureBlob::Http.new(URI.parse(IDENTITY_ENDPOINT)).post(
+        URI.encode_www_form(
+          client_id: CLIENT_ID,
+          scope: SCOPE,
+          client_assertion_type: CLIENT_ASSERTION_TYPE,
+          client_assertion: federated_token,
+          grant_type: GRANT_TYPE
+        )
+      )
+    end
+
+    private
+
+    def federated_token
+      File.read(FEDERATED_TOKEN_FILE).strip
+    end
+  end
+end