From eeae4f6133a9c42fa6ea24c9e10b233f29978df5 Mon Sep 17 00:00:00 2001
From: Daniel Azuma <dazuma@google.com>
Date: Wed, 9 Oct 2024 00:01:01 +0000
Subject: [PATCH] feat: Disable universe domain check if credentials include an
 explicit request to do so

---
 gapic-common/.rubocop.yml                          | 12 +++---------
 gapic-common/lib/gapic/headers.rb                  |  5 +++++
 gapic-common/lib/gapic/universe_domain_concerns.rb |  4 ++--
 gapic-common/test/gapic/grpc/service_stub_test.rb  | 14 ++++++++++++++
 4 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/gapic-common/.rubocop.yml b/gapic-common/.rubocop.yml
index 6f1fb6038..d2f82c186 100644
--- a/gapic-common/.rubocop.yml
+++ b/gapic-common/.rubocop.yml
@@ -6,19 +6,13 @@ AllCops:
     - "test/**/*"
 
 Metrics/ClassLength:
-  Exclude:
-    - "lib/gapic/generic_lro/operation.rb"
-    - "lib/gapic/rest/grpc_transcoder.rb"
+  Max: 200
 
 Metrics/CyclomaticComplexity:
-  Exclude:
-    - "lib/gapic/headers.rb"
-    - "lib/gapic/rest/client_stub.rb"
+  Max: 15
 
 Metrics/PerceivedComplexity:
-  Exclude:
-    - "lib/gapic/headers.rb"
-    - "lib/gapic/rest/client_stub.rb"
+  Max: 15
 
 Naming/FileName:
   Exclude:
diff --git a/gapic-common/lib/gapic/headers.rb b/gapic-common/lib/gapic/headers.rb
index a5f9cb945..84b744e51 100644
--- a/gapic-common/lib/gapic/headers.rb
+++ b/gapic-common/lib/gapic/headers.rb
@@ -19,6 +19,8 @@
 module Gapic
   # A collection of common header values.
   module Headers
+    # rubocop:disable Metrics/CyclomaticComplexity,Metrics/PerceivedComplexity
+
     ##
     # @param ruby_version [String] The ruby version. Defaults to `RUBY_VERSION`.
     # @param lib_name [String] The client library name.
@@ -32,6 +34,7 @@ module Headers
     #     `:grpc` to send the GRPC library version (if defined)
     #     `:rest` to send the REST library version (if defined)
     #   Defaults to `[:grpc]`
+    #
     def self.x_goog_api_client ruby_version: nil, lib_name: nil, lib_version: nil, gax_version: nil,
                                gapic_version: nil, grpc_version: nil, rest_version: nil, protobuf_version: nil,
                                transports_version_send: [:grpc]
@@ -52,4 +55,6 @@ def self.x_goog_api_client ruby_version: nil, lib_name: nil, lib_version: nil, g
       x_goog_api_client_header.join " ".freeze
     end
   end
+
+  # rubocop:enable Metrics/CyclomaticComplexity,Metrics/PerceivedComplexity
 end
diff --git a/gapic-common/lib/gapic/universe_domain_concerns.rb b/gapic-common/lib/gapic/universe_domain_concerns.rb
index 2ef13d255..02cc048ec 100644
--- a/gapic-common/lib/gapic/universe_domain_concerns.rb
+++ b/gapic-common/lib/gapic/universe_domain_concerns.rb
@@ -60,11 +60,11 @@ def setup_universe_domain universe_domain: nil,
       raise ArgumentError, "endpoint or endpoint_template is required" if endpoint.nil? && endpoint_template.nil?
       raise ArgumentError, "credentials is required" if credentials.nil?
 
-      # TODO: The env logic should live in google-cloud-env
       universe_domain ||= ENV["GOOGLE_CLOUD_UNIVERSE_DOMAIN"] || "googleapis.com"
       endpoint ||= endpoint_template.sub ENDPOINT_SUBSTITUTION, universe_domain
 
-      if credentials.respond_to?(:universe_domain) && credentials.universe_domain != universe_domain
+      if credentials.respond_to?(:universe_domain) && credentials.universe_domain != universe_domain &&
+         !(credentials.respond_to?(:disable_universe_domain_check) && credentials.disable_universe_domain_check)
         raise UniverseDomainMismatch,
               "Universe domain is #{universe_domain} but credentials are in #{credentials.universe_domain}"
       end
diff --git a/gapic-common/test/gapic/grpc/service_stub_test.rb b/gapic-common/test/gapic/grpc/service_stub_test.rb
index ce0b02f21..0f7583213 100644
--- a/gapic-common/test/gapic/grpc/service_stub_test.rb
+++ b/gapic-common/test/gapic/grpc/service_stub_test.rb
@@ -52,6 +52,12 @@ def universe_domain
     end
   end
 
+  FakeCredentialsWithDisabledCheck = Class.new FakeCredentials do
+    def disable_universe_domain_check
+      true
+    end
+  end
+
   FakeRpcCall = Class.new do
     def initialize method_stub
       @method_stub = method_stub
@@ -270,6 +276,14 @@ def test_universe_domain_credentials_mismatch
     end
   end
 
+  def test_universe_domain_credentials_with_mismatch_disabled
+    creds = FakeCredentialsWithDisabledCheck.new universe_domain: "myuniverse.com"
+    service_stub = Gapic::ServiceStub.new FakeGrpcStub,
+                                          endpoint_template: "myservice.$UNIVERSE_DOMAIN$",
+                                          credentials: creds
+    assert_equal "googleapis.com", service_stub.universe_domain
+  end
+
   def test_endpoint_override
     creds = FakeCredentials.new universe_domain: "myuniverse.com"
     service_stub = Gapic::ServiceStub.new FakeGrpcStub,