From 1026d8aec73e0aad8949a86ee7575e3edb3d56be Mon Sep 17 00:00:00 2001
From: Brent Shaffer <betterbrent@google.com>
Date: Fri, 16 Feb 2024 13:38:55 -0700
Subject: [PATCH] fix: universe domain check for grpc transport (#534)

---
 src/CredentialsWrapper.php                       |  2 +-
 src/Transport/GrpcTransport.php                  | 16 ++++++++++++++++
 tests/Tests/Unit/Transport/GrpcTransportTest.php |  2 ++
 3 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/CredentialsWrapper.php b/src/CredentialsWrapper.php
index 3ea2c8ed8..6988188fa 100644
--- a/src/CredentialsWrapper.php
+++ b/src/CredentialsWrapper.php
@@ -272,7 +272,7 @@ public function getAuthorizationHeaderCallback($audience = null)
     /**
      * Verify that the expected universe domain matches the universe domain from the credentials.
      */
-    private function checkUniverseDomain()
+    public function checkUniverseDomain()
     {
         if (false === $this->hasCheckedUniverse) {
             $credentialsUniverse = $this->credentialsFetcher instanceof GetUniverseDomainInterface
diff --git a/src/Transport/GrpcTransport.php b/src/Transport/GrpcTransport.php
index d71eae272..3bc3328e8 100644
--- a/src/Transport/GrpcTransport.php
+++ b/src/Transport/GrpcTransport.php
@@ -159,6 +159,8 @@ public static function build(string $apiEndpoint, array $config = [])
      */
     public function startBidiStreamingCall(Call $call, array $options)
     {
+        $this->verifyUniverseDomain($options);
+
         return new BidiStream(
             $this->_bidiRequest(
                 '/' . $call->getMethod(),
@@ -175,6 +177,9 @@ public function startBidiStreamingCall(Call $call, array $options)
      */
     public function startClientStreamingCall(Call $call, array $options)
     {
+
+        $this->verifyUniverseDomain($options);
+
         return new ClientStream(
             $this->_clientStreamRequest(
                 '/' . $call->getMethod(),
@@ -191,6 +196,8 @@ public function startClientStreamingCall(Call $call, array $options)
      */
     public function startServerStreamingCall(Call $call, array $options)
     {
+        $this->verifyUniverseDomain($options);
+
         $message = $call->getMessage();
 
         if (!$message) {
@@ -216,6 +223,8 @@ public function startServerStreamingCall(Call $call, array $options)
      */
     public function startUnaryCall(Call $call, array $options)
     {
+        $this->verifyUniverseDomain($options);
+
         $unaryCall = $this->_simpleRequest(
             '/' . $call->getMethod(),
             $call->getMessage(),
@@ -245,6 +254,13 @@ function () use ($unaryCall, $options, &$promise) {
         return $promise;
     }
 
+    private function verifyUniverseDomain(array $options)
+    {
+        if (isset($options['credentialsWrapper'])) {
+            $options['credentialsWrapper']->checkUniverseDomain();
+        }
+    }
+
     private function getCallOptions(array $options)
     {
         $callOptions = $options['transportOptions']['grpcOptions'] ?? [];
diff --git a/tests/Tests/Unit/Transport/GrpcTransportTest.php b/tests/Tests/Unit/Transport/GrpcTransportTest.php
index 8464460dd..9d8e37d70 100644
--- a/tests/Tests/Unit/Transport/GrpcTransportTest.php
+++ b/tests/Tests/Unit/Transport/GrpcTransportTest.php
@@ -387,6 +387,8 @@ public function testAudienceOption()
         $call->getDecodeType()->shouldBeCalled();
 
         $credentialsWrapper = $this->prophesize(CredentialsWrapper::class);
+        $credentialsWrapper->checkUniverseDomain()
+            ->shouldBeCalledOnce();
         $credentialsWrapper->getAuthorizationHeaderCallback('an-audience')
             ->shouldBeCalledOnce();
         $hostname = '';