/*
 *  kexbasen\shell32\SHGetFolderPath.c
 *
 *  Copyright (C) 2009-2010, Xeno86, Tihiy
 *  Copyright (C) 2019, jumper
 *
 *  This file is part of KernelEx source code.
 *
 *  KernelEx is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published
 *  by the Free Software Foundation; version 2 of the License.
 *
 *  KernelEx is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with GNU Make; see the file COPYING.  If not, write to
 *  the Free Software Foundation, 675 Mass Ave, Cambridge, MA 02139, USA.
 */

#include <windows.h>
#include <shlobj.h>
#include <shlwapi.h>
#include "delayload.h"
#include "kexcoresdk.h"
#include "folderfix.h"


typedef int (WINAPI *PROC5)(PVOID, int, PVOID, int, PVOID);

#define RELOAD  (PROC)(-1)
static  PROC    pSHGetFolderPathA = RELOAD;
static  PROC    pSHGetFolderPathW = RELOAD;
static  HMODULE hShfolder;


/*
 * Hook Shfolder.dll and process DLL_PROCESS_DETACH to protect against 
 * buggy apps that call FreeLibrary too many times.
 */
#include <pshpack1.h>
typedef struct
{	BYTE jmp;
	DWORD func;
} LONGJMP, *PLONGJMP;
#include <poppack.h>

static BYTE prev_entry[sizeof(LONGJMP)];

static BOOL APIENTRY shfolder_entry(HINSTANCE inst, DWORD reason, BOOL load_static)
{
	if (reason == DLL_PROCESS_DETACH)
	{
		DBGPRINTF(("kexbasen: shfolder detached\n"));
		pSHGetFolderPathA = RELOAD;
		pSHGetFolderPathW = RELOAD;
		hShfolder = NULL;
	}
	return TRUE;
}

static void protect_shfolder()
{
	DWORD fold;
	DWORD fnew;
	DWORD entry_addr;
	PLONGJMP ljmp;
	DWORD dwShfolder = (DWORD) hShfolder;
	IMAGE_DOS_HEADER* dosh = (IMAGE_DOS_HEADER*) hShfolder;
	IMAGE_NT_HEADERS* nth = (IMAGE_NT_HEADERS*) (dwShfolder + dosh->e_lfanew);

	entry_addr = dwShfolder + nth->OptionalHeader.AddressOfEntryPoint;
	memcpy(prev_entry, (PVOID) entry_addr, sizeof(LONGJMP));
	VirtualProtect((PVOID) entry_addr, sizeof(LONGJMP), PAGE_READWRITE, &fold);
	ljmp = (PLONGJMP) entry_addr;
	ljmp->jmp = 0xe9; //jmp near rel
	ljmp->func = (DWORD) shfolder_entry - (entry_addr + sizeof(LONGJMP));
	VirtualProtect((PVOID) entry_addr, sizeof(LONGJMP), fold, &fnew);
}

static void unprotect_shfolder()
{
	DWORD fold;
	DWORD fnew;
	DWORD entry_addr;
	DWORD dwShfolder = (DWORD) hShfolder;
	IMAGE_DOS_HEADER* dosh = (IMAGE_DOS_HEADER*) hShfolder;
	IMAGE_NT_HEADERS* nth = (IMAGE_NT_HEADERS*) (dwShfolder + dosh->e_lfanew);

	entry_addr = dwShfolder + nth->OptionalHeader.AddressOfEntryPoint;
	VirtualProtect((PVOID) entry_addr, sizeof(LONGJMP), PAGE_READWRITE, &fold);
	memcpy((PVOID) entry_addr, prev_entry, sizeof(LONGJMP));
	VirtualProtect((PVOID) entry_addr, sizeof(LONGJMP), fold, &fnew);
}

void uninit_SHGetFolderPath()
{
	if (!hShfolder) return;

	unprotect_shfolder();
	hShfolder = NULL;
}


static PROC LoadShProc(LPSTR proc)
{
	static	const char ShfolderFn[] = "SHFOLDER.DLL";
	PROC	ret;
	DWORD	lasterr = GetLastError();
	
	//first try with shell32
	ret = GetDelayProc(Shell32, proc);
	
	//fallback to shfolder
	if (!ret)
	{
		if (!hShfolder) 
		{
			hShfolder = LoadLibrary(ShfolderFn);
			if (hShfolder)
				protect_shfolder();
		}
		if (hShfolder) 
			ret = kexGetProcAddress(hShfolder, proc);
	}
	SetLastError(lasterr);
	return ret;
}


/* MAKE_EXPORT SHGetFolderPathA_fix=SHGetFolderPathA */
HRESULT WINAPI SHGetFolderPathA_fix(
		HWND hwndOwner, int nFolder, HANDLE hToken, DWORD dwFlags, LPSTR pszPath)
{
	if (pSHGetFolderPathA == RELOAD)
		pSHGetFolderPathA =  LoadShProc("SHGetFolderPathA");
	if (pSHGetFolderPathA == NULL)
		return E_NOTIMPL;
	nFolder = folder_fix(nFolder);
	return ((PROC5)pSHGetFolderPathA)(hwndOwner, nFolder, NULL, dwFlags, pszPath);
}

/* MAKE_EXPORT SHGetFolderPathW_fix=SHGetFolderPathW */
HRESULT WINAPI SHGetFolderPathW_fix(
		HWND hwndOwner, int nFolder, HANDLE hToken, DWORD dwFlags, LPWSTR pszPath)
{
	if (pSHGetFolderPathW == RELOAD)
		pSHGetFolderPathW =  LoadShProc("SHGetFolderPathW");
	if (pSHGetFolderPathW == NULL)
		return E_NOTIMPL;
	nFolder = folder_fix(nFolder);
	return ((PROC5)pSHGetFolderPathW)(hwndOwner, nFolder, NULL, dwFlags, pszPath);
}

/* MAKE_EXPORT SHGetFolderPathAndSubDirA_xp=SHGetFolderPathAndSubDirA */
HRESULT SHGetFolderPathAndSubDirA_xp (
  HWND hwnd, int csidl, HANDLE hToken, DWORD dwFlags, LPCSTR pszSubDir, LPSTR pszPath
)
{
  HRESULT hr = SHGetFolderPathA_fix (hwnd, csidl, hToken, dwFlags, pszPath);
  if (hr == S_OK) PathAppendA (pszPath, pszSubDir);
  return hr;
}

/* MAKE_EXPORT SHGetFolderPathAndSubDirW_xp=SHGetFolderPathAndSubDirW */
HRESULT SHGetFolderPathAndSubDirW_xp (
  HWND hwnd, int csidl, HANDLE hToken, DWORD dwFlags, LPCWSTR pwszSubDir, LPWSTR pwszPath
)
{
  HRESULT hr = SHGetFolderPathW_fix (hwnd, csidl, hToken, dwFlags, pwszPath);
  if (hr == S_OK) PathAppendW (pwszPath, pwszSubDir);
  return hr;
}