mirrors/fort
1
0
Fork 0
mirror of https://github.com/tnodir/fort synced 2024-11-15 06:15:15 +00:00
Code Issues Packages Projects Releases Wiki Activity

Driver: Robust buffer handling.

Browse Source
This commit is contained in:
Nodir Temirkhodjaev 2019-03-12 10:52:35 +05:00
parent 86b78424af
commit 734017b73f
3 changed files with 44 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

39
src/driver/fortbuf.c
View File

@ -76,7 +76,7 @@ fort_buffer_data_alloc (PFORT_BUFFER buf, UINT32 len)
}
static void
fort_buffer_data_free (PFORT_BUFFER buf)
fort_buffer_data_shift (PFORT_BUFFER buf)
{
PFORT_BUFFER_DATA data = buf->data_head;
@ -123,13 +123,16 @@ static NTSTATUS
fort_buffer_prepare (PFORT_BUFFER buf, UINT32 len, PCHAR *out,
PIRP *irp, ULONG_PTR *info)
{
const ULONG out_len = buf->out_len;
/* Check pending buffer */
if (buf->out_len && buf->out_top < buf->out_len) {
if (buf->data_head == NULL
&& out_len != 0 && buf->out_top < out_len) {
const UINT32 out_top = buf->out_top;
UINT32 new_top = out_top + len;
/* Is it time to flush logs? */
if (buf->out_len - new_top < FORT_LOG_SIZE_MAX) {
if (out_len - new_top < FORT_LOG_SIZE_MAX) {
if (irp != NULL) {
buf->out_len = 0;
@ -234,15 +237,14 @@ fort_buffer_xmove (PFORT_BUFFER buf, PIRP irp, PVOID out, ULONG out_len,
data = buf->data_head;
*info = buf_top = (data ? data->top : 0);
if (!buf_top) {
if (buf->out_len) {
status = STATUS_UNSUCCESSFUL;
} else if (out_len < FORT_LOG_SIZE_MAX) {
status = STATUS_BUFFER_TOO_SMALL;
if (buf_top == 0) {
if (buf->out_len != 0) {
status = STATUS_UNSUCCESSFUL; /* collision */
} else {
buf->irp = irp;
buf->out = out;
buf->out_len = out_len;
buf->out_top = 0;
status = STATUS_PENDING;
}
goto end;
@ -255,7 +257,7 @@ fort_buffer_xmove (PFORT_BUFFER buf, PIRP irp, PVOID out, ULONG out_len,
RtlCopyMemory(out, data->p, buf_top);
fort_buffer_data_free(buf);
fort_buffer_data_shift(buf);
end:
KeReleaseInStackQueuedSpinLock(&lock_queue);
@ -276,7 +278,7 @@ fort_buffer_cancel_pending (PFORT_BUFFER buf, PIRP irp, ULONG_PTR *info)
buf->irp = NULL;
buf->out_len = 0;
if (buf->out_top) {
if (buf->out_top != 0) {
*info = buf->out_top;
buf->out_top = 0;
@ -303,9 +305,22 @@ fort_buffer_dpc_end (PKLOCK_QUEUE_HANDLE lock_queue)
static void
fort_buffer_dpc_flush_pending (PFORT_BUFFER buf, PIRP *irp, ULONG_PTR *info)
{
const UINT32 out_top = buf->out_top;
UINT32 out_top = buf->out_top;
if (out_top) {
/* Move data from buffer to pending */
if (out_top == 0 && buf->out_len != 0) {
PFORT_BUFFER_DATA data = buf->data_head;
out_top = (data ? data->top : 0);
if (out_top != 0) {
RtlCopyMemory(buf->out, data->p, out_top);
fort_buffer_data_shift(buf);
}
}
if (out_top != 0) {
*info = out_top;
buf->out_top = 0;

26
src/driver/fortdrv.c
View File

@ -946,7 +946,7 @@ fort_callout_timer (void)
fort_stat_dpc_begin(stat, &stat_lock_queue);
/* Flush traffic statistics */
while (stat->proc_active_count) {
while (stat->proc_active_count != 0) {
const UINT16 proc_count =
(stat->proc_active_count < FORT_LOG_STAT_BUFFER_PROC_COUNT)
? stat->proc_active_count : FORT_LOG_STAT_BUFFER_PROC_COUNT;
@ -974,7 +974,9 @@ fort_callout_timer (void)
fort_stat_dpc_end(&stat_lock_queue);
/* Flush pending buffer */
fort_buffer_dpc_flush_pending(buf, &irp, &info);
if (irp == NULL) {
fort_buffer_dpc_flush_pending(buf, &irp, &info);
}
/* Unlock buffer */
fort_buffer_dpc_end(&buf_lock_queue);
@ -1152,18 +1154,22 @@ fort_device_control (PDEVICE_OBJECT device, PIRP irp)
PVOID out = irp->AssociatedIrp.SystemBuffer;
const ULONG out_len = irp_stack->Parameters.DeviceIoControl.OutputBufferLength;
status = fort_buffer_xmove(&g_device->buffer, irp, out, out_len, &info);
if (out_len < FORT_BUFFER_SIZE) {
status = STATUS_BUFFER_TOO_SMALL;
} else {
status = fort_buffer_xmove(&g_device->buffer, irp, out, out_len, &info);
if (status == STATUS_PENDING) {
KIRQL cirq;
if (status == STATUS_PENDING) {
KIRQL cirq;
IoMarkIrpPending(irp);
IoMarkIrpPending(irp);
IoAcquireCancelSpinLock(&cirq);
IoSetCancelRoutine(irp, fort_device_cancel_pending);
IoReleaseCancelSpinLock(cirq);
IoAcquireCancelSpinLock(&cirq);
IoSetCancelRoutine(irp, fort_device_cancel_pending);
IoReleaseCancelSpinLock(cirq);
return STATUS_PENDING;
return STATUS_PENDING;
}
}
break;
}

2
src/driver/fortstat.c
View File

@ -582,7 +582,7 @@ fort_stat_dpc_traf_flush (PFORT_STAT stat, UINT16 proc_count, PCHAR out)
/* Write process_id */
*out_proc = proc->process_id;
if (!proc->refcount) {
if (proc->refcount == 0) {
/* The process is inactive */
*out_proc |= 1;