Driver: fortpkt: Simplify injection handles

This commit is contained in:
Nodir Temirkhodjaev 2023-04-06 13:00:01 +03:00
parent 6b13e5df02
commit f7415e60d2
2 changed files with 21 additions and 32 deletions

View File

@ -335,11 +335,9 @@ static NTSTATUS fort_packet_inject_out(const PFORT_PACKET_IO pkt, PNET_BUFFER_LI
(FWPS_INJECT_COMPLETE0) &fort_packet_inject_complete, pkt); (FWPS_INJECT_COMPLETE0) &fort_packet_inject_complete, pkt);
} }
inline static HANDLE fort_shaper_injection_id(PFORT_SHAPER shaper, BOOL isIPv6, BOOL inbound) inline static HANDLE fort_shaper_injection_id(PFORT_SHAPER shaper, BOOL isIPv6)
{ {
return inbound return isIPv6 ? shaper->injection_transport6_id : shaper->injection_transport4_id;
? (isIPv6 ? shaper->injection_in_transport6_id : shaper->injection_in_transport4_id)
: (isIPv6 ? shaper->injection_out_transport6_id : shaper->injection_out_transport4_id);
} }
inline static NTSTATUS fort_packet_clone( inline static NTSTATUS fort_packet_clone(
@ -372,19 +370,19 @@ static NTSTATUS fort_shaper_packet_base_inject(
NTSTATUS status; NTSTATUS status;
const BOOL inbound = (pkt->flags & FORT_PACKET_INBOUND) != 0; const BOOL inbound = (pkt->flags & FORT_PACKET_INBOUND) != 0;
const BOOL isIPv6 = (pkt->flags & FORT_PACKET_IP6) != 0;
const HANDLE injection_id = fort_shaper_injection_id(shaper, isIPv6, inbound);
const ADDRESS_FAMILY addressFamily = (isIPv6 ? AF_INET6 : AF_INET);
status = fort_packet_clone(pkt, clonedNetBufList, inbound); status = fort_packet_clone(pkt, clonedNetBufList, inbound);
if (!NT_SUCCESS(status)) { if (!NT_SUCCESS(status)) {
LOG("Shaper: Packet clone error: %x\n", status); LOG("Shaper: Packet clone error: %x\n", status);
TRACE(FORT_SHAPER_PACKET_CLONE_ERROR, status, 0, 0); TRACE(FORT_SHAPER_PACKET_CLONE_ERROR, status, 0, 0);
return status;
} }
if (NT_SUCCESS(status)) { const BOOL isIPv6 = (pkt->flags & FORT_PACKET_IP6) != 0;
status = inbound const ADDRESS_FAMILY addressFamily = (isIPv6 ? AF_INET6 : AF_INET);
? fort_packet_inject_in(pkt, *clonedNetBufList, injection_id, addressFamily) const HANDLE injection_id = fort_shaper_injection_id(shaper, isIPv6);
status = inbound ? fort_packet_inject_in(pkt, *clonedNetBufList, injection_id, addressFamily)
: fort_packet_inject_out(pkt, *clonedNetBufList, injection_id, addressFamily); : fort_packet_inject_out(pkt, *clonedNetBufList, injection_id, addressFamily);
if (!NT_SUCCESS(status)) { if (!NT_SUCCESS(status)) {
LOG("Shaper: Packet injection call error: %x\n", status); LOG("Shaper: Packet injection call error: %x\n", status);
@ -392,7 +390,6 @@ static NTSTATUS fort_shaper_packet_base_inject(
(*clonedNetBufList)->Status = STATUS_SUCCESS; (*clonedNetBufList)->Status = STATUS_SUCCESS;
} }
}
return status; return status;
} }
@ -814,13 +811,9 @@ FORT_API void fort_shaper_open(PFORT_SHAPER shaper)
g_RandomSeed = now.LowPart; g_RandomSeed = now.LowPart;
FwpsInjectionHandleCreate0( FwpsInjectionHandleCreate0(
AF_INET, FWPS_INJECTION_TYPE_TRANSPORT, &shaper->injection_in_transport4_id); AF_INET, FWPS_INJECTION_TYPE_TRANSPORT, &shaper->injection_transport4_id);
FwpsInjectionHandleCreate0( FwpsInjectionHandleCreate0(
AF_INET6, FWPS_INJECTION_TYPE_TRANSPORT, &shaper->injection_in_transport6_id); AF_INET6, FWPS_INJECTION_TYPE_TRANSPORT, &shaper->injection_transport6_id);
FwpsInjectionHandleCreate0(
AF_INET, FWPS_INJECTION_TYPE_TRANSPORT, &shaper->injection_out_transport4_id);
FwpsInjectionHandleCreate0(
AF_INET6, FWPS_INJECTION_TYPE_TRANSPORT, &shaper->injection_out_transport6_id);
tommy_arrayof_init(&shaper->pending.procs, sizeof(FORT_PENDING_PROC)); tommy_arrayof_init(&shaper->pending.procs, sizeof(FORT_PENDING_PROC));
tommy_arrayof_init(&shaper->pending.packets, sizeof(FORT_PENDING_PACKET)); tommy_arrayof_init(&shaper->pending.packets, sizeof(FORT_PENDING_PACKET));
@ -844,10 +837,8 @@ FORT_API void fort_shaper_close(PFORT_SHAPER shaper)
tommy_arrayof_done(&shaper->packets); tommy_arrayof_done(&shaper->packets);
FwpsInjectionHandleDestroy0(shaper->injection_in_transport4_id); FwpsInjectionHandleDestroy0(shaper->injection_transport4_id);
FwpsInjectionHandleDestroy0(shaper->injection_in_transport6_id); FwpsInjectionHandleDestroy0(shaper->injection_transport6_id);
FwpsInjectionHandleDestroy0(shaper->injection_out_transport4_id);
FwpsInjectionHandleDestroy0(shaper->injection_out_transport6_id);
} }
FORT_API void fort_shaper_conf_update(PFORT_SHAPER shaper, const PFORT_CONF_IO conf_io) FORT_API void fort_shaper_conf_update(PFORT_SHAPER shaper, const PFORT_CONF_IO conf_io)
@ -1001,7 +992,7 @@ inline static NTSTATUS fort_shaper_packet_queue(
inline static BOOL fort_shaper_injected_by_self(PFORT_SHAPER shaper, PCFORT_CALLOUT_ARG ca) inline static BOOL fort_shaper_injected_by_self(PFORT_SHAPER shaper, PCFORT_CALLOUT_ARG ca)
{ {
const HANDLE injection_id = fort_shaper_injection_id(shaper, ca->isIPv6, ca->inbound); const HANDLE injection_id = fort_shaper_injection_id(shaper, ca->isIPv6);
return fort_packet_injected_by_self(injection_id, ca->netBufList); return fort_packet_injected_by_self(injection_id, ca->netBufList);
} }

View File

@ -31,7 +31,7 @@ typedef struct fort_packet_out
#define FORT_PACKET_INBOUND 0x01 #define FORT_PACKET_INBOUND 0x01
#define FORT_PACKET_IP6 0x02 #define FORT_PACKET_IP6 0x02
#define FORT_PACKET_IPSEC_PROTECTED 0x04 #define FORT_PACKET_IPSEC_PROTECTED 0x08
typedef struct fort_packet_io typedef struct fort_packet_io
{ {
@ -127,10 +127,8 @@ typedef struct fort_shaper
LONG volatile group_io_bits; LONG volatile group_io_bits;
LONG volatile active_io_bits; LONG volatile active_io_bits;
HANDLE injection_in_transport4_id; HANDLE injection_transport4_id;
HANDLE injection_in_transport6_id; HANDLE injection_transport6_id;
HANDLE injection_out_transport4_id;
HANDLE injection_out_transport6_id;
FORT_PENDING pending; FORT_PENDING pending;