/////////////////////////////////////////////////////////////////////////////
// Name:        nvwa_loader.c
// Purpose:		Loader process that analyses program output and determines
//				whether memory leaks have been detected
// Author:      Cesar Mauri Loba (cesar at crea-si dot com)
// Modified by: 
// Created:     13/01/2011
// Copyright:   (C) 2011 Cesar Mauri Loba - CREA Software Systems
// License:
//
// This software is provided 'as-is', without any express or implied
// warranty.  In no event will the authors be held liable for any
// damages arising from the use of this software.
//
// Permission is granted to anyone to use this software for any purpose,
// including commercial applications, and to alter it and redistribute
// it freely, subject to the following restrictions:
//
// 1. The origin of this software must not be misrepresented; you must
//    not claim that you wrote the original software.  If you use this
//    software in a product, an acknowledgement in the product
//    documentation would be appreciated but is not required.
// 2. Altered source versions must be plainly marked as such, and must
//    not be misrepresented as being the original software.
// 3. This notice may not be removed or altered from any source
//    distribution.
/////////////////////////////////////////////////////////////////////////////
#include <stdio.h>
#include <sys/types.h>
#include <signal.h>
#include <assert.h>
#include <string>
#include <algorithm>
#include <iostream>

#if defined (WIN32)
#include <windows.h>
#include <winbase.h>
#include <process.h>
#include <io.h>
#include <fcntl.h>
#else
#include <sys/wait.h>
#include <unistd.h>
#endif


#define BUFFER_SIZE 4096

/* 
	parses buffer to find specific strings.
	Return:
		0 -> no string found
		1 -> last string found was "nvwa: no memory leaks detected\n"
		2 -> last string found was "nvwa: memory leaks detected\n"
*/
void parse_buffer (char* buff, size_t size, int* result)
{
	/* quick and dirty implementation. FIXME: use regular expressions */
	const char expr1[]= "nvwa: no memory leaks detected";
	static int expr1_still_matches= 1;
	const char expr2[]= "nvwa: memory leaks detected";
	static int expr2_still_matches= 1;
	static size_t i_expr= 0;	

	//int result= 0;
	size_t i_buff= 0;
		
	while (i_buff< size) {
		// Matches 1
		if (expr1_still_matches){
			if (i_expr< sizeof(expr1)-2) {
				if (expr1[i_expr]!= buff[i_buff])
					// Match not found
					expr1_still_matches= 0;
			}			
			else if (i_expr== sizeof(expr1)-2 && expr1[i_expr]== buff[i_buff]) {
				//printf ("expr1 match\n");
				*result= 1; // Match expr1 found!
				expr1_still_matches= 0;
			}
			else
				expr1_still_matches= 0;
		}
		
		// Matches 2
		if (expr2_still_matches){
			if (i_expr< sizeof(expr2)-2) {
				if (expr2[i_expr]!= buff[i_buff])
					// Match not found
					expr2_still_matches= 0;
			}			
			else if (i_expr== sizeof(expr2)-2 && expr2[i_expr]== buff[i_buff]) {
				//printf ("expr2 match\n");
				*result= 2; // Match expr2 found!
				expr2_still_matches= 0;
			}
			else
				expr2_still_matches= 0;
		}
		
		if (!expr1_still_matches && !expr2_still_matches) {
			expr1_still_matches= expr2_still_matches= 1;
			if (i_expr== 0) ++i_buff;
			else i_expr= 0;
		}
		else {
			++i_expr;
			++i_buff;
		}
	}
}

#if defined(WIN32)
int process_child_output(HANDLE handle)
#else
int process_child_output(int handle)
#endif
{
	/* parse child's process output */
	char buffer[BUFFER_SIZE + 1];
	int leak_condition= 0;

#if defined(WIN32)
	DWORD read_bytes;
	BOOL bSuccess;

	DWORD errDEBUG;
#else
	size_t read_bytes;	
#endif
	
	for (;;) {		
#if defined(WIN32)
		read_bytes= 0;
		bSuccess = ReadFile( handle, buffer, BUFFER_SIZE, &read_bytes, NULL);
		errDEBUG= GetLastError();
		if (!bSuccess) read_bytes= 0;
#else
		read_bytes= read(handle, buffer, BUFFER_SIZE);
#endif		
		if (read_bytes== 0) break;
		buffer[read_bytes]= '\0';
		fwrite (buffer, 1, read_bytes, stderr);
		parse_buffer (buffer, read_bytes, &leak_condition);			
	}

	return leak_condition;
}


#if defined(WIN32)
int do_win32 (const char* cmdline)
{
	SECURITY_ATTRIBUTES saAttr;
	HANDLE hChildStderrRd = NULL;
	HANDLE hChildStderrWr = NULL;
	STARTUPINFO si;
	PROCESS_INFORMATION pi;
	int len = 0;
	BSTR unicodestr;
	int retval= 0;
	DWORD exitCode= 0;


	// Set the bInheritHandle flag so pipe handles are inherited 
	saAttr.nLength = sizeof(SECURITY_ATTRIBUTES);
	saAttr.bInheritHandle = TRUE;
	saAttr.lpSecurityDescriptor = NULL;

	// Create a pipe for the child process's STDERR
	if ( ! CreatePipe(&hChildStderrRd, &hChildStderrWr, &saAttr, 0) ) {
		fprintf (stderr, "pipe creation failed\n");
		return -1;
	}
	
	// Ensure the read handle to the pipe for STDERR is not inherited.
	if (!SetHandleInformation(hChildStderrRd, HANDLE_FLAG_INHERIT, 0) ) {
		fprintf (stderr, "SetHandleInformation failed\n");
		return -1;
	}

	// Create child process
	len = lstrlenA(cmdline);
	unicodestr = SysAllocStringLen(NULL, len);
	MultiByteToWideChar(CP_ACP, 0, cmdline, len, unicodestr, len);

	ZeroMemory(&si,sizeof(si));
	ZeroMemory(&pi,sizeof(pi));

	si.cb=sizeof(si);
	si.hStdError= hChildStderrWr;
	si.dwFlags |= STARTF_USESTDHANDLES;	
	if(!CreateProcess(unicodestr, unicodestr,0,0,TRUE,0,0,0,&si,&pi)) {
		fprintf (stderr, "could not start process %s. Error: %d\n", cmdline, GetLastError());
		return -1;
	}
	SysFreeString(unicodestr);
    
	// Close handles to the child process and its primary thread.
	// Some applications might keep these handles to monitor the status
	// of the child process, for example
	//CloseHandle(piProcInfo.hProcess);
	//CloseHandle(piProcInfo.hThread);

	// Close the write end of the pipe
	CloseHandle(hChildStderrWr);

	retval= process_child_output(hChildStderrRd);

	if (WaitForSingleObject(pi.hProcess,INFINITE)!= 0) {
		fprintf (stderr, "WaitForSingleObject error\n");
		return -1;		
	}

	if (!GetExitCodeProcess(pi.hProcess, &exitCode)) {
		fprintf (stderr, "GetExitCodeProcess failed\n");
		return -1;
	}

	if (exitCode!= 0) return exitCode;

	if (retval== 2) return -1;	// Leak detected
	
	return 0;
}
#else
int do_unix (const char* cmdline)
{
	int pipefd[2];	
	/* init pipe */
	if (pipe(pipefd)!= 0) {	
		fprintf (stderr, "pipe creation failed\n");
		return -1;
	}

	/* create child process to run program */
	pid_t pid= fork();
	if (pid== -1) {
		perror ("fork");
		return -1;
	}
	
	if (pid== 0) {
		/* child process */
		
		/* close read end of the pipe */
		close (pipefd[0]);
				
		/* redirect stderr */
		if (dup2(pipefd[1], STDERR_FILENO)== -1) {
			close (pipefd[1]);
			return -1;
		}
		
		return execl (cmdline, cmdline, NULL);		
	}
	else
	{
		/* parent process */
		
		/* close write end of the pipe */
		close (pipefd[1]);
		
		/* parse child's process output */
		int retval= process_child_output(pipefd[0]);
			
		/* wait for termination */
		int status;
		if (wait(&status)== -1) {
			perror ("wait");
			return -1;
		}
		
		if (!WIFEXITED(status)) return WEXITSTATUS(status);
	
		if (retval== 2) return -1;	// Leak detected
	
		return 0;
	}
	
	assert (0);
	return 0;
}
#endif

/*
bool hasEnding (std::string const &fullString, std::string const &ending)
{
    if (fullString.length() > ending.length()) {
        return (0 == fullString.compare (fullString.length() - ending.length(), ending.length(), ending));
    } else {
        return false;
    }
}
*/

int main (int argc, char **argv)
{	
	if (argc!= 2) {
		fprintf (stderr, "not enough arguments\n");
		return -1;
	}
#if defined(WIN32)
	std::string cmd(argv[1]);

	std::replace(cmd.begin(), cmd.end(), '/', '\\');
	
	// Append .exe extension when needed
	//if (!hasEnding (cmd, ".exe")) cmd.append(".exe");

	//std::cout << cmd << std::endl;

	return do_win32(cmd.c_str());
#else
	return do_unix(argv[1]);
#endif
}
