poverlapped error

[571451370]

#include <winsock2.h>
#include <windows.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "WinDivert.h"

#define MAX_THREADS 64
#define MAX_BATCH_SIZE WINDIVERT_BATCH_MAX
#define DEFAULT_BUFFER_SIZE (0x10000 + 4096) // Default buffer size

typedef struct {
HANDLE handle;
HANDLE event;
HANDLE ioport;
UINT8* packet;
UINT packet_len;
WINDIVERT_ADDRESS* addr;
UINT addr_len;
} AsyncPacket;

/*

• Options.error
  */
  static int threads = 1;
  static int batch = WINDIVERT_BATCH_MAX;
  static int priority = 0;
  static WINDIVERT_LAYER layer = WINDIVERT_LAYER_NETWORK;
  static int size = DEFAULT_BUFFER_SIZE;

/*

• Print usage and exit.
  /
  static void usage(const char prog) {
  fprintf(stderr, "usage: %s [OPTIONS] filter-string\n\n", prog);
  fprintf(stderr, "OPTIONS:\n");
  fprintf(stderr, "\t--batch N\n");
  fprintf(stderr, "\t\tSet the batch size to N (default=%u)\n", WINDIVERT_BATCH_MAX);
  fprintf(stderr, "\t--layer LAYER\n");
  fprintf(stderr, "\t\tSet the filter layer to LAYER (default=network).\n");
  fprintf(stderr, "\t\tValid values are {ethernet,network,forward}.\n");
  fprintf(stderr, "\t--priority N\n");
  fprintf(stderr, "\t\tSet the filter priority to N (default=0)\n");
  fprintf(stderr, "\t--size N\n");
  fprintf(stderr, "\t\tSet the buffer size to N (default=%u)\n", DEFAULT_BUFFER_SIZE);
  fprintf(stderr, "\t--threads N\n");
  fprintf(stderr, "\t\tSet the number of threads to be N (default=1)\n");
  exit(EXIT_FAILURE);
  }

/*

• Parse options.
  /
  static const char parse_options(int argc, char** argv) {
  int i;
  size_t n;
  const char* filter = NULL, * opt, * arg;

  for (i = 1; i < argc; i++) {
  opt = argv[i];
  if (opt[0] != '-' || opt[1] != '-') {
  if (filter != NULL) {
  usage(argv[0]);
  }
  filter = opt;
  continue;
  }
  opt += 2;
  arg = strchr(opt, '=');
  if (arg == NULL) {
  i++;
  if (i >= argc) {
  usage(argv[0]);
  }
  arg = argv[i];
  n = strlen(opt);
  }
  else {
  n = arg - opt;
  arg++;
  }
  if (strncmp(opt, "threads", n) == 0) {
  threads = atoi(arg);
  }
  else if (strncmp(opt, "batch", n) == 0) {
  batch = atoi(arg);
  }
  else if (strncmp(opt, "priority", n) == 0) {
  priority = atoi(arg);
  }
  else if (strncmp(opt, "size", n) == 0) {
  size = atoi(arg);
  }
  else if (strncmp(opt, "layer", n) == 0) {
  if (strcmp(arg, "ethernet") == 0) {
  layer = WINDIVERT_LAYER_ETHERNET;
  }
  else if (strcmp(arg, "network") == 0) {
  layer = WINDIVERT_LAYER_NETWORK;
  }
  else if (strcmp(arg, "forward") == 0) {
  layer = WINDIVERT_LAYER_NETWORK_FORWARD;
  }
  else {
  usage(argv[0]);
  }
  }
  else {
  usage(argv[0]);
  }
  }
  return (filter == NULL ? "true" : filter);
  }

/*

• Cleanup completed I/O requests.
  /
  static void cleanup(HANDLE ioport, OVERLAPPED ignore) {
  OVERLAPPED* overlapped;
  DWORD iolen;
  ULONG_PTR iokey = 0;

  while (GetQueuedCompletionStatus(ioport, &iolen, &iokey, &overlapped, 0)) {
  if (overlapped != ignore) {
  free(overlapped);
  }
  }
  }

// Forward declaration of worker_thread
static DWORD WINAPI worker_thread(LPVOID arg);

/*

• Entry.
  /
  int __cdecl main(int argc, char* argv) {
  const char* filter;
  int i;
  HANDLE handle;
  HANDLE threads_arr[MAX_THREADS];
  AsyncPacket* packets[MAX_THREADS];

  filter = parse_options(argc, argv);

  if (threads < 1 || threads > MAX_THREADS) {
  fprintf(stderr, "error: number of threads must be within the range 1..%d, found %d\n", MAX_THREADS, threads);
  exit(EXIT_FAILURE);
  }
  if (batch < 1 || batch > WINDIVERT_BATCH_MAX) {
  fprintf(stderr, "error: batch size must be within the range 1..%u, found %d\n", WINDIVERT_BATCH_MAX, batch);
  exit(EXIT_FAILURE);
  }
  if (priority < WINDIVERT_PRIORITY_LOWEST || priority > WINDIVERT_PRIORITY_HIGHEST) {
  fprintf(stderr, "error: priority must be within the range %d..%d, found %d\n", WINDIVERT_PRIORITY_LOWEST, WINDIVERT_PRIORITY_HIGHEST, priority);
  exit(EXIT_FAILURE);
  }
  if (size < 1 || size >= MAX_BATCH_SIZE * WINDIVERT_MTU_MAX) {
  fprintf(stderr, "error: buffer size must be within the range 1..%d, found %d\n", MAX_BATCH_SIZE * WINDIVERT_MTU_MAX, size);
  exit(EXIT_FAILURE);
  }

  // Divert traffic matching the filter:
  handle = WinDivertOpen(filter, layer, (INT16)priority, 0);
  if (handle == INVALID_HANDLE_VALUE) {
  if (GetLastError() == ERROR_INVALID_PARAMETER) {
  fprintf(stderr, "error: filter syntax error\n");
  exit(EXIT_FAILURE);
  }
  fprintf(stderr, "error: failed to open the WinDivert device (%d)\n", GetLastError());
  exit(EXIT_FAILURE);
  }

  // Create IOCP
  HANDLE iocp = CreateIoCompletionPort(INVALID_HANDLE_VALUE, NULL, 0, 0);
  if (iocp == NULL) {
  fprintf(stderr, "error: failed to create IOCP (%d)\n", GetLastError());
  WinDivertClose(handle);
  exit(EXIT_FAILURE);
  }

  // Associate handle with IOCP
  if (CreateIoCompletionPort(handle, iocp, 0, 0) == NULL) {
  fprintf(stderr, "error: failed to associate handle with IOCP (%d)\n", GetLastError());
  CloseHandle(iocp);
  WinDivertClose(handle);
  exit(EXIT_FAILURE);
  }

  // Allocate packets
  for (i = 0; i < threads; i++) {
  packets[i] = (AsyncPacket*)malloc(sizeof(AsyncPacket));
  if (packets[i] == NULL) {
  fprintf(stderr, "error: failed to allocate packet (%d)\n", GetLastError());
  for (int j = 0; j < i; j++) {
  CloseHandle(packets[j]->event);
  free(packets[j]->packet);
  free(packets[j]->addr);
  free(packets[j]);
  }
  CloseHandle(iocp);
  WinDivertClose(handle);
  exit(EXIT_FAILURE);
  }
  packets[i]->handle = handle; // Initialize handle
  packets[i]->ioport = iocp;
  packets[i]->packet = (UINT8*)malloc(size);
  packets[i]->addr = (WINDIVERT_ADDRESS*)malloc(batch * sizeof(WINDIVERT_ADDRESS));
  if (packets[i]->packet == NULL || packets[i]->addr == NULL) {
  fprintf(stderr, "error: failed to allocate buffer (%d)\n", GetLastError());
  CloseHandle(packets[i]->event);
  free(packets[i]->packet);
  free(packets[i]->addr);
  free(packets[i]);
  for (int j = 0; j < i; j++) {
  CloseHandle(packets[j]->event);
  free(packets[j]->packet);
  free(packets[j]->addr);
  free(packets[j]);
  }
  CloseHandle(iocp);
  WinDivertClose(handle);
  exit(EXIT_FAILURE);
  }
  packets[i]->event = CreateEvent(NULL, FALSE, FALSE, NULL);
  if (packets[i]->event == NULL) {
  fprintf(stderr, "error: failed to create event (%d)\n", GetLastError());
  free(packets[i]->packet);
  free(packets[i]->addr);
  free(packets[i]);
  for (int j = 0; j < i; j++) {
  CloseHandle(packets[j]->event);
  free(packets[j]->packet);
  free(packets[j]->addr);
  free(packets[j]);
  }
  CloseHandle(iocp);
  WinDivertClose(handle);
  exit(EXIT_FAILURE);
  }

  }

  // Start worker threads
  for (i = 0; i < threads; i++) {
  threads_arr[i] = CreateThread(NULL, 0, worker_thread, packets[i], 0, NULL);
  if (threads_arr[i] == NULL) {
  fprintf(stderr, "error: failed to start worker thread (%d)\n", GetLastError());
  for (int j = 0; j < i; j++) {
  CloseHandle(threads_arr[j]);
  }
  for (int j = 0; j < threads; j++) {
  CloseHandle(packets[j]->event);
  free(packets[j]->packet);
  free(packets[j]->addr);
  free(packets[j]);
  }
  CloseHandle(iocp);
  WinDivertClose(handle);
  exit(EXIT_FAILURE);
  }
  }

  // Wait for threads to complete
  WaitForMultipleObjects(threads, threads_arr, TRUE, INFINITE);

  // Cleanup
  for (i = 0; i < threads; i++) {
  CloseHandle(threads_arr[i]);
  CloseHandle(packets[i]->event);
  free(packets[i]->packet);
  free(packets[i]->addr);
  free(packets[i]);
  }
  CloseHandle(iocp);
  WinDivertClose(handle);

  printf("Done.\n");
  return 0;
  }

// Passthru thread.
static DWORD WINAPI worker_thread(LPVOID arg) {
AsyncPacket* packet = (AsyncPacket*)arg;
HANDLE handle = packet->handle;
HANDLE event = packet->event;
OVERLAPPED overlapped;
DWORD len;
DWORD bytes_transferred;
ULONG_PTR completion_key;
OVERLAPPED* poverlapped;

while (TRUE) {
    memset(&overlapped, 0, sizeof(overlapped));
    ResetEvent(event);
    overlapped.hEvent = event;
    packet->addr_len = batch * sizeof(WINDIVERT_ADDRESS);

    // Post read operation
    if (!WinDivertRecvEx(handle, packet->packet, size, &packet->packet_len, 0, packet->addr, &packet->addr_len, &overlapped)) {
        if (GetLastError() != ERROR_IO_PENDING) {
        read_failed:
            fprintf(stderr, "failed to read packet (%d)\n", GetLastError());
            continue;
        }

        // Timeout = 1s
        while (WaitForSingleObject(event, 1000) == WAIT_TIMEOUT) {
            cleanup(packet->ioport, &overlapped);
        }

        if (!GetOverlappedResult(handle, &overlapped, &len, FALSE)) {
            goto read_failed;
        }

        packet->addr_len = len;
    }
    cleanup(packet->ioport, &overlapped);

    // Log packet size
    fprintf(stderr, "Received packet of size: %u\n", packet->packet_len);
    WinDivertHelperCalcChecksums(packet->packet, packet->packet_len, packet->addr->Layer, packet->addr, 0);

    // Post send operation
    poverlapped = (OVERLAPPED*)malloc(sizeof(OVERLAPPED));
    if (poverlapped == NULL) {
        fprintf(stderr, "failed to allocate memory\n");
        continue;
    }
    memset(poverlapped, 0, sizeof(OVERLAPPED));
    if (WinDivertSendEx(handle, packet->packet, packet->packet_len, NULL, 0, packet->addr, packet->addr_len, poverlapped)) {
        free(poverlapped);
        continue;
    }

    if (GetLastError() != ERROR_IO_PENDING) {
        fprintf(stderr, "warning: failed to post send operation (%d)\n", GetLastError());
        free(poverlapped);
        continue;
    }
}

return 0;

}