Fix memory monitor safety improvements (#725)

* Fix memory safety issues in memory monitor

* Add null safety checks and file protection to memory monitor

* Add tests for mmap/munmap and memory monitor enable/disable cycles

* Update KSCrashMonitor_Memory.m

Author: naftaly

Sources/KSCrashRecording/Monitors/KSCrashMonitor_Memory.m

@@ -1,5 +1,5 @@
 //
-//  KSCrashMonitor_Memory.h
+//  KSCrashMonitor_Memory.m
 //
 //  Created by Alexander Cohen on 2024-05-20.
 //
@@ -54,8 +54,8 @@
 
 static const int32_t KSCrash_Memory_Magic = 'kscm';
 
-static const uint8_t KSCrash_Memory_Version_1 = 1;
-const uint8_t KSCrash_Memory_CurrentVersion = KSCrash_Memory_Version_1;
+const uint8_t KSCrash_Memory_Version_1_0 = 1;
+const uint8_t KSCrash_Memory_CurrentVersion = KSCrash_Memory_Version_1_0;
 
 const uint8_t KSCrash_Memory_NonFatalReportLevelNone = KSCrashAppMemoryStateTerminal + 1;
 
@@ -77,6 +77,8 @@
 static void notifyPostSystemEnable(void);
 static void ksmemory_read(const char *path);
 static void ksmemory_map(const char *path);
+static void ksmemory_unmap(void);
+static void ksmemory_applyNoFileProtection(NSString *path);
 
 // ============================================================================
 #pragma mark - Globals -
@@ -140,12 +142,21 @@ static void _ks_memory_update(void (^block)(KSCrash_Memory *mem))
 static void _ks_memory_set(KSCrash_Memory *mem)
 {
     os_unfair_lock_lock(&g_memoryLock);
+    void *old = g_memory;
     g_memory = mem;
     os_unfair_lock_unlock(&g_memoryLock);
+
+    if (old) {
+        ksfu_munmap(old, sizeof(KSCrash_Memory));
+    }
 }
 
 static void _ks_memory_update_from_app_memory(KSCrashAppMemory *const memory)
 {
+    if (!memory) {
+        return;
+    }
+
     _ks_memory_update(^(KSCrash_Memory *mem) {
         *mem = (KSCrash_Memory) {
             .magic = KSCrash_Memory_Magic,
@@ -268,9 +279,10 @@ static void setEnabled(bool isEnabled)
             }];
 
     } else {
-        g_memoryTracker = nil;
         [KSCrashAppStateTracker.sharedInstance removeObserver:g_appStateObserver];
         g_appStateObserver = nil;
+        g_memoryTracker = nil;
+        ksmemory_unmap();
     }
 }
 
@@ -291,14 +303,16 @@ static void addContextualInfoToEvent(KSCrash_MonitorContext *eventContext)
         // since we're in a signal or something that can only
         // use async safe functions, we can't lock.
         // It's "ok" though, since no other threads should be running.
-        g_memory->fatal = eventContext->requirements.isFatal;
+        if (g_memory) {
+            g_memory->fatal = eventContext->requirements.isFatal;
+        }
     } else {
         _ks_memory_update(^(KSCrash_Memory *mem) {
             mem->fatal = eventContext->requirements.isFatal;
         });
     }
 
-    if (g_isEnabled) {
+    if (g_isEnabled && g_memory) {
         // same as above re: not locking when _asyncSafeOnly_ is set.
         KSCrash_Memory memCopy = asyncSafeOnly ? *g_memory : _ks_memory_copy();
         eventContext->AppMemory.footprint = memCopy.footprint;
@@ -337,7 +351,7 @@ static void kscm_memory_check_for_oom_in_previous_session(void)
     // a programming error and receiving a Mach event or signal that
     // indicates a crash, we should process that on startup and ignore
     // and indication of an OOM.
-    bool userPerceivedOOM = NO;
+    bool userPerceivedOOM = false;
     if (g_FatalReportsEnabled && ksmemory_previous_session_was_terminated_due_to_memory(&userPerceivedOOM)) {
         // We only report an OOM that the user might have seen.
         // Ignore this check if we want to report all OOM, foreground and background.
@@ -459,10 +473,10 @@ static void ksmemory_read(const char *path)
 
     // check the timestamp, let's say it's valid for the last week
     // do we really want crash reports older than a week anyway??
-    const uint64_t kUS_in_day = 8.64e+10;
+    const uint64_t kUS_in_day = 86400000000ULL;  // 24 * 60 * 60 * 1000000
     const uint64_t kUS_in_week = kUS_in_day * 7;
     uint64_t now = ksdate_microseconds();
-    if (memory.timestamp <= 0 || memory.timestamp == INT64_MAX || memory.timestamp < now - kUS_in_week) {
+    if (memory.timestamp == 0 || memory.timestamp > now || memory.timestamp < now - kUS_in_week) {
         return;
     }
 
@@ -516,9 +530,18 @@ static void ksmemory_map(const char *path)
     }
 
     _ks_memory_set(ptr);
-    _ks_memory_update_from_app_memory(g_memoryTracker.memory);
+
+    KSCrashAppMemory *currentMemory = g_memoryTracker.memory;
+    if (currentMemory) {
+        _ks_memory_update_from_app_memory(currentMemory);
+    }
 }
 
+/**
+ Unmaps the memory-mapped file and clears the global pointer.
+ */
+static void ksmemory_unmap(void) { _ks_memory_set(NULL); }
+
 /**
  What we're doing here is writing a file out that can be reused
  on restart if the data shows us there was a memory issue.
@@ -562,12 +585,27 @@ static void ksmemory_write_possible_oom(void)
     g_callbacks.handle(ctx);
 }
 
+static void ksmemory_applyNoFileProtection(NSString *path)
+{
+    if (!path) {
+        return;
+    }
+
+    NSDictionary *attrs = @ { NSFileProtectionKey : NSFileProtectionNone };
+    [[NSFileManager defaultManager] setAttributes:attrs ofItemAtPath:path error:nil];
+}
+
 void ksmemory_initialize(const char *dataPath)
 {
-    g_hasPostEnable = 0;
+    g_hasPostEnable = false;
     g_dataURL = [NSURL fileURLWithPath:@(dataPath)];
     g_memoryURL = [g_dataURL URLByAppendingPathComponent:@"memory.bin"];
 
+    // Ensure files we touch stay readable while the app is locked.
+    ksmemory_applyNoFileProtection(g_dataURL.path);
+    ksmemory_applyNoFileProtection(g_memoryURL.path);
+    ksmemory_applyNoFileProtection(kscm_memory_oom_breadcrumb_URL().path);
+
     // load up the old memory data
     ksmemory_read(g_memoryURL.path.UTF8String);
 }
@@ -579,7 +617,7 @@ bool ksmemory_previous_session_was_terminated_due_to_memory(bool *userPerceptibl
     // exception/event occured that terminated/crashed the app. We don't want to report
     // that as an OOM.
     if (g_previousSessionMemory.fatal) {
-        return NO;
+        return false;
     }
 
     // We might care if the user might have seen the OOM

Sources/KSCrashRecordingCore/KSFileUtils.c

@@ -583,8 +583,22 @@ void *ksfu_mmap(const char *path, size_t size)
         // This comes before close which is ok since it'll happen
         // when all fd's are closed.
         unlink(path);
+        close(fd);
+        return NULL;
     }
 
     close(fd);
     return ptr;
 }
+
+bool ksfu_munmap(void *ptr, size_t size)
+{
+    if (ptr == NULL) {
+        return true;
+    }
+    if (munmap(ptr, size) == -1) {
+        KSLOG_ERROR("Could not munmap: %s", strerror(errno));
+        return false;
+    }
+    return true;
+}

Sources/KSCrashRecordingCore/include/KSFileUtils.h

@@ -274,12 +274,22 @@ bool ksfu_readBufferedReaderUntilChar(KSBufferedReader *reader, int ch, char *ds
  * @param size The size of the map.
  *
  * @return the mapped pointer if successful, NULL otherwise.
- * The return value must be unmapped using `munmap` when done
+ * The return value must be unmapped using `ksfu_munmap` when done
  * with the returned pointer. It is ok to let the pointer live up to termination,
  * the system will unmap on termination if required.
  */
 void *ksfu_mmap(const char *path, size_t size);
 
+/** Unmaps a memory-mapped region previously created by ksfu_mmap.
+ *
+ * @param ptr The pointer returned by ksfu_mmap.
+ *
+ * @param size The size that was passed to ksfu_mmap.
+ *
+ * @return true if successful, false otherwise.
+ */
+bool ksfu_munmap(void *ptr, size_t size);
+
 #ifdef __cplusplus
 }
 #endif

Tests/KSCrashRecordingCoreTests/KSFileUtils_Tests.m

@@ -648,4 +648,58 @@ - (void)testWriteBuffered_FlushAndLargeWriteOrder
     XCTAssertEqualObjects(actualFileContents, expectedContents);
 }
 
+#pragma mark - Memory Mapping Tests
+
+- (void)testMunmap_NullPointer
+{
+    // munmap with NULL should return true (no-op)
+    XCTAssertTrue(ksfu_munmap(NULL, 100));
+}
+
+- (void)testMmap_InvalidPath
+{
+    // mmap with non-existent path should return NULL
+    NSString *invalidPath = @"/nonexistent/path/to/file.bin";
+    void *ptr = ksfu_mmap(invalidPath.UTF8String, 100);
+    XCTAssertTrue(ptr == NULL);
+}
+
+- (void)testMmap_Munmap_RoundTrip
+{
+    // Test mapping a file, writing to it, and unmapping
+    NSString *path = [self generateTempFilePath];
+    size_t size = 1024;
+
+    void *ptr = ksfu_mmap(path.UTF8String, size);
+    XCTAssertTrue(ptr != NULL);
+
+    // Write some data to the mapped memory
+    memset(ptr, 0xAB, size);
+
+    // Verify the data was written
+    uint8_t *bytes = (uint8_t *)ptr;
+    XCTAssertEqual(bytes[0], 0xAB);
+    XCTAssertEqual(bytes[size - 1], 0xAB);
+
+    // Unmap should succeed
+    XCTAssertTrue(ksfu_munmap(ptr, size));
+}
+
+- (void)testMmap_Munmap_MultipleRoundTrips
+{
+    // Test multiple map/unmap cycles to ensure no leaks or issues
+    NSString *path = [self generateTempFilePath];
+    size_t size = 256;
+
+    for (int i = 0; i < 5; i++) {
+        void *ptr = ksfu_mmap(path.UTF8String, size);
+        XCTAssertTrue(ptr != NULL);
+
+        // Write iteration number
+        memset(ptr, (uint8_t)i, size);
+
+        XCTAssertTrue(ksfu_munmap(ptr, size));
+    }
+}
+
 @end

Tests/KSCrashRecordingTests/KSCrashMonitor_Memory_Tests.m

@@ -319,4 +319,35 @@ - (void)testNonFatalReportLevel
     XCTAssertEqual(ksmemory_get_nonfatal_report_level(), KSCrashAppMemoryStateUrgent);
 }
 
+- (void)testRapidEnableDisableCycles
+{
+    // Test rapid enable/disable cycles to ensure proper cleanup
+    // This helps catch issues with memory mapping/unmapping
+    KSCrashMonitorAPI *api = kscm_memory_getAPI();
+
+    for (int i = 0; i < 10; i++) {
+        api->setEnabled(true);
+        XCTAssertTrue(api->isEnabled());
+        api->setEnabled(false);
+        XCTAssertFalse(api->isEnabled());
+    }
+}
+
+- (void)testEnableDisableWithoutInitialize
+{
+    // Test that enable/disable handles edge cases gracefully
+    KSCrashMonitorAPI *api = kscm_memory_getAPI();
+
+    // Should handle being disabled when already disabled
+    api->setEnabled(false);
+    XCTAssertFalse(api->isEnabled());
+
+    // Enable then disable
+    api->setEnabled(true);
+    XCTAssertTrue(api->isEnabled());
+    [NSThread sleepForTimeInterval:0.05];
+    api->setEnabled(false);
+    XCTAssertFalse(api->isEnabled());
+}
+
 @end