From 85c09a1289ddc360b47575961461e968af81bcd6 Mon Sep 17 00:00:00 2001
From: John Pedrie <jdpedrie@users.noreply.github.com>
Date: Thu, 8 Aug 2019 12:22:36 -0400
Subject: [PATCH] Fix: Provide headers required for Firestore emulator. (#2239)

---
 Firestore/src/Connection/Grpc.php            | 35 +++++++++-----
 Firestore/tests/Unit/Connection/GrpcTest.php | 49 ++++++++++++++++++--
 2 files changed, 68 insertions(+), 16 deletions(-)

diff --git a/Firestore/src/Connection/Grpc.php b/Firestore/src/Connection/Grpc.php
index 3890153e79eb..0a29ed02b4b6 100644
--- a/Firestore/src/Connection/Grpc.php
+++ b/Firestore/src/Connection/Grpc.php
@@ -52,6 +52,11 @@ class Grpc implements ConnectionInterface
      */
     private $resourcePrefixHeader;
 
+    /**
+     * @var bool
+     */
+    private $isUsingEmulator = false;
+
     /**
      * @param array $config [optional]
      */
@@ -95,6 +100,7 @@ public function __construct(array $config = [])
         //@codeCoverageIgnoreStart
         $config += ['emulatorHost' => null];
         if ((bool) $config['emulatorHost']) {
+            $this->isUsingEmulator = true;
             $grpcConfig += $this->emulatorGapicConfig($config['emulatorHost']);
         }
         //@codeCoverageIgnoreEnd
@@ -115,7 +121,7 @@ public function batchGetDocuments(array $args)
         return $this->send([$this->firestore, 'batchGetDocuments'], [
             $this->pluck('database', $args),
             $this->pluck('documents', $args),
-            $this->addResourcePrefixHeader($args)
+            $this->addRequestHeaders($args)
         ]);
     }
 
@@ -135,7 +141,7 @@ public function beginTransaction(array $args)
 
         return $this->send([$this->firestore, 'beginTransaction'], [
             $this->pluck('database', $args),
-            $this->addResourcePrefixHeader($args)
+            $this->addRequestHeaders($args)
         ]);
     }
 
@@ -152,7 +158,7 @@ public function commit(array $args)
         return $this->send([$this->firestore, 'commit'], [
             $this->pluck('database', $args),
             $writes,
-            $this->addResourcePrefixHeader($args)
+            $this->addRequestHeaders($args)
         ]);
     }
 
@@ -163,7 +169,7 @@ public function listCollectionIds(array $args)
     {
         return $this->send([$this->firestore, 'listCollectionIds'], [
             $this->pluck('parent', $args),
-            $this->addResourcePrefixHeader($args)
+            $this->addRequestHeaders($args)
         ]);
     }
 
@@ -181,7 +187,7 @@ public function listDocuments(array $args)
         return $this->send([$this->firestore, 'listDocuments'], [
             $this->pluck('parent', $args),
             $this->pluck('collectionId', $args),
-            $this->addResourcePrefixHeader([
+            $this->addRequestHeaders([
                 'showMissing' => true
             ] + $args)
         ]);
@@ -195,7 +201,7 @@ public function rollback(array $args)
         return $this->send([$this->firestore, 'rollback'], [
             $this->pluck('database', $args),
             $this->pluck('transaction', $args),
-            $this->addResourcePrefixHeader($args)
+            $this->addRequestHeaders($args)
         ]);
     }
 
@@ -211,7 +217,7 @@ public function runQuery(array $args)
 
         return $this->send([$this->firestore, 'runQuery'], [
             $this->pluck('parent', $args),
-            $this->addResourcePrefixHeader($args)
+            $this->addRequestHeaders($args)
         ]);
     }
 
@@ -223,17 +229,24 @@ private function documentMask(array $mask)
     }
 
     /**
-     * Add the `google-cloud-resource-prefix` header value to the request.
+     * Add required headers to the request.
      *
      * @param array $args
      * @return array
      */
-    private function addResourcePrefixHeader(array $args)
+    private function addRequestHeaders(array $args)
     {
-        $args['headers'] = [
-            'google-cloud-resource-prefix' => [$this->resourcePrefixHeader]
+        $args += [
+            'headers' => []
         ];
 
+        $args['headers']['google-cloud-resource-prefix'] = [$this->resourcePrefixHeader];
+
+        // Provide authentication header for requests when emulator is enabled.
+        if ($this->isUsingEmulator) {
+            $args['headers']['Authorization'] = ['Bearer owner'];
+        }
+
         return $args;
     }
 }
diff --git a/Firestore/tests/Unit/Connection/GrpcTest.php b/Firestore/tests/Unit/Connection/GrpcTest.php
index 62c9109848a6..7e7696d1a019 100644
--- a/Firestore/tests/Unit/Connection/GrpcTest.php
+++ b/Firestore/tests/Unit/Connection/GrpcTest.php
@@ -20,6 +20,7 @@
 use Google\Cloud\Core\GrpcRequestWrapper;
 use Google\Cloud\Core\GrpcTrait;
 use Google\Cloud\Core\Testing\GrpcTestTrait;
+use Google\Cloud\Core\Testing\TestHelpers;
 use Google\Cloud\Firestore\Connection\Grpc;
 use Google\Cloud\Firestore\V1\Document;
 use Google\Cloud\Firestore\V1\DocumentMask;
@@ -239,6 +240,43 @@ public function testRunQuery()
         $this->sendAndAssert('runQuery', $args, $expected);
     }
 
+    public function testCustomRequestHeaders()
+    {
+        $args = [
+            'parent' => sprintf('projects/%s/databases/%s/documents', self::PROJECT, self::DATABASE),
+            'headers' => [
+                'foo' => ['bar']
+            ]
+        ];
+
+        $headers = $this->header();
+        $headers['headers']['foo'] = ['bar'];
+        $expected = [$args['parent'], $headers];
+
+        $this->sendAndAssert('listCollectionIds', $args, $expected);
+    }
+
+    public function testProvidesAuthorizationHeaderWithEmulator()
+    {
+        $args = [
+            'parent' => sprintf('projects/%s/databases/%s/documents', self::PROJECT, self::DATABASE),
+        ];
+
+        $headers = $this->header();
+        $headers['headers']['Authorization'] = ['Bearer owner'];
+        $expected = [$args['parent'], $headers];
+
+        $connection = TestHelpers::stub(Grpc::class, [
+            [
+                'projectId' => 'test',
+                'database' => '(default)'
+            ]
+        ], ['isUsingEmulator']);
+
+        $connection->___setProperty('isUsingEmulator', true);
+        $this->sendAndAssert('listCollectionIds', $args, $expected, $connection);
+    }
+
     private function header()
     {
         return [
@@ -248,18 +286,19 @@ private function header()
         ];
     }
 
-    private function sendAndAssert($method, array $args, array $expectedArgs)
+    private function sendAndAssert($method, array $args, array $expectedArgs, Grpc $connection = null)
     {
+        $connection = $connection ?: new Grpc([
+            'projectId' => 'test',
+            'database' => '(default)'
+        ]);
+
         $this->requestWrapper->send(
             Argument::type('callable'),
             $expectedArgs,
             Argument::type('array')
         )->willReturn($this->successMessage);
 
-        $connection = new Grpc([
-            'projectId' => 'test',
-            'database' => '(default)'
-        ]);
         $connection->setRequestWrapper($this->requestWrapper->reveal());
 
         $this->assertEquals($this->successMessage, $connection->$method($args));