protect from background writing in iOS

Core/Core.xcodeproj/project.pbxproj

@@ -66,6 +66,7 @@
 
 /* Begin PBXFileReference section */
 		CB0DCC70242B57E5009AFE59 /* libc++.tbd */ = {isa = PBXFileReference; lastKnownFileType = "sourcecode.text-based-dylib-definition"; name = "libc++.tbd"; path = "Platforms/MacOSX.platform/Developer/SDKs/MacOSX10.15.sdk/usr/lib/libc++.tbd"; sourceTree = DEVELOPER_DIR; };
+		CB467F862431D3ED00FD7421 /* MMKV_OSX.h */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.c.h; path = MMKV_OSX.h; sourceTree = "<group>"; };
 		CB58B3FF23AB3035002457F1 /* Foundation.framework */ = {isa = PBXFileReference; lastKnownFileType = wrapper.framework; name = Foundation.framework; path = Platforms/MacOSX.platform/Developer/SDKs/MacOSX10.15.sdk/System/Library/Frameworks/Foundation.framework; sourceTree = DEVELOPER_DIR; };
 		CB9563D823AB2D9500ACCD39 /* libMMKVCore.a */ = {isa = PBXFileReference; explicitFileType = archive.ar; includeInIndex = 0; path = libMMKVCore.a; sourceTree = BUILT_PRODUCTS_DIR; };
 		CB9563E423AB2E9100ACCD39 /* CodedInputData.h */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.c.h; path = CodedInputData.h; sourceTree = "<group>"; };
@@ -162,6 +163,7 @@
 				CB9563F823AB2E9100ACCD39 /* MMBuffer.cpp */,
 				CB9563FA23AB2E9100ACCD39 /* MMBuffer.h */,
 				CBD723BE23B5C22800D3CDAF /* MMKV_OSX.cpp */,
+				CB467F862431D3ED00FD7421 /* MMKV_OSX.h */,
 				CB9563F223AB2E9100ACCD39 /* MMKV.cpp */,
 				CB9563ED23AB2E9100ACCD39 /* MMKV.h */,
 				CB95641123AB2E9100ACCD39 /* MMKVLog.cpp */,

Core/MMKV.cpp

@@ -654,8 +654,13 @@ void MMKV::oldStyleWriteActualSize(size_t actualSize) {
     MMKV_ASSERT(m_file->getMemory());
 
     m_actualSize = actualSize;
-    // TODO: background protection
+#ifdef MMKV_IOS
+    protectFromBackgroundWriting(m_file->getMemory(), Fixed32Size, ^{
+      memcpy(m_file->getMemory(), &actualSize, Fixed32Size);
+    });
+#else
     memcpy(m_file->getMemory(), &actualSize, Fixed32Size);
+#endif
 }
 
 bool MMKV::writeActualSize(size_t size, uint32_t crcDigest, const void *iv, bool increaseSequence) {
@@ -666,7 +671,6 @@ bool MMKV::writeActualSize(size_t size, uint32_t crcDigest, const void *iv, bool
         return false;
     }
 
-    // TODO: background protection
     bool needsFullWrite = false;
     m_actualSize = size;
     m_metaInfo->m_actualSize = static_cast<uint32_t>(size);
@@ -692,13 +696,22 @@ bool MMKV::writeActualSize(size_t size, uint32_t crcDigest, const void *iv, bool
         }
         needsFullWrite = true;
     }
+#ifdef MMKV_IOS
+    return protectFromBackgroundWriting(m_metaFile->getMemory(), sizeof(MMKVMetaInfo), ^{
+      if (unlikely(needsFullWrite)) {
+          m_metaInfo->write(m_metaFile->getMemory());
+      } else {
+          m_metaInfo->writeCRCAndActualSizeOnly(m_metaFile->getMemory());
+      }
+    });
+#else
     if (unlikely(needsFullWrite)) {
         m_metaInfo->write(m_metaFile->getMemory());
     } else {
         m_metaInfo->writeCRCAndActualSizeOnly(m_metaFile->getMemory());
     }
-
     return true;
+#endif
 }
 
 const MMBuffer &MMKV::getDataForKey(MMKVKey_t key) {
@@ -769,9 +782,9 @@ bool MMKV::appendDataWithKey(const MMBuffer &data, MMKVKey_t key) {
     }
 
 #ifdef MMKV_IOS
-    auto ret = protectFromBackgroundWriting(size, ^(CodedOutputData *output) {
-      output->writeString(key);
-      output->writeData(data); // note: write size of data
+    auto ret = protectFromBackgroundWriting(m_output->curWritePointer(), size, ^{
+      m_output->writeString(key);
+      m_output->writeData(data); // note: write size of data
     });
     if (!ret) {
         return false;
@@ -842,8 +855,8 @@ bool MMKV::doFullWriteBack(MMBuffer &&allData) {
     delete m_output;
     m_output = new CodedOutputData(ptr + Fixed32Size, m_file->getFileSize() - Fixed32Size);
 #ifdef MMKV_IOS
-    auto ret = protectFromBackgroundWriting(allData.length(), ^(CodedOutputData *output) {
-      output->writeRawData(allData); // note: don't write size of data
+    auto ret = protectFromBackgroundWriting(m_output->curWritePointer(), allData.length(), ^{
+      m_output->writeRawData(allData); // note: don't write size of data
     });
     if (!ret) {
         // revert everything

Core/MMKV.h

@@ -133,9 +133,9 @@ class MMKV {
     void checkReSetCryptKey(int fd, int metaFD, std::string *cryptKey);
 #endif
 
-#ifdef MMKV_IOS
-    typedef void (^WriteBlock)(mmkv::CodedOutputData *output);
-    bool protectFromBackgroundWriting(size_t size, WriteBlock block);
+#if defined(MMKV_IOS)
+    typedef void (^WriteBlock)(void);
+    bool protectFromBackgroundWriting(void *ptr, size_t size, WriteBlock block);
 #endif
 
 public:

Core/MMKVPredef.h

@@ -35,9 +35,6 @@
 #        define MMKV_IOS_OR_MAC
 #        ifdef __ENVIRONMENT_IPHONE_OS_VERSION_MIN_REQUIRED__
 #            define MMKV_IOS
-#           if __has_feature(attribute_availability_app_extension)
-#               define MMKV_IOS_EXTENSION
-#           endif
 #        else
 #            define MMKV_MAC
 #        endif

Core/MMKV_OSX.cpp

@@ -34,6 +34,7 @@
 
 #    ifdef MMKV_IOS
 #        include "Checksum.h"
+#        include "MMKV_OSX.h"
 #        include <sys/mman.h>
 #    endif
 
@@ -50,6 +51,29 @@ extern MMKVPath_t g_rootDir;
 enum { UnKnown = 0, PowerMac = 1, Mac, iPhone, iPod, iPad, AppleTV, AppleWatch };
 static void GetAppleMachineInfo(int &device, int &version);
 
+#    ifdef MMKV_IOS
+MLockPtr::MLockPtr(void *ptr, size_t size) : m_lockDownSize(0), m_lockedPtr(nullptr) {
+    // calc ptr to be mlock()
+    auto writePtr = (size_t) ptr;
+    auto lockPtr = (writePtr / DEFAULT_MMAP_SIZE) * DEFAULT_MMAP_SIZE;
+    auto lockDownSize = writePtr - lockPtr + size;
+    if (mlock((void *) lockPtr, lockDownSize) == 0) {
+        m_lockedPtr = (uint8_t *) lockPtr;
+        m_lockDownSize = lockDownSize;
+    } else {
+        MMKVError("fail to mlock [%p], %s", m_lockedPtr, strerror(errno));
+        // just fail on this condition, otherwise app will crash anyway
+    }
+}
+
+MLockPtr::~MLockPtr() {
+    if (m_lockedPtr) {
+        munlock(m_lockedPtr, m_lockDownSize);
+    }
+}
+
+#    endif
+
 MMKV_NAMESPACE_BEGIN
 
 extern ThreadOnceToken_t once_control;
@@ -85,38 +109,28 @@ void MMKV::setIsInBackground(bool isInBackground) {
     MMKVInfo("g_isInBackground:%d", g_isInBackground);
 }
 
-// @finally in C++ stype
-template <typename F>
-struct AtExit {
-    AtExit(F f) : m_func{f} {}
-    ~AtExit() { m_func(); }
-
-private:
-    F m_func;
-};
-
-bool MMKV::protectFromBackgroundWriting(size_t size, WriteBlock block) {
+bool MMKV::protectFromBackgroundWriting(void *ptr, size_t size, WriteBlock block) {
     if (g_isInBackground) {
-        // calc ptr to be mlock()
-        auto writePtr = (size_t) m_output->curWritePointer();
-        auto ptr = (writePtr / DEFAULT_MMAP_SIZE) * DEFAULT_MMAP_SIZE;
-        size_t lockDownSize = writePtr - ptr + size;
-        if (mlock((void *) ptr, lockDownSize) != 0) {
-            MMKVError("fail to mlock [%s], %s", m_mmapID.c_str(), strerror(errno));
-            // just fail on this condition, otherwise app will crash anyway
-            //block(m_output);
-            return false;
-        } else {
-            AtExit cleanup([=] { munlock((void *) ptr, lockDownSize); });
+        MLockPtr mlockPtr(ptr, size);
+        if (mlockPtr.isLocked()) {
             try {
-                block(m_output);
+                block();
             } catch (std::exception &exception) {
                 MMKVError("%s", exception.what());
                 return false;
             }
+        } else {
+            // just fail on this condition, otherwise app will crash anyway
+            //block(m_output);
+            return false;
         }
     } else {
-        block(m_output);
+        try {
+            block();
+        } catch (std::exception &exception) {
+            MMKVError("%s", exception.what());
+            return false;
+        }
     }
 
     return true;

Core/MMKV_OSX.h

@@ -0,0 +1,44 @@
+/*
+ * Tencent is pleased to support the open source community by making
+ * MMKV available.
+ *
+ * Copyright (C) 2019 THL A29 Limited, a Tencent company.
+ * All rights reserved.
+ *
+ * Licensed under the BSD 3-Clause License (the "License"); you may not use
+ * this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ *       https://opensource.org/licenses/BSD-3-Clause
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+#include "MMKVPredef.h"
+
+#if defined(MMKV_IOS) && defined(__cplusplus)
+
+class MLockPtr {
+    size_t m_lockDownSize;
+    uint8_t *m_lockedPtr;
+
+public:
+    MLockPtr(void *ptr, size_t size);
+
+    ~MLockPtr();
+
+    bool isLocked() const {
+        return (m_lockedPtr != nullptr);
+    }
+
+    // just forbid it for possibly misuse
+    explicit MLockPtr(const MLockPtr &other) = delete;
+    MLockPtr &operator=(const MLockPtr &other) = delete;
+};
+
+#endif

iOS/MMKV/MMKV.xcodeproj/project.pbxproj

@@ -393,6 +393,11 @@
 				DYLIB_COMPATIBILITY_VERSION = 1;
 				DYLIB_CURRENT_VERSION = 1;
 				DYLIB_INSTALL_NAME_BASE = "@rpath";
+				GCC_PREPROCESSOR_DEFINITIONS = (
+					"DEBUG=1",
+					"$(inherited)",
+					MMKV_IOS_EXTENSION,
+				);
 				GENERATE_MASTER_OBJECT_FILE = NO;
 				INFOPLIST_FILE = "MMKV copy-Info.plist";
 				INSTALL_PATH = "$(LOCAL_LIBRARY_DIR)/Frameworks";
@@ -435,6 +440,10 @@
 				DYLIB_COMPATIBILITY_VERSION = 1;
 				DYLIB_CURRENT_VERSION = 1;
 				DYLIB_INSTALL_NAME_BASE = "@rpath";
+				GCC_PREPROCESSOR_DEFINITIONS = (
+					"NDEBUG=1",
+					MMKV_IOS_EXTENSION,
+				);
 				GENERATE_MASTER_OBJECT_FILE = NO;
 				INFOPLIST_FILE = "MMKV copy-Info.plist";
 				INSTALL_PATH = "$(LOCAL_LIBRARY_DIR)/Frameworks";

iOS/MMKV/MMKV/libMMKV.mm

@@ -38,6 +38,10 @@
 static NSString *g_basePath = nil;
 static NSString *g_groupPath = nil;
 
+#if defined(MMKV_IOS) && !defined(MMKV_IOS_EXTENSION)
+static BOOL g_isRunningInAppExtension = NO;
+#endif
+
 static void LogHandler(mmkv::MMKVLogLevel level, const char *file, int line, const char *function, NSString *message);
 static mmkv::MMKVRecoverStrategic ErrorHandler(const string &mmapID, mmkv::MMKVErrorType errorType);
 static void ContentChangeHandler(const string &mmapID);
@@ -60,13 +64,18 @@ + (void)initialize {
         mmkv::MMKV::minimalInit([self mmkvBasePath].UTF8String);
 
 #if defined(MMKV_IOS) && !defined(MMKV_IOS_EXTENSION)
-        auto appState = [UIApplication sharedApplication].applicationState;
-        auto isInBackground = (appState == UIApplicationStateBackground);
-        mmkv::MMKV::setIsInBackground(isInBackground);
-        MMKVInfo("appState:%ld", (long) appState);
-
-        [[NSNotificationCenter defaultCenter] addObserver:self selector:@selector(didEnterBackground) name:UIApplicationDidEnterBackgroundNotification object:nil];
-        [[NSNotificationCenter defaultCenter] addObserver:self selector:@selector(didBecomeActive) name:UIApplicationDidBecomeActiveNotification object:nil];
+        if ([[[NSBundle mainBundle] bundlePath] hasSuffix:@".appex"]) {
+            g_isRunningInAppExtension = YES;
+        }
+        if (g_isRunningInAppExtension) {
+            auto appState = [UIApplication sharedApplication].applicationState;
+            auto isInBackground = (appState == UIApplicationStateBackground);
+            mmkv::MMKV::setIsInBackground(isInBackground);
+            MMKVInfo("appState:%ld", (long) appState);
+
+            [[NSNotificationCenter defaultCenter] addObserver:self selector:@selector(didEnterBackground) name:UIApplicationDidEnterBackgroundNotification object:nil];
+            [[NSNotificationCenter defaultCenter] addObserver:self selector:@selector(didBecomeActive) name:UIApplicationDidBecomeActiveNotification object:nil];
+        }
 #endif
     }
 }
@@ -183,10 +192,12 @@ - (instancetype)initWithMMapID:(NSString *)mmapID cryptKey:(NSData *)cryptKey re
         m_mmapID = [NSString stringWithUTF8String:m_mmkv->mmapID().c_str()];
 
 #if defined(MMKV_IOS) && !defined(MMKV_IOS_EXTENSION)
-        [[NSNotificationCenter defaultCenter] addObserver:self
-                                                 selector:@selector(onMemoryWarning)
-                                                     name:UIApplicationDidReceiveMemoryWarningNotification
-                                                   object:nil];
+        if (g_isRunningInAppExtension) {
+            [[NSNotificationCenter defaultCenter] addObserver:self
+                                                     selector:@selector(onMemoryWarning)
+                                                         name:UIApplicationDidReceiveMemoryWarningNotification
+                                                       object:nil];
+        }
 #endif
     }
     return self;