pid_pname.c

基于TDI驱动编写的个人防火墙程序。包括驱动模块、应用层规则配置及加载模块。

// -*- mode: C++; tab-width: 4; indent-tabs-mode: nil -*- (for GNU Emacs)
//
// $Id: pid_pname.c,v 1.2 2002/12/03 13:32:17 dev Exp $

/*
 * Get process name by its pid and all related routines
 */

#include <ntddk.h>

#include "memtrack.h"
#include "pid_pname.h"

/* process list entry */
struct plist_entry {
	struct	plist_entry *next;
	
	// id & name
	ULONG	pid;
	char	*pname;
	KEVENT	*pname_event;

	int		context;
};

/* process list */
static struct {
	struct		plist_entry *head;
	struct		plist_entry *tail;
	KSPIN_LOCK	guard;
} g_plist;

static struct plist_entry *add_ple(ULONG pid, KIRQL *irql);
static struct plist_entry *find_ple(ULONG pid, KIRQL *irql, struct plist_entry **prev);

static VOID ProcessNotifyProc(IN HANDLE ParentId, IN HANDLE ProcessId, IN BOOLEAN Create);

void
pid_pname_init(void)
{
	KeInitializeSpinLock(&g_plist.guard);
	g_plist.head = g_plist.tail = NULL;

	PsSetCreateProcessNotifyRoutine(ProcessNotifyProc, FALSE);
}

// free plist
void
pid_pname_free(void)
{
	KIRQL irql;
	struct plist_entry *ple;

	PsSetCreateProcessNotifyRoutine(ProcessNotifyProc, TRUE);

	KeAcquireSpinLock(&g_plist.guard, &irql);
	for (ple = g_plist.head; ple != NULL;) {
		struct plist_entry *ple2 = ple->next;
		if (ple->pname != NULL)
			free(ple->pname);
		free(ple);
		ple = ple2;
	}
	g_plist.head = g_plist.tail = NULL;
	KeReleaseSpinLock(&g_plist.guard, irql);
}

// try to get pname by pid
BOOLEAN
pid_pname_resolve(ULONG pid, char *buf, int buf_size)
{
	BOOLEAN result;
	KIRQL irql;
	struct plist_entry *ple = find_ple(pid, &irql, NULL);

	if (ple == NULL)
		return FALSE;

	if (ple->pname != NULL) {
		if (buf_size > 0) {
			strncpy(buf, ple->pname, buf_size);
			buf[buf_size - 1] = '\0';
		}
		result = TRUE;
	} else
		result = FALSE;
	
	KeReleaseSpinLock(&g_plist.guard, irql);
	return result;
}

// set pname_event by pid
NTSTATUS
pid_pname_set_event(ULONG pid, KEVENT *event)
{
	KIRQL irql;
	struct plist_entry *ple = find_ple(pid, &irql, NULL);

	if (ple == NULL) {
		// try to add
		ple = add_ple(pid, &irql);
		if (ple == NULL)
			return STATUS_INSUFFICIENT_RESOURCES;
	}

	ple->pname_event = event;

	KeReleaseSpinLock(&g_plist.guard, irql);
	return STATUS_SUCCESS;
}

NTSTATUS
pid_pname_set(ULONG pid, const char *pname, int context)
{
	KIRQL irql;
	struct plist_entry *ple = find_ple(pid, &irql, NULL);
	NTSTATUS status;

	if (ple == NULL) {
		// try to add
		ple = add_ple(pid, &irql);
		if (ple == NULL)
			return STATUS_INSUFFICIENT_RESOURCES;
	}
	
	if (ple->pname == NULL) {
		ple->pname = (char *)malloc_np(strlen(pname) + 1);
		if (ple->pname != NULL) {
			strcpy(ple->pname, pname);
			status = STATUS_SUCCESS;
		} else {
			KdPrint(("[tdi_fw] set_pid_name: malloc_np!\n"));
			status = STATUS_INSUFFICIENT_RESOURCES;
		}

		// and signal event we have name! (almost)
		if (ple->pname_event != NULL)
			KeSetEvent(ple->pname_event, IO_NO_INCREMENT, FALSE);
	}

	ple->context = context;

	KeReleaseSpinLock(&g_plist.guard, irql);
	return status;
}

int
pid_pname_get_context(ULONG pid)
{
	KIRQL irql;
	int context;
	struct plist_entry *ple = find_ple(pid, &irql, NULL);
	if (ple == NULL)
		return 0;

	context = ple->context;
	
	KeReleaseSpinLock(&g_plist.guard, irql);
	return context;
}


struct plist_entry *
add_ple(ULONG pid, KIRQL *irql)
{
	struct plist_entry *ple;

	if (irql != NULL)
		KeAcquireSpinLock(&g_plist.guard, irql);

	// add new entry to g_plist
	ple = (struct plist_entry *)malloc_np(sizeof(*ple));
	if (ple != NULL) {
		memset(ple, 0, sizeof(*ple));
		ple->pid = pid;

		// append
		if (g_plist.tail != NULL) {
			g_plist.tail->next = ple;
			g_plist.tail = ple;
		} else
			g_plist.head = g_plist.tail = ple;

	} else {
		KdPrint(("[tdi_fw] add_ple: malloc_np!\n"));
	
		if (irql != NULL)
			KeReleaseSpinLock(&g_plist.guard, *irql);
	}

	return ple;
}

struct plist_entry *
find_ple(ULONG pid, KIRQL *irql, struct plist_entry **prev)
{
	struct plist_entry *ple, *prev_ple;

	if (irql != NULL)
		KeAcquireSpinLock(&g_plist.guard, irql);

	prev_ple = NULL;
	for (ple = g_plist.head; ple != NULL; ple = ple->next) {
		if (ple->pid == pid) {
			if (prev != NULL)
				*prev = prev_ple;
			return ple;
		}
		prev_ple = ple;
	}

	if (irql != NULL)
		KeReleaseSpinLock(&g_plist.guard, *irql);
	
	return NULL;
}

// notify routine on process creation or removing
VOID
ProcessNotifyProc(IN HANDLE ParentId, IN HANDLE ProcessId, IN BOOLEAN Create)
{
	KIRQL irql;
	struct plist_entry *ple, *prev_ple;

	if (Create) {
		KdPrint(("[tdi_fw] ProcessNotifyProc: create process with pid:%u\n", ProcessId));

		add_ple((ULONG)ProcessId, &irql);
		KeReleaseSpinLock(&g_plist.guard, irql);

	} else {
		// remove entry from plist

		KdPrint(("[tdi_fw] ProcessNotifyProc: remove process with pid:%u\n", ProcessId));

		ple = find_ple((ULONG)ProcessId, &irql, &prev_ple);
		if (ple == NULL)
			return;

		if (prev_ple != NULL)
			prev_ple->next = ple->next;
		else
			g_plist.head = ple->next;

		if (ple->next == NULL)
			g_plist.tail = prev_ple;

		if (ple->pname != NULL) {
			KdPrint(("[tdi_fw] ProcessNotifyProc: pname was %s\n", ple->pname));
			free(ple->pname);
		}

		free(ple);
		KeReleaseSpinLock(&g_plist.guard, irql);
	}
}