Add DiskArbitrationTestUtil to shim out DiskArbitration for unit testing (#720)

Source/santad/BUILD

@@ -107,6 +107,23 @@ objc_library(
     ],
 )
 
+objc_library(
+    name = "DiskArbitrationTestLib",
+    testonly = 1,
+    srcs = [
+        "EventProviders/DiskArbitrationTestUtil.h",
+        "EventProviders/DiskArbitrationTestUtil.mm",
+    ],
+    sdk_dylibs = [
+        "EndpointSecurity",
+        "bsm",
+    ],
+    sdk_frameworks = [
+        "DiskArbitration",
+        "IOKit",
+    ],
+)
+
 macos_bundle(
     name = "com.google.santa.daemon",
     bundle_extension = "systemextension",
@@ -230,6 +247,7 @@ santa_unit_test(
         "bsm",
     ],
     deps = [
+        ":DiskArbitrationTestLib",
         ":EndpointSecurityTestLib",
         ":santad_lib",
         "//Source/common:SNTKernelCommon",

Source/santad/EventProviders/DiskArbitrationTestUtil.h

@@ -0,0 +1,84 @@
+/// Copyright 2021 Google Inc. All rights reserved.
+///
+/// Licensed under the Apache License, Version 2.0 (the "License");
+/// you may not use this file except in compliance with the License.
+/// You may obtain a copy of the License at
+///
+///    http://www.apache.org/licenses/LICENSE-2.0
+///
+///    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.
+
+#include <CoreFoundation/CFDictionary.h>
+#include <CoreFoundation/CoreFoundation.h>
+#include <DiskArbitration/DiskArbitration.h>
+#include <Foundation/Foundation.h>
+
+NS_ASSUME_NONNULL_BEGIN
+
+// Mock object to point the opaque DADiskRefs to instead.
+// Note that this will have undefined behavior for DA functions that aren't
+// shimmed out by this utility, as the original DADiskRef refers to a completely
+// different struct managed by the CFRuntime.
+// https://opensource.apple.com/source/DiskArbitration/DiskArbitration-297.70.1/DiskArbitration/DADisk.c.auto.html
+@interface MockDADisk : NSObject
+@property(nonatomic) NSDictionary *diskDescription;
+@property(nonatomic, readwrite) NSString *name;
+@end
+
+typedef void (^MockDADiskAppearedCallback)(DADiskRef ref);
+// Singleton mock fixture around all of the DiskArbitration framework functions
+@interface MockDiskArbitration : NSObject
+@property(nonatomic, readwrite, nonnull)
+  NSMutableDictionary<NSString *, MockDADisk *> *insertedDevices;
+@property(nonatomic, readwrite, nonnull)
+  NSMutableArray<MockDADiskAppearedCallback> *diskAppearedCallbacks;
+@property(nonatomic) BOOL wasRemounted;
+@property(nonatomic, nullable) dispatch_queue_t sessionQueue;
+
+- (instancetype _Nonnull)init;
+- (void)reset;
+
+// Also triggers DADiskRegisterDiskAppearedCallback
+- (void)insert:(MockDADisk *)ref bsdName:(NSString *)bsdName;
+
+// Retrieve an initialized singleton MockDiskArbitration object
++ (instancetype _Nonnull)mockDiskArbitration;
+@end
+
+//
+// All DiskArbitration functions used in SNTDeviceManager and shimmed out accordingly.
+//
+CF_EXTERN_C_BEGIN
+
+void DADiskMountWithArguments(DADiskRef _Nonnull disk, CFURLRef __nullable path,
+                              DADiskMountOptions options, DADiskMountCallback __nullable callback,
+                              void *__nullable context,
+                              CFStringRef __nullable arguments[_Nullable]);
+
+DADiskRef __nullable DADiskCreateFromBSDName(CFAllocatorRef __nullable allocator,
+                                             DASessionRef session, const char *name);
+
+CFDictionaryRef __nullable DADiskCopyDescription(DADiskRef disk);
+
+void DARegisterDiskAppearedCallback(DASessionRef session, CFDictionaryRef __nullable match,
+                                    DADiskAppearedCallback callback, void *__nullable context);
+
+void DARegisterDiskDisappearedCallback(DASessionRef session, CFDictionaryRef __nullable match,
+                                       DADiskDisappearedCallback callback,
+                                       void *__nullable context);
+
+void DARegisterDiskDescriptionChangedCallback(DASessionRef session,
+                                              CFDictionaryRef __nullable match,
+                                              CFArrayRef __nullable watch,
+                                              DADiskDescriptionChangedCallback callback,
+                                              void *__nullable context);
+
+void DASessionSetDispatchQueue(DASessionRef session, dispatch_queue_t __nullable queue);
+DASessionRef __nullable DASessionCreate(CFAllocatorRef __nullable allocator);
+
+CF_EXTERN_C_END
+NS_ASSUME_NONNULL_END

Source/santad/EventProviders/DiskArbitrationTestUtil.mm

@@ -0,0 +1,120 @@
+/// Copyright 2021 Google Inc. All rights reserved.
+///
+/// Licensed under the Apache License, Version 2.0 (the "License");
+/// you may not use this file except in compliance with the License.
+/// You may obtain a copy of the License at
+///
+///    http://www.apache.org/licenses/LICENSE-2.0
+///
+///    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.
+#import <Foundation/Foundation.h>
+
+#include <stdlib.h>
+
+#import "Source/santad/EventProviders/DiskArbitrationTestUtil.h"
+
+NS_ASSUME_NONNULL_BEGIN
+
+@implementation MockDADisk
+@end
+
+@implementation MockDiskArbitration
+
+- (instancetype _Nonnull)init {
+  self = [super init];
+  if (self) {
+    _insertedDevices = [NSMutableDictionary dictionary];
+    _diskAppearedCallbacks = [NSMutableArray array];
+  }
+  return self;
+}
+
+- (void)reset {
+  [self.insertedDevices removeAllObjects];
+  [self.diskAppearedCallbacks removeAllObjects];
+  self.sessionQueue = dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0);
+  self.wasRemounted = NO;
+}
+
+- (void)insert:(MockDADisk *)ref bsdName:(NSString *)bsdName {
+  self.insertedDevices[bsdName] = ref;
+
+  for (MockDADiskAppearedCallback callback in self.diskAppearedCallbacks) {
+    dispatch_sync(self.sessionQueue, ^{
+      callback((__bridge DADiskRef)ref);
+    });
+  }
+}
+
+// Retrieve an initialized singleton MockDiskArbitration object
++ (instancetype _Nonnull)mockDiskArbitration {
+  static MockDiskArbitration *sharedES;
+  static dispatch_once_t onceToken;
+  dispatch_once(&onceToken, ^{
+    sharedES = [[MockDiskArbitration alloc] init];
+  });
+  return sharedES;
+};
+
+@end
+
+void DADiskMountWithArguments(DADiskRef _Nonnull disk, CFURLRef __nullable path,
+                              DADiskMountOptions options, DADiskMountCallback __nullable callback,
+                              void *__nullable context,
+                              CFStringRef __nullable arguments[_Nullable]) {
+  MockDiskArbitration *mockDA = [MockDiskArbitration mockDiskArbitration];
+  mockDA.wasRemounted = YES;
+}
+
+DADiskRef __nullable DADiskCreateFromBSDName(CFAllocatorRef __nullable allocator,
+                                             DASessionRef session, const char *name) {
+  NSString *nsName = [NSString stringWithUTF8String:name];
+
+  MockDiskArbitration *mockDA = [MockDiskArbitration mockDiskArbitration];
+  MockDADisk *got = mockDA.insertedDevices[nsName];
+  DADiskRef ref = (__bridge DADiskRef)got;
+  CFRetain(ref);
+  return ref;
+}
+
+CFDictionaryRef __nullable DADiskCopyDescription(DADiskRef disk) {
+  CFDictionaryRef description = NULL;
+  if (disk) {
+    MockDADisk *mockDisk = (__bridge MockDADisk *)disk;
+    description = (__bridge_retained CFDictionaryRef)mockDisk.diskDescription;
+  }
+  return description;
+}
+
+void DARegisterDiskAppearedCallback(DASessionRef session, CFDictionaryRef __nullable match,
+                                    DADiskAppearedCallback callback, void *__nullable context) {
+  MockDiskArbitration *mockDA = [MockDiskArbitration mockDiskArbitration];
+  [mockDA.diskAppearedCallbacks addObject:^(DADiskRef ref) {
+    callback(ref, context);
+  }];
+}
+
+void DARegisterDiskDisappearedCallback(DASessionRef session, CFDictionaryRef __nullable match,
+                                       DADiskDisappearedCallback callback,
+                                       void *__nullable context){};
+
+void DARegisterDiskDescriptionChangedCallback(DASessionRef session,
+                                              CFDictionaryRef __nullable match,
+                                              CFArrayRef __nullable watch,
+                                              DADiskDescriptionChangedCallback callback,
+                                              void *__nullable context){};
+
+void DASessionSetDispatchQueue(DASessionRef session, dispatch_queue_t __nullable queue) {
+  MockDiskArbitration *mockDA = [MockDiskArbitration mockDiskArbitration];
+  mockDA.sessionQueue = queue;
+};
+
+DASessionRef __nullable DASessionCreate(CFAllocatorRef __nullable allocator) {
+  return (__bridge DASessionRef)[MockDiskArbitration mockDiskArbitration];
+};
+
+NS_ASSUME_NONNULL_END

Source/santad/EventProviders/SNTDeviceManager.h

@@ -24,9 +24,9 @@
 
 @property(nonatomic, readwrite) BOOL subscribed;
 @property(nonatomic, readwrite) BOOL blockUSBMount;
-@property(nonatomic, readwrite) NSArray<NSString *> *remountArgs;
+@property(nonatomic, readwrite, nullable) NSArray<NSString *> *remountArgs;
 
-- (instancetype)init;
+- (instancetype _Nonnull)init;
 - (void)listen;
 - (BOOL)subscribed;
 

Source/santad/EventProviders/SNTDeviceManager.mm

@@ -44,7 +44,8 @@ void diskMountedCallback(DADiskRef disk, DADissenterRef dissenter, void *context
 void diskAppearedCallback(DADiskRef disk, void *context) {
   NSDictionary *props = CFBridgingRelease(DADiskCopyDescription(disk));
   if (![props[@"DAVolumeMountable"] boolValue]) return;
-  [[SNTEventLog logger] logDiskAppeared:props];
+  SNTEventLog *logger = [SNTEventLog logger];
+  if (logger) [logger logDiskAppeared:props];
 }
 
 void diskDescriptionChangedCallback(DADiskRef disk, CFArrayRef keys, void *context) {
@@ -52,7 +53,8 @@ void diskDescriptionChangedCallback(DADiskRef disk, CFArrayRef keys, void *conte
   if (![props[@"DAVolumeMountable"] boolValue]) return;
 
   if (props[@"DAVolumePath"]) {
-    [[SNTEventLog logger] logDiskAppeared:props];
+    SNTEventLog *logger = [SNTEventLog logger];
+    if (logger) [logger logDiskAppeared:props];
   }
 }
 
@@ -60,7 +62,8 @@ void diskDisappearedCallback(DADiskRef disk, void *context) {
   NSDictionary *props = CFBridgingRelease(DADiskCopyDescription(disk));
   if (![props[@"DAVolumeMountable"] boolValue]) return;
 
-  [[SNTEventLog logger] logDiskDisappeared:props];
+  SNTEventLog *logger = [SNTEventLog logger];
+  if (logger) [logger logDiskDisappeared:props];
 }
 
 NSArray<NSString *> *maskToMountArgs(long remountOpts) {
@@ -105,25 +108,25 @@ long mountArgsToMask(NSArray<NSString *> *args) {
 @interface SNTDeviceManager ()
 
 @property DASessionRef diskArbSession;
-@property(readonly, nonatomic) es_client_t *client;
+@property(nonatomic, readonly) es_client_t *client;
 @property(nonatomic, readonly) dispatch_queue_t esAuthQueue;
+@property(nonatomic, readonly) dispatch_queue_t diskQueue;
 @end
 
 @implementation SNTDeviceManager
 
-- (instancetype)init API_AVAILABLE(macos(10.15)) {
+- (instancetype _Nonnull)init API_AVAILABLE(macos(10.15)) {
   self = [super init];
   if (self) {
     _blockUSBMount = false;
 
-    dispatch_queue_t disk_queue =
-      dispatch_queue_create("com.google.santad.disk_queue", DISPATCH_QUEUE_SERIAL);
+    _diskQueue = dispatch_queue_create("com.google.santad.disk_queue", DISPATCH_QUEUE_SERIAL);
 
     _esAuthQueue =
       dispatch_queue_create("com.google.santa.daemon.es_device_auth", DISPATCH_QUEUE_CONCURRENT);
 
     _diskArbSession = DASessionCreate(NULL);
-    DASessionSetDispatchQueue(_diskArbSession, disk_queue);
+    DASessionSetDispatchQueue(_diskArbSession, _diskQueue);
 
     if (@available(macos 10.15, *)) [self initES];
   }
@@ -205,16 +208,17 @@ long mountArgsToMask(NSArray<NSString *> *args) {
   // TODO(tnek): Log all of the other attributes available in diskInfo into a structured log format.
   NSDictionary *diskInfo = CFBridgingRelease(DADiskCopyDescription(disk));
   BOOL isRemovable = [diskInfo[(__bridge NSString *)kDADiskDescriptionMediaRemovableKey] boolValue];
-  BOOL isUSB = [diskInfo[@"DADeviceProtocol"] isEqualTo:@"USB"];
+  BOOL isUSB =
+    [diskInfo[(__bridge NSString *)kDADiskDescriptionDeviceProtocolKey] isEqualTo:@"USB"];
 
   if (!isRemovable || !isUSB) {
     es_respond_auth_result(self.client, m, ES_AUTH_RESULT_ALLOW, false);
     return;
   }
 
-  es_respond_auth_result(self.client, m, ES_AUTH_RESULT_DENY, false);
+  BOOL shouldRemount = self.remountArgs != nil && [self.remountArgs count] > 0;
 
-  if (self.remountArgs != nil && [self.remountArgs count] > 0) {
+  if (shouldRemount) {
     long remountOpts = mountArgsToMask(self.remountArgs);
     if (mountMode & remountOpts) {
       es_respond_auth_result(self.client, m, ES_AUTH_RESULT_ALLOW, false);
@@ -227,6 +231,7 @@ long mountArgsToMask(NSArray<NSString *> *args) {
          newMode);
     [self remount:disk mountMode:newMode];
   }
+  es_respond_auth_result(self.client, m, ES_AUTH_RESULT_DENY, false);
 }
 
 - (void)remount:(DADiskRef)disk mountMode:(long)remountMask {
@@ -237,6 +242,7 @@ long mountArgsToMask(NSArray<NSString *> *args) {
 
   DADiskMountWithArguments(disk, NULL, kDADiskMountOptionDefault, diskMountedCallback,
                            (__bridge void *)self, (CFStringRef *)argv);
+
   free(argv);
 }
 

Source/santad/EventProviders/SNTDeviceManagerTest.mm

@@ -19,21 +19,29 @@
 #include <sys/mount.h>
 
 #import "Source/common/SNTConfigurator.h"
-#import "Source/santad/EventProviders/EndpointSecurityTestUtil.h"
 #import "Source/santad/EventProviders/SNTDeviceManager.h"
 
+#import "Source/santad/EventProviders/DiskArbitrationTestUtil.h"
+#import "Source/santad/EventProviders/EndpointSecurityTestUtil.h"
+
 @interface SNTDeviceManagerTest : XCTestCase
+@property id mockConfigurator;
 @end
 
 @implementation SNTDeviceManagerTest
 
 - (void)setUp {
   [super setUp];
+  self.mockConfigurator = OCMClassMock([SNTConfigurator class]);
+  OCMStub([self.mockConfigurator configurator]).andReturn(self.mockConfigurator);
+  OCMStub([self.mockConfigurator eventLogType]).andReturn(-1);
+
   fclose(stdout);
 }
 
 - (ESResponse *)triggerTestMount:(SNTDeviceManager *)deviceManager
-                          mockES:(MockEndpointSecurity *)mockES {
+                          mockES:(MockEndpointSecurity *)mockES
+                          mockDA:(MockDiskArbitration *)mockDA {
   if (!deviceManager.subscribed) {
     // [deviceManager listen] is synchronous, but we want to asynchronously dispatch it
     // with an enforced timeout to ensure that we never run into issues where the client
@@ -54,10 +62,28 @@
   }
 
   struct statfs *fs = static_cast<struct statfs *>(calloc(1, sizeof(struct statfs)));
-  const char test_mntfromname[] = "/dev/disk2s1";
-  const char test_mntonname[] = "/Volumes/KATE'S 4G";
-  strncpy(fs->f_mntfromname, test_mntfromname, sizeof(test_mntfromname));
-  strncpy(fs->f_mntonname, test_mntonname, sizeof(test_mntonname));
+  NSString *test_mntfromname = @"/dev/disk2s1";
+  NSString *test_mntonname = @"/Volumes/KATE'S 4G";
+  const char *c_mntfromname = [test_mntfromname UTF8String];
+  const char *c_mntonname = [test_mntonname UTF8String];
+
+  strncpy(fs->f_mntfromname, c_mntfromname, MAXPATHLEN);
+  strncpy(fs->f_mntonname, c_mntonname, MAXPATHLEN);
+
+  MockDADisk *disk = [[MockDADisk alloc] init];
+  disk.diskDescription = @{
+    (__bridge NSString *)kDADiskDescriptionDeviceProtocolKey : @"USB",
+    (__bridge NSString *)kDADiskDescriptionMediaRemovableKey : @YES,
+    @"DAVolumeMountable" : @YES,
+    @"DAVolumePath" : test_mntonname,
+    @"DADeviceModel" : @"Some device model",
+    @"DADevicePath" : test_mntonname,
+    @"DADeviceVendor" : @"Some vendor",
+    @"DAAppearanceTime" : @0,
+    @"DAMediaBSDName" : test_mntfromname,
+  };
+
+  [mockDA insert:disk bsdName:test_mntfromname];
 
   ESMessage *m = [[ESMessage alloc] initWithBlock:^(ESMessage *m) {
     m.binaryPath = @"/System/Library/Filesystems/msdos.fs/Contents/Resources/mount_msdos";
@@ -78,6 +104,7 @@
 
   [self waitForExpectations:@[ expectation ] timeout:60.0];
   free(fs);
+
   return got;
 }
 
@@ -85,14 +112,46 @@
   MockEndpointSecurity *mockES = [MockEndpointSecurity mockEndpointSecurity];
   [mockES reset];
 
+  MockDiskArbitration *mockDA = [MockDiskArbitration mockDiskArbitration];
+  [mockDA reset];
+
   SNTDeviceManager *deviceManager = [[SNTDeviceManager alloc] init];
   deviceManager.blockUSBMount = NO;
-  ESResponse *got = [self triggerTestMount:deviceManager mockES:mockES];
-
+  ESResponse *got = [self triggerTestMount:deviceManager mockES:mockES mockDA:mockDA];
   XCTAssertEqual(got.result, ES_AUTH_RESULT_ALLOW);
 }
 
-// TODO(tnek): Write a DiskArbitrationTestUtil similar to the EndpointSecurityTestUtil for
-// verifying that DiskArbitration callbacks get correctly called on device discovery.
+- (void)testRemount {
+  MockEndpointSecurity *mockES = [MockEndpointSecurity mockEndpointSecurity];
+  [mockES reset];
+
+  MockDiskArbitration *mockDA = [MockDiskArbitration mockDiskArbitration];
+  [mockDA reset];
+
+  SNTDeviceManager *deviceManager = [[SNTDeviceManager alloc] init];
+  deviceManager.blockUSBMount = YES;
+  deviceManager.remountArgs = @[ @"noexec", @"rdonly" ];
+
+  ESResponse *got = [self triggerTestMount:deviceManager mockES:mockES mockDA:mockDA];
+
+  XCTAssertEqual(got.result, ES_AUTH_RESULT_DENY);
+  XCTAssertEqual(mockDA.wasRemounted, YES);
+}
+
+- (void)testBlockNoRemount {
+  MockEndpointSecurity *mockES = [MockEndpointSecurity mockEndpointSecurity];
+  [mockES reset];
+
+  MockDiskArbitration *mockDA = [MockDiskArbitration mockDiskArbitration];
+  [mockDA reset];
+
+  SNTDeviceManager *deviceManager = [[SNTDeviceManager alloc] init];
+  deviceManager.blockUSBMount = YES;
+
+  ESResponse *got = [self triggerTestMount:deviceManager mockES:mockES mockDA:mockDA];
+
+  XCTAssertEqual(got.result, ES_AUTH_RESULT_DENY);
+  XCTAssertEqual(mockDA.wasRemounted, NO);
+}
 
 @end