/*++


Module Name:

    sync.c

Abstract:

    Program to flush disk cache.

Author:

    Ahdrey Shedel <andreys@cr.cyco.com>

   You may distribute under the terms of the GNU General Public License
   You can contact author at andreys@cr.cyco.com or http://www.chat.ru/~ashedel

--*/


#define WIN32_LEAN_AND_MEAN
#define _DLL

#include <stdarg.h>
#include <windows.h>
#include <shellapi.h>

#pragma warning(push, 4)

#pragma warning(disable:4204) // nonstandard extension used : non-constant aggregate initializer

#pragma comment(lib,"shell32.lib")
#pragma comment(lib,"kernel32.lib")
#pragma comment(lib,"ntdll.lib")
#pragma comment(linker, "-entry:raw_main")
#pragma comment(linker, "-opt:nowin98")
#pragma comment(linker, "-opt:ref")
#pragma comment(linker, "-merge:.rdata=.text")
#pragma comment(linker, "-subsystem:console")


typedef ULONG NTSTATUS;

typedef struct {
	USHORT Length;
	USHORT MaximumLength;
	PWSTR  Buffer;
} UNICODE_STRING, *PUNICODE_STRING;


typedef struct {
	NTSTATUS Status;
	ULONG Information;
} IO_STATUS_BLOCK, *PIO_STATUS_BLOCK;


typedef struct {
	ULONG Length;
	HANDLE RootDirectory;
	PUNICODE_STRING ObjectName;
	ULONG Attributes;
	PVOID SecurityDescriptor;
	PVOID SecurityQualityOfService;
} OBJECT_ATTRIBUTES, *POBJECT_ATTRIBUTES;


NTSYSAPI
VOID
NTAPI
RtlInitUnicodeString(
    PUNICODE_STRING DestinationString,
    PCWSTR SourceString
    );


NTSYSAPI
NTSTATUS
NTAPI
NtClose(
    IN HANDLE Handle
    );

NTSYSAPI
NTSTATUS
NTAPI
NtFlushBuffersFile(
    IN HANDLE FileHandle,
    OUT PIO_STATUS_BLOCK IoStatusBlock
   );

NTSYSAPI
NTSTATUS
NTAPI
NtOpenFile(
    OUT PHANDLE FileHandle,
    IN ACCESS_MASK DesiredAccess,
    IN POBJECT_ATTRIBUTES ObjectAttributes,
    OUT PIO_STATUS_BLOCK IoStatusBlock,
    IN ULONG ShareAccess,
    IN ULONG OpenOptions
    );


#define InitializeObjectAttributes( p, n, a, r, s ) { \
    (p)->Length = sizeof( OBJECT_ATTRIBUTES );\
    (p)->RootDirectory = r;\
    (p)->Attributes = a;\
    (p)->ObjectName = n;\
    (p)->SecurityDescriptor = s;\
    (p)->SecurityQualityOfService = NULL;\
    }



#define OBJ_CASE_INSENSITIVE    0x00000040L
#define FILE_SYNCHRONOUS_IO_NONALERT            0x00000020
#define STATUS_SUCCESS 0




static const WCHAR NTTLL_DLL[]  = L"ntdll.dll";

static const char szFailed[]    = "failed: ";
static const char szDriveOpen[] = "Drive open ";
static const char szFlush[]     = "Flush ";
static const char szSpace[]     = " ";
static const char szFlushing[]  = "Flushing ";
static const char szSucceeded[] = "Succeeded.\n";
static const char szCRLF[]      = "\n";
static const char szUsage[]     = "SYNC: Disk Flusher for Windows NT\nCopyright (C) 1999 Andrey Shedel (andreys@cr.cyco.com)\nUsage: sync <path_name>\n";






void
__fastcall
Puts(IN const char* msg)
{
	ULONG Written;
	WriteFile(GetStdHandle(STD_OUTPUT_HANDLE), msg,
			  strlen(msg), &Written, NULL);
}


__inline
char*
__fastcall
GetErrorString(NTSTATUS error)
{
	char* lpMsgBuf;

	FormatMessageA(
				   FORMAT_MESSAGE_ALLOCATE_BUFFER |
				   FORMAT_MESSAGE_IGNORE_INSERTS |
				   FORMAT_MESSAGE_FROM_SYSTEM |
				   FORMAT_MESSAGE_FROM_HMODULE,
				   GetModuleHandleW(NTTLL_DLL),
				   error,
				   0,
				   (char*)&lpMsgBuf, 0, NULL);

	return lpMsgBuf;
}



void
__fastcall
pstatus(NTSTATUS Status)
{
	if(STATUS_SUCCESS != Status)
	{
		char* s = GetErrorString(Status);

		Puts(szFailed);

		if(NULL != s)
		{
			Puts(s);
			LocalFree((HLOCAL)s);
		}
	}
}







NTSTATUS
__fastcall
_Flush (
    IN int DriveNr
   )
{
	NTSTATUS Status;
	UNICODE_STRING UnicodeString;
	OBJECT_ATTRIBUTES ObjectAttributes;
	IO_STATUS_BLOCK IoStatusBlock;
	HANDLE Handle;
	WCHAR Buffer[6] = {L'\\', L'?', L'?', L'\\', (WCHAR)(L'A' + DriveNr), L':'};

	UnicodeString.Length = UnicodeString.MaximumLength = 12;
	UnicodeString.Buffer = &Buffer[0];


	InitializeObjectAttributes(&ObjectAttributes,
							   &UnicodeString,
							   OBJ_CASE_INSENSITIVE,
							   NULL,
							   NULL );


	Status = NtOpenFile(&Handle,
						SYNCHRONIZE | FILE_READ_DATA | FILE_WRITE_DATA,
						&ObjectAttributes,
						&IoStatusBlock,
						FILE_SHARE_WRITE | FILE_SHARE_READ,
						FILE_SYNCHRONOUS_IO_NONALERT);

	if(STATUS_SUCCESS != Status)
	{
		Puts(szDriveOpen); pstatus(Status);
		return Status;
	}


	//__try {

		Status = NtFlushBuffersFile(Handle, &IoStatusBlock);

		if(STATUS_SUCCESS != Status)
		{
			Puts(szFlush); pstatus(Status);
		}

	//}
	//__finally {

		(void)NtClose(Handle);

	//}

	return Status;
}



void
__fastcall
_FlushLocal(void)
{
	ULONG drives = GetLogicalDrives();
	int i;
	WCHAR root[4] = {L'A', L':', L'\\', L'\0'};
	NTSTATUS Status;
	BOOLEAN Succeeded = TRUE;
	char ansi_root[3] = {'A', ':', '\0'};

	Puts(szFlushing);

	for(i = 0; i < 32; i++)
	{
		if(1 == (drives & 1))
		{
			root[0] = (WCHAR)(L'A' + i);

			if(DRIVE_FIXED == (DRIVE_FIXED & GetDriveTypeW(root)))
			{
				ansi_root[0] = (char)('A' + i);
				Puts(ansi_root);
				Puts(szSpace);

				Status = _Flush(i);

				if(Succeeded)
				{
					Succeeded = (BOOLEAN)(STATUS_SUCCESS == Status);
				}
			}
		}

		drives = drives >> 1;
	}

	if(Succeeded)
	{
		Puts(szSucceeded);
	}

	Puts(szCRLF);
}





int
__fastcall
real_wmain(int argc, unsigned short** argv)
{
	PCWSTR p;
	int DriveNr;
	char ansi_drive[3];

	if(argc < 2)
	{
		_FlushLocal();
		return 0;
	}

	p = argv[1];

	if(!((((*p <= L'z') && (*p >= L'a')) || ((*p <= L'Z') && (*p >= L'A'))) &&
		 (L':' == *(p + 1)) &&
		 (L'\0' == *(p + 2))))
	{
		Puts(szUsage);
		return 0;
	}

	DriveNr = *p - ((*p > L'a') ? L'a' : L'A');

	ansi_drive[0] = (char)('A' + DriveNr);
	ansi_drive[1] = ':';
	ansi_drive[2] = '\0';

	Puts(szFlushing);
	Puts(ansi_drive);
	Puts(szSpace);

	if(STATUS_SUCCESS == _Flush(DriveNr))
	{
		Puts(szSucceeded);
	}

	Puts(szCRLF);
	return 0;
}



void __cdecl raw_main(void)
{
	int argc;
	LPWSTR* argv = CommandLineToArgvW(GetCommandLineW(), &argc);

	ExitProcess(real_wmain(argc, argv));
}