From b3c69e85d78fb221f271d0eb608c7bb274ecbd6d Mon Sep 17 00:00:00 2001
From: Pieter De Baets <pieterdb@fb.com>
Date: Fri, 28 Jan 2022 02:54:09 -0800
Subject: [PATCH] Use AssetManager buffer instead of copying bundle

Summary:
The current JSLoader implementation (on Android) will copy the buffer into a JSBigBufferString during startup. This duplicates work Android is already doing (it allocates creates a copy of the bundle in memory while decompressing the asset from the APK) or worse, if the APK is mmap'ed this will unnecessarily create dirty memory.

The risk with this approach is that the bundle loaded from disk is no longer \0 terminated, but that's also the case for the bundles we load with `JSBigFileString` and is not an issue as  far as I can tell.

Changelog: [Internal]

Reviewed By: mhorowitz

Differential Revision: D33792735

fbshipit-source-id: 61fc089a223f3602d3575340d79a8de2ec92d8a0
---
 .../src/main/jni/react/jni/JSLoader.cpp       | 56 +++++++++++++------
 1 file changed, 40 insertions(+), 16 deletions(-)

diff --git a/ReactAndroid/src/main/jni/react/jni/JSLoader.cpp b/ReactAndroid/src/main/jni/react/jni/JSLoader.cpp
index 1c669181a51..16fcc43d4fa 100644
--- a/ReactAndroid/src/main/jni/react/jni/JSLoader.cpp
+++ b/ReactAndroid/src/main/jni/react/jni/JSLoader.cpp
@@ -9,14 +9,9 @@
 
 #include <android/asset_manager_jni.h>
 #include <cxxreact/JSBigString.h>
-#include <fb/log.h>
+#include <cxxreact/JSBundleType.h>
 #include <fbjni/fbjni.h>
 #include <folly/Conv.h>
-#include <fstream>
-#include <memory>
-#include <sstream>
-#include <streambuf>
-#include <string>
 
 #ifdef WITH_FBSYSTRACE
 #include <fbsystrace.h>
@@ -28,6 +23,31 @@ using namespace facebook::jni;
 namespace facebook {
 namespace react {
 
+class AssetManagerString : public JSBigString {
+ public:
+  AssetManagerString(AAsset *asset) : asset_(asset){};
+
+  virtual ~AssetManagerString() {
+    AAsset_close(asset_);
+  }
+
+  bool isAscii() const override {
+    return false;
+  }
+
+  const char *c_str() const override {
+    return (const char *)AAsset_getBuffer(asset_);
+  }
+
+  // Length of the c_str without the NULL byte.
+  size_t size() const override {
+    return AAsset_getLength(asset_);
+  }
+
+ private:
+  AAsset *asset_;
+};
+
 __attribute__((visibility("default"))) AAssetManager *extractAssetManager(
     alias_ref<JAssetManager::javaobject> assetManager) {
   auto env = Environment::current();
@@ -50,17 +70,21 @@ loadScriptFromAssets(AAssetManager *manager, const std::string &assetName) {
         AASSET_MODE_STREAMING); // Optimized for sequential read: see
                                 // AssetManager.java for docs
     if (asset) {
-      auto buf = std::make_unique<JSBigBufferString>(AAsset_getLength(asset));
-      size_t offset = 0;
-      int readbytes;
-      while ((readbytes = AAsset_read(
-                  asset, buf->data() + offset, buf->size() - offset)) > 0) {
-        offset += readbytes;
-      }
-      AAsset_close(asset);
-      if (offset == buf->size()) {
-        return std::move(buf);
+      auto script = std::make_unique<AssetManagerString>(asset);
+      if (script->size() >= sizeof(BundleHeader)) {
+        // When using bytecode, it's safe for the underlying buffer to not be \0
+        // terminated. In all other scenarios, we will force a copy of the
+        // script to ensure we have a terminator.
+        const BundleHeader *header =
+            reinterpret_cast<const BundleHeader *>(script->c_str());
+        if (parseTypeFromHeader(*header) == ScriptTag::HBCBundle) {
+          return script;
+        }
       }
+
+      auto buf = std::make_unique<JSBigBufferString>(script->size());
+      memcpy(buf->data(), script->c_str(), script->size());
+      return buf;
     }
   }