Fix storage content-type on desktop, and add test. (#911)

jonsimantov · web-flow · commit 0a0c4fc22961 · 2022-04-25T15:23:31.000-07:00
* Fix storage content-type on desktop, and add test.

* Fix data thread access.

* Set content_type in a consistent way.
diff --git a/release_build_files/readme.md b/release_build_files/readme.md
@@ -567,6 +567,11 @@ workflow use only during the development of your app, not for publicly shipping
 code.
 
 ## Release Notes
+### Upcoming release
+-   Changes
+    -   Storage (Desktop): Set Content-Type HTTP header when uploading with
+        custom metadata.
+
 ### 8.11.0
 -   Changes
     -   Firestore/Database (Desktop): Upgrade LevelDb dependency to 1.23
diff --git a/storage/integration_test/src/integration_test.cc b/storage/integration_test/src/integration_test.cc
@@ -509,6 +509,66 @@ TEST_F(FirebaseStorageTest, TestWriteAndReadFileWithCustomMetadata) {
   }
 }
 
+// 1x1 transparent PNG file
+static const unsigned char kEmptyPngFileBytes[] = {
+    0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a, 0x00, 0x00, 0x00, 0x0d,
+    0x49, 0x48, 0x44, 0x52, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01,
+    0x08, 0x06, 0x00, 0x00, 0x00, 0x1f, 0x15, 0xc4, 0x89, 0x00, 0x00, 0x00,
+    0x0d, 0x49, 0x44, 0x41, 0x54, 0x78, 0xda, 0x63, 0xfc, 0xcf, 0xc0, 0x50,
+    0x0f, 0x00, 0x04, 0x85, 0x01, 0x80, 0x84, 0xa9, 0x8c, 0x21, 0x00, 0x00,
+    0x00, 0x00, 0x49, 0x45, 0x4e, 0x44, 0xae, 0x42, 0x60, 0x82};
+
+TEST_F(FirebaseStorageTest, TestWriteAndReadCustomContentType) {
+  SignIn();
+
+  firebase::storage::StorageReference ref =
+      CreateFolder().Child("TestFile-CustomContentType.png");
+  LogDebug("Storage URL: gs://%s%s", ref.bucket().c_str(),
+           ref.full_path().c_str());
+  cleanup_files_.push_back(ref);
+  std::string content_type = "image/png";
+  // Write to a simple file.
+  {
+    LogDebug("Write a sample file with custom content-type from byte buffer.");
+    firebase::storage::Metadata metadata;
+    metadata.set_content_type(content_type.c_str());
+    firebase::Future<firebase::storage::Metadata> future =
+        ref.PutBytes(kEmptyPngFileBytes, sizeof(kEmptyPngFileBytes), metadata);
+    WaitForCompletion(future, "PutBytes");
+    const firebase::storage::Metadata* metadata_written = future.result();
+    ASSERT_NE(metadata_written, nullptr);
+    EXPECT_EQ(metadata_written->content_type(), content_type);
+  }
+  // Now read back the file.
+  {
+    LogDebug("Download sample file with custom content-type to memory.");
+    const size_t kBufferSize = 1024;
+    char buffer[kBufferSize];
+    memset(buffer, 0, sizeof(buffer));
+
+    firebase::Future<size_t> future = RunWithRetry<size_t>(
+        [&]() { return ref.GetBytes(buffer, kBufferSize); });
+    WaitForCompletion(future, "GetBytes");
+    ASSERT_NE(future.result(), nullptr);
+    size_t file_size = *future.result();
+    EXPECT_EQ(file_size, sizeof(kEmptyPngFileBytes));
+    EXPECT_THAT(kEmptyPngFileBytes, ElementsAreArray(buffer, file_size))
+        << "Download failed, file contents did not match.";
+  }
+  // And read the custom content type
+  {
+    LogDebug("Read custom content-type.");
+    firebase::Future<firebase::storage::Metadata> future =
+        RunWithRetry<firebase::storage::Metadata>(
+            [&]() { return ref.GetMetadata(); });
+    WaitForCompletion(future, "GetFileMetadata");
+    const firebase::storage::Metadata* metadata = future.result();
+    ASSERT_NE(metadata, nullptr);
+
+    EXPECT_EQ(metadata->content_type(), content_type);
+  }
+}
+
 const char kPutFileTestFile[] = "PutFileTest.txt";
 const char kGetFileTestFile[] = "GetFileTest.txt";
 const char kFileUriScheme[] = "file://";
diff --git a/storage/src/desktop/storage_reference_desktop.cc b/storage/src/desktop/storage_reference_desktop.cc
@@ -230,7 +230,8 @@ Future<void> StorageReferenceInternal::DeleteLastResult() {
 // Handy utility function, since REST calls have similar setup and teardown.
 void StorageReferenceInternal::PrepareRequest(rest::Request* request,
                                               const char* url,
-                                              const char* method) {
+                                              const char* method,
+                                              const char* content_type) {
   request->set_url(url);
   request->set_method(method);
 
@@ -240,6 +241,10 @@ void StorageReferenceInternal::PrepareRequest(rest::Request* request,
     std::string auth_header = "Bearer " + token;
     request->add_header("Authorization", auth_header.c_str());
   }
+  // if content_type was specified, add a header.
+  if (content_type != nullptr && *content_type != '\0') {
+    request->add_header("Content-Type", content_type);
+  }
   // Unfortunately the storage backend rejects requests with the complete
   // user agent specified by the x-goog-api-client header so we only use
   // the X-Firebase-Storage-Version header to attribute the client.
@@ -404,10 +409,12 @@ Future<Metadata> StorageReferenceInternal::PutBytes(
 
 Future<Metadata> StorageReferenceInternal::PutBytesInternal(
     const void* buffer, size_t buffer_size, Listener* listener,
-    Controller* controller_out) {
+    Controller* controller_out, const char* content_type) {
   auto* future_api = future();
   auto handle = future_api->SafeAlloc<Metadata>(kStorageReferenceFnPutBytes);
-  auto send_request_funct{[&, buffer, buffer_size, listener,
+
+  std::string content_type_str = content_type ? content_type : "";
+  auto send_request_funct{[&, content_type_str, buffer, buffer_size, listener,
                            controller_out]() -> BlockingResponse* {
     auto* future_api = future();
     auto handle =
@@ -416,7 +423,8 @@ Future<Metadata> StorageReferenceInternal::PutBytesInternal(
     storage::internal::RequestBinary* request =
         new storage::internal::RequestBinary(static_cast<const char*>(buffer),
                                              buffer_size);
-    PrepareRequest(request, storageUri_.AsHttpUrl().c_str(), rest::util::kPost);
+    PrepareRequest(request, storageUri_.AsHttpUrl().c_str(), rest::util::kPost,
+                   content_type_str.c_str());
     ReturnedMetadataResponse* response =
         new ReturnedMetadataResponse(handle, future_api, AsStorageReference());
     RestCall(request, request->notifier(), response, handle.get(), listener,
@@ -442,8 +450,9 @@ Future<Metadata> StorageReferenceInternal::PutBytes(
   // different storage reference than the original, so the caller of this
   // function can't access it via PutFileLastResult.
   Future<Metadata> putbytes_internal =
-      data->storage_ref.internal_->PutBytesInternal(buffer, buffer_size,
-                                                    listener, controller_out);
+      data->storage_ref.internal_->PutBytesInternal(
+          buffer, buffer_size, listener, controller_out,
+          metadata ? metadata->content_type() : nullptr);
 
   SetupMetadataChain(putbytes_internal, data);
 
@@ -466,36 +475,38 @@ Future<Metadata> StorageReferenceInternal::PutFile(const char* path,
 }
 
 Future<Metadata> StorageReferenceInternal::PutFileInternal(
-    const char* path, Listener* listener, Controller* controller_out) {
+    const char* path, Listener* listener, Controller* controller_out,
+    const char* content_type) {
   auto* future_api = future();
   auto handle = future_api->SafeAlloc<Metadata>(kStorageReferenceFnPutFile);
 
   std::string final_path = StripProtocol(path);
-  auto send_request_funct{
-      [&, final_path, listener, controller_out]() -> BlockingResponse* {
-        auto* future_api = future();
-        auto handle =
-            future_api->SafeAlloc<Metadata>(kStorageReferenceFnPutFileInternal);
-
-        // Open the file, calculate the length.
-        storage::internal::RequestFile* request(
-            new storage::internal::RequestFile(final_path.c_str(), 0));
-        if (!request->IsFileOpen()) {
-          delete request;
-          future_api->Complete(handle, kErrorUnknown, "Could not read file.");
-          return nullptr;
-        } else {
-          // Everything is good.  Fire off the request.
-          ReturnedMetadataResponse* response = new ReturnedMetadataResponse(
-              handle, future_api, AsStorageReference());
-
-          PrepareRequest(request, storageUri_.AsHttpUrl().c_str(),
-                         rest::util::kPost);
-          RestCall(request, request->notifier(), response, handle.get(),
-                   listener, controller_out);
-          return response;
-        }
-      }};
+  std::string content_type_str = content_type ? content_type : "";
+  auto send_request_funct{[&, final_path, content_type_str, listener,
+                           controller_out]() -> BlockingResponse* {
+    auto* future_api = future();
+    auto handle =
+        future_api->SafeAlloc<Metadata>(kStorageReferenceFnPutFileInternal);
+
+    // Open the file, calculate the length.
+    storage::internal::RequestFile* request(
+        new storage::internal::RequestFile(final_path.c_str(), 0));
+    if (!request->IsFileOpen()) {
+      delete request;
+      future_api->Complete(handle, kErrorUnknown, "Could not read file.");
+      return nullptr;
+    } else {
+      // Everything is good.  Fire off the request.
+      ReturnedMetadataResponse* response = new ReturnedMetadataResponse(
+          handle, future_api, AsStorageReference());
+
+      PrepareRequest(request, storageUri_.AsHttpUrl().c_str(),
+                     rest::util::kPost, content_type_str.c_str());
+      RestCall(request, request->notifier(), response, handle.get(), listener,
+               controller_out);
+      return response;
+    }
+  }};
   SendRequestWithRetry(kStorageReferenceFnPutFileInternal, send_request_funct,
                        handle, storage_->max_upload_retry_time());
   return PutFileLastResult();
@@ -516,8 +527,9 @@ Future<Metadata> StorageReferenceInternal::PutFile(const char* path,
   // different storage reference than the original, so the caller of this
   // function can't access it via PutFileLastResult.
   Future<Metadata> putfile_internal =
-      data->storage_ref.internal_->PutFileInternal(path, listener,
-                                                   controller_out);
+      data->storage_ref.internal_->PutFileInternal(
+          path, listener, controller_out,
+          metadata ? metadata->content_type() : nullptr);
 
   SetupMetadataChain(putfile_internal, data);
 
@@ -579,11 +591,11 @@ Future<Metadata> StorageReferenceInternal::UpdateMetadata(
         new ReturnedMetadataResponse(handle, future_api, AsStorageReference());
 
     storage::internal::Request* request = new storage::internal::Request();
-    PrepareRequest(request, storageUri_.AsHttpUrl().c_str(), "PATCH");
+    PrepareRequest(request, storageUri_.AsHttpUrl().c_str(), "PATCH",
+                   "application/json");
 
     std::string metadata_json = metadata->internal_->ExportAsJson();
     request->set_post_fields(metadata_json.c_str(), metadata_json.length());
-    request->add_header("Content-Type", "application/json");
 
     RestCall(request, request->notifier(), response, handle.get(), nullptr,
              nullptr);
diff --git a/storage/src/desktop/storage_reference_desktop.h b/storage/src/desktop/storage_reference_desktop.h
@@ -176,17 +176,19 @@ class StorageReferenceInternal {
   // Upload data without metadata.
   Future<Metadata> PutBytesInternal(const void* buffer, size_t buffer_size,
                                     Listener* listener,
-                                    Controller* controller_out);
+                                    Controller* controller_out,
+                                    const char* content_type = nullptr);
   // Upload file without metadata.
   Future<Metadata> PutFileInternal(const char* path, Listener* listener,
-                                   Controller* controller_out);
+                                   Controller* controller_out,
+                                   const char* content_type = nullptr);
 
   void RestCall(rest::Request* request, internal::Notifier* request_notifier,
                 BlockingResponse* response, FutureHandle handle,
                 Listener* listener, Controller* controller_out);
 
   void PrepareRequest(rest::Request* request, const char* url,
-                      const char* method);
+                      const char* method, const char* content_type = nullptr);
 
   void SetupMetadataChain(Future<Metadata> starting_future,
                           MetadataChainData* data);
