/*
 * Copyright (c) 2006 Andy Polyakov <appro@fy.chalmers.se>
 *
 * Build with:
 *
 * cl -Ox -GD -GF -Zl -MD -LD vsync.c ws2_32.lib kernel32.lib ntdll.lib
 *
 * Pre-load as:
 *
 * [HKLM\SOFTWARE\Microsoft\Windows NT\CurrentVersion\DLL_PRELOAD]
 * "wcescomm.exe"="vsync.dll"
 * "wcesmgr.exe"="vsync.dll"
 * "rapimgr.exe"="vsync.dll"
 * [HKLM\SOFTWARE\Microsoft\Windows CE Services\Http server]
 * "HttpSysOverride"=dword:1
 * [-HKLM\SOFTWARE\Citrix\CtxHook\AppInit_Dlls\VIPHook\wcescomm.exe]
 * [-HKLM\SOFTWARE\Citrix\CtxHook\AppInit_Dlls\VIPHook\wcesmgr.exe]
 *
 * See http://fy.chalmers.se/~appro/nt/DLL_PRELOAD/ for further details.
 *
 */
#ifndef _DLL
#error "_DLL is not defined."
#endif

#pragma comment(linker,"/entry:DllMain@12")
#pragma comment(linker,"/section:.text,erw")
#pragma comment(linker,"/merge:.rdata=.text")
#pragma comment(linker,"/merge:.data=.text")

#define UNICODE
#define _UNICODE

#if defined(WIN32) && !defined(_WIN32)
#define _WIN32
#endif

#define _WIN32_WINNT 0x0500

#include <winsock2.h>
#include <ws2spi.h>
#include <ws2tcpip.h>
#include <windows.h>
#include <winbase.h>
#include <winnt.h>
#include <winternl.h>
#include <tchar.h>
#include <stdarg.h>

#ifdef _DEBUG
static VOID DebugOutputA (const char *fmt,...)
{ va_list argv;
  char buf[256];

  va_start(argv,fmt);
    _vsnprintf (buf,sizeof(buf)/sizeof(buf[0])-1,fmt,argv);
    buf[sizeof(buf)/sizeof(buf[0])-1]='\0';
    OutputDebugStringA (buf);
  va_end(argv);
}
#endif

/* Native API wrappers */

__declspec(dllimport) NTSTATUS WINAPI ZwCreateMutant (
    OUT PHANDLE MutantHandle,
    IN ACCESS_MASK DesiredAccess,
    IN POBJECT_ATTRIBUTES ObjectAttributes,
    IN BOOLEAN InitialOwner
    );
static NTSTATUS WINAPI NtCreateMutant_ (
    OUT PHANDLE MutantHandle,
    IN ACCESS_MASK DesiredAccess,
    IN POBJECT_ATTRIBUTES ObjectAttributes,
    IN BOOLEAN InitialOwner
    )
{ OBJECT_ATTRIBUTES *p = ObjectAttributes,ob;
  UNICODE_STRING    us;

    if (p && p->ObjectName &&
        p->ObjectName->Length>14 &&
        !_wcsnicmp(L"global\\",p->ObjectName->Buffer,7))
    {   us = *p->ObjectName;
        us.Buffer += 7, us.Length -= 14, us.MaximumLength =- 14;
        ob = *p;
        ob.ObjectName = &us;
        p = &ob;
    }

  return ZwCreateMutant(
            MutantHandle,
            DesiredAccess,
            p,
            InitialOwner);
}

__declspec(dllimport) NTSTATUS WINAPI ZwOpenMutant (
    OUT PHANDLE MutantHandle,
    IN ACCESS_MASK DesiredAccess,
    IN POBJECT_ATTRIBUTES ObjectAttributes
    );
static NTSTATUS WINAPI NtOpenMutant_ (
    OUT PHANDLE MutantHandle,
    IN ACCESS_MASK DesiredAccess,
    IN POBJECT_ATTRIBUTES ObjectAttributes
    )
{ OBJECT_ATTRIBUTES *p = ObjectAttributes,ob;
  UNICODE_STRING    us;

    if (p && p->ObjectName &&
        p->ObjectName->Length>14 &&
        !_wcsnicmp(L"global\\",p->ObjectName->Buffer,7))
    {   us = *p->ObjectName;
        us.Buffer += 7, us.Length -= 14, us.MaximumLength =- 14;
        ob = *p;
        ob.ObjectName = &us;
        p = &ob;
    }

  return  ZwOpenMutant(
            MutantHandle,
            DesiredAccess,
            p);
}

__declspec(dllimport) NTSTATUS WINAPI ZwCreateSemaphore (
    OUT PHANDLE SemaphoreHandle,
    IN ACCESS_MASK DesiredAccess,
    IN POBJECT_ATTRIBUTES ObjectAttributes,
    IN LONG InitialCount,
    IN LONG MaxCount
    );
static NTSTATUS WINAPI NtCreateSemaphore_ (
    OUT PHANDLE SemaphoreHandle,
    IN ACCESS_MASK DesiredAccess,
    IN POBJECT_ATTRIBUTES ObjectAttributes,
    IN LONG InitialCount,
    IN LONG MaxCount
    )
{ OBJECT_ATTRIBUTES *p = ObjectAttributes,ob;
  UNICODE_STRING    us;

    if (p && p->ObjectName &&
        p->ObjectName->Length>14 &&
        !_wcsnicmp(L"global\\",p->ObjectName->Buffer,7))
    {   us = *p->ObjectName;
        us.Buffer += 7, us.Length -= 14, us.MaximumLength =- 14;
        ob = *p;
        ob.ObjectName = &us;
        p = &ob;
    }

  return ZwCreateSemaphore(
            SemaphoreHandle,
            DesiredAccess,
            p,
            InitialCount,
            MaxCount);
}

__declspec(dllimport) NTSTATUS WINAPI ZwOpenSemaphore (
    OUT PHANDLE SemaphoreHandle,
    IN ACCESS_MASK DesiredAccess,
    IN POBJECT_ATTRIBUTES ObjectAttributes
    );
static NTSTATUS WINAPI NtOpenSemaphore_ (
    OUT PHANDLE SemaphoreHandle,
    IN ACCESS_MASK DesiredAccess,
    IN POBJECT_ATTRIBUTES ObjectAttributes
    )
{ OBJECT_ATTRIBUTES *p = ObjectAttributes,ob;
  UNICODE_STRING    us;

    if (p && p->ObjectName &&
        p->ObjectName->Length>14 &&
        !_wcsnicmp(L"global\\",p->ObjectName->Buffer,7))
    {   us = *p->ObjectName;
        us.Buffer += 7, us.Length -= 14, us.MaximumLength =- 14;
        ob = *p;
        ob.ObjectName = &us;
        p = &ob;
    }

  return ZwOpenSemaphore(
            SemaphoreHandle,
            DesiredAccess,
            p);
}

/* WSP wrappers */

static LPWSPSTARTUP  _WSPStartup;
static WSPPROC_TABLE WSPDispatch_;
static struct in_addr vloopback;

static SOCKET WSPAPI WSPAccept_(
    SOCKET s,
    struct sockaddr FAR * addr,
    LPINT addrlen,
    LPCONDITIONPROC lpfnCondition,
    DWORD_PTR dwCallbackData,
    LPINT lpErrno
    )
{ SOCKET ret=WSPDispatch_.lpWSPAccept (s,addr,addrlen,NULL,(DWORD_PTR)NULL,lpErrno);
  struct sockaddr_in *sin=(void *)addr;

#ifdef _DEBUG
    DebugOutputA("accept %p %p %s",ret,addr,addr?inet_ntoa(sin->sin_addr):"(null)");
#endif
    if (ret!=INVALID_SOCKET && addr!=NULL && addr->sa_family==AF_INET)
    {   if (sin->sin_addr.s_addr == vloopback.s_addr ||
	    sin->sin_addr.s_addr == (vloopback.s_addr|0x0000807F))
            sin->sin_addr.s_addr = 0x0100007F;
    }

  return ret;
}

static int WSPAPI WSPBind_(
    SOCKET s,
    const struct sockaddr FAR * name,
    int namelen,
    LPINT lpErrno
    )
{ struct sockaddr_in *sin=(void *)name;

    if (name->sa_family == AF_INET)
    {   if (sin->sin_addr.s_addr == 0x0100007F)
        {   sin = _alloca(namelen);
            memcpy (sin,name,namelen);
            sin->sin_addr = vloopback;
        }
        else if (sin->sin_addr.s_addr == 0 && sin->sin_port!=0)
        {   sin = _alloca(namelen);
            memcpy (sin,name,namelen);
            sin->sin_addr.s_addr = vloopback.s_addr|0x0000807F;
        }
#ifdef _DEBUG
        DebugOutputA("bind %s:%d",inet_ntoa(sin->sin_addr),ntohs(sin->sin_port));
#endif
    }

  return WSPDispatch_.lpWSPBind (s,(void *)sin,namelen,lpErrno);
}

static int WSPAPI WSPConnect_(
    SOCKET s,
    const struct sockaddr FAR * name,
    int namelen,
    LPWSABUF lpCallerData,
    LPWSABUF lpCalleeData,
    LPQOS lpSQOS,
    LPQOS lpGQOS,
    LPINT lpErrno
    )
{ struct sockaddr_in *sin=(void *)name;
  int ret;

    if (name->sa_family == AF_INET)
    {   if (sin->sin_addr.s_addr == 0x0100007F)
        {   sin = _alloca(namelen);
            memcpy (sin,name,namelen);
            sin->sin_addr = vloopback;
        }
#ifdef _DEBUG
        DebugOutputA("connect %s:%d",inet_ntoa(sin->sin_addr),ntohs(sin->sin_port));
#endif
    }

    ret = WSPDispatch_.lpWSPConnect (s,(void *)sin,namelen,
                    lpCallerData,lpCalleeData,lpSQOS,lpGQOS,lpErrno);

    if (ret!=0 && (void*)sin!=(void*)name && *lpErrno==WSAECONNREFUSED)
    {   sin->sin_addr.s_addr |= 0x0000807F;
#ifdef _DEBUG
        DebugOutputA("retry connect %s:%d",inet_ntoa(sin->sin_addr),ntohs(sin->sin_port));
#endif
        ret = WSPDispatch_.lpWSPConnect (s,(void *)sin,namelen,
                    lpCallerData,lpCalleeData,lpSQOS,lpGQOS,lpErrno);
    }

  return ret;
}

static int WSPAPI WSPGetPeerName_(
    SOCKET s,
    struct sockaddr FAR * name,
    LPINT namelen,
    LPINT lpErrno
    )
{ int ret=WSPDispatch_.lpWSPGetPeerName(s,name,namelen,lpErrno);
  struct sockaddr_in *sin=(void *)name;

    if (ret==0 && name->sa_family==AF_INET)
    {   if (sin->sin_addr.s_addr == vloopback.s_addr ||
            sin->sin_addr.s_addr == (vloopback.s_addr|0x0000807F))
            sin->sin_addr.s_addr = 0x0100007F;
    }

  return ret;
}

static int WSPAPI WSPGetSockName_(
    SOCKET s,
    struct sockaddr FAR * name,
    LPINT namelen,
    LPINT lpErrno
    )
{ int ret=WSPDispatch_.lpWSPGetSockName(s,name,namelen,lpErrno);
  struct sockaddr_in *sin=(void *)name;

    if (ret==0 && name->sa_family==AF_INET)
    {   if (sin->sin_addr.s_addr == vloopback.s_addr ||
            sin->sin_addr.s_addr == (vloopback.s_addr|0x0000807F))
            sin->sin_addr.s_addr = 0x0100007F;
    }

  return ret;
}

static int WSPAPI WSPStartup_(
    IN WORD wVersionRequested,
    OUT LPWSPDATA lpWSPData,
    IN LPWSAPROTOCOL_INFOW lpProtocolInfo,
    IN WSPUPCALLTABLE UpcallTable,
    OUT LPWSPPROC_TABLE lpProcTable
    )
{ int ret = (*_WSPStartup)(
                    wVersionRequested,
                    lpWSPData,
                    lpProtocolInfo,
                    UpcallTable,
                    &WSPDispatch_);

    /* Copy whole table... */
    *lpProcTable                    = WSPDispatch_;
    /* ... and subclass selected functions */
    lpProcTable->lpWSPAccept        = WSPAccept_;
    lpProcTable->lpWSPBind          = WSPBind_;
    lpProcTable->lpWSPConnect       = WSPConnect_;
    lpProcTable->lpWSPGetPeerName   = WSPGetPeerName_;
    lpProcTable->lpWSPGetSockName   = WSPGetSockName_;

  return ret;
}

static PVOID WINAPI GetProcAddress_ (
    HMODULE hModule,
    LPCSTR lpProcName
    )
{ PVOID ret = GetProcAddress(hModule,lpProcName);

    if (ret && !strcmp(lpProcName,"WSPStartup"))
    {   /* Place myself between WS2_32 and WSP */
        _WSPStartup = ret;
        ret = WSPStartup_;
    }

  return ret;
}

PVOID   *__NtCreateMutant       =NULL,
        *__NtOpenMutant         =NULL,
        *__NtCreateSemaphore    =NULL,
        *__NtOpenSemaphore      =NULL,
        *__GetProcAddress       =NULL;

BOOL WINAPI DllMain (HINSTANCE h, DWORD reason, LPVOID junk)
{ DWORD                     acc0,acc1,sess;
  HMODULE                   hmod;
  IMAGE_DOS_HEADER         *dos_header;
  IMAGE_NT_HEADERS         *nt_headers;
  IMAGE_DATA_DIRECTORY     *dir;
  IMAGE_IMPORT_DESCRIPTOR  *idesc;
  IMAGE_THUNK_DATA         *thunk;
  static void              *page0,*page1;
  static size_t             plen0,plen1;

    switch (reason)
    { case DLL_PROCESS_ATTACH:

        DisableThreadLibraryCalls(h);

        if (!ProcessIdToSessionId(GetCurrentProcessId(),&sess) || sess==0)
            return FALSE;

        /* 0x7F524450 is 127.'R'.'D'.'P' */
        vloopback.s_addr = htonl(0x7F524450+sess);

        if (!(hmod=GetModuleHandle(_T("KERNEL32.DLL"))))
        {   OutputDebugStringA("KERNEL32.DLL not found?");
            return FALSE;
        }

        dos_header = (IMAGE_DOS_HEADER *)hmod;
        nt_headers = (IMAGE_NT_HEADERS *)((char *)hmod + dos_header->e_lfanew);

        dir=&nt_headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];

        idesc=(IMAGE_IMPORT_DESCRIPTOR *)((char *)hmod + dir->VirtualAddress);
        while (idesc->Name)
        {   if (!_stricmp((char *)hmod+idesc->Name,"NTDLL.DLL")) break;
            idesc++;
        }

        if (!idesc->Name)
        {   OutputDebugStringA("Can't locate NTDLL.DLL import descriptor");
            return FALSE;
        }

        page0 = (char *)hmod+idesc->FirstThunk;

        for (thunk=(IMAGE_THUNK_DATA *)page0;
             thunk->u1.Function && thunk->u1.Function!=(ULONG_PTR)ZwCreateMutant;
             thunk++) ;
        if (thunk->u1.Function) __NtCreateMutant=(PVOID *)thunk;

        for (thunk=(IMAGE_THUNK_DATA *)page0;
             thunk->u1.Function && thunk->u1.Function!=(ULONG_PTR)ZwOpenMutant;
             thunk++) ;
        if (thunk->u1.Function) __NtOpenMutant=(PVOID *)thunk;

        for (thunk=(IMAGE_THUNK_DATA *)page0;
             thunk->u1.Function && thunk->u1.Function!=(ULONG_PTR)ZwCreateSemaphore;
             thunk++) ;
        if (thunk->u1.Function) __NtCreateSemaphore=(PVOID *)thunk;

        for (thunk=(IMAGE_THUNK_DATA *)page0;
             thunk->u1.Function && thunk->u1.Function!=(ULONG_PTR)ZwOpenSemaphore;
             thunk++) ;
        if (thunk->u1.Function) __NtOpenSemaphore=(PVOID *)thunk;

        for (thunk=(IMAGE_THUNK_DATA *)page0;thunk->u1.Function;thunk++) ;

        plen0 = (size_t)thunk-(size_t)page0;

        if (!(hmod=GetModuleHandle(_T("WS2_32.DLL"))))
        {   OutputDebugStringA("WS2_32.DLL not found?");
            return FALSE;
        }

        dos_header = (IMAGE_DOS_HEADER *)hmod;
        nt_headers = (IMAGE_NT_HEADERS *)((char *)hmod + dos_header->e_lfanew);

        dir=&nt_headers->OptionalHeader.DataDirectory[IMAGE_DIRECTORY_ENTRY_IMPORT];

        idesc=(IMAGE_IMPORT_DESCRIPTOR *)((char *)hmod + dir->VirtualAddress);
        while (idesc->Name)
        {   if (!_stricmp((char *)hmod+idesc->Name,"KERNEL32.DLL")) break;
            idesc++;
        }

        if (!idesc->Name)
        {   OutputDebugStringA("Can't locate KERNEL32.DLL import descriptor");
            return FALSE;
        }

        page1 = (char *)hmod+idesc->FirstThunk;

        for (thunk=(IMAGE_THUNK_DATA *)page1;
             thunk->u1.Function && thunk->u1.Function!=(ULONG_PTR)GetProcAddress;
             thunk++) ;
        if (thunk->u1.Function) __GetProcAddress=(PVOID *)thunk;

        for (thunk=(IMAGE_THUNK_DATA *)page1;thunk->u1.Function;thunk++) ;

        plen1 = (size_t)thunk-(size_t)page1;

        if (!VirtualProtect (page0,plen0,PAGE_EXECUTE_READWRITE,&acc0))
        {   OutputDebugStringA("Unable to unlock KERNEL32.DLL Thunk Table");
            return FALSE;
        }
        if (!VirtualProtect (page1,plen1,PAGE_EXECUTE_READWRITE,&acc1))
        {   OutputDebugStringA("Unable to unlock WS2_32.DLL Thunk Table");
            return FALSE;
        }

        if (__NtCreateMutant)       *__NtCreateMutant   =NtCreateMutant_;
        if (__NtOpenMutant)         *__NtOpenMutant     =NtOpenMutant_;
        if (__NtCreateSemaphore)    *__NtCreateSemaphore=NtCreateSemaphore_;
        if (__NtOpenSemaphore)      *__NtOpenSemaphore  =NtOpenSemaphore_;

        if (__GetProcAddress)       *__GetProcAddress   =GetProcAddress_;

        VirtualProtect (page0,plen0,acc0,&acc0);
        VirtualProtect (page1,plen1,acc1,&acc1);

        break;

      case DLL_PROCESS_DETACH:

        VirtualProtect (page0,plen0,PAGE_EXECUTE_READWRITE,&acc0);
        VirtualProtect (page1,plen1,PAGE_EXECUTE_READWRITE,&acc1);

        if (__NtCreateMutant)       *__NtCreateMutant   =ZwCreateMutant;
        if (__NtOpenMutant)         *__NtOpenMutant     =ZwOpenMutant;
        if (__NtCreateSemaphore)    *__NtCreateSemaphore=ZwCreateSemaphore;
        if (__NtOpenSemaphore)      *__NtOpenSemaphore  =ZwOpenSemaphore;

        if (__GetProcAddress)       *__GetProcAddress   =GetProcAddress;

        VirtualProtect (page0,plen0,acc0,&acc0);
        VirtualProtect (page1,plen1,acc1,&acc1);

        break;
    }

  return TRUE;
}