diff --git a/src/driver/fortbuf.c b/src/driver/fortbuf.c index db271346..c26eb6ce 100644 --- a/src/driver/fortbuf.c +++ b/src/driver/fortbuf.c @@ -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; diff --git a/src/driver/fortdrv.c b/src/driver/fortdrv.c index 6866217e..9661d8de 100644 --- a/src/driver/fortdrv.c +++ b/src/driver/fortdrv.c @@ -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; } diff --git a/src/driver/fortstat.c b/src/driver/fortstat.c index d13a2aeb..c8f48d48 100644 --- a/src/driver/fortstat.c +++ b/src/driver/fortstat.c @@ -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;