/* * File: nettool_dual.c
 * Author: LULU
 * Description: Standalone TCP/UDP Latency Tester with DNS support.
 * Compile: gcc -static nettool_dual.c -o nettool
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <sys/select.h>
#include <netdb.h>
#include <errno.h>
#include <signal.h>

#define BUF_SIZE 1024

// --- 辅助函数：获取当前时间戳 (毫秒) ---
double get_timestamp() {
    struct timeval tv;
    gettimeofday(&tv, NULL);
    return (double)tv.tv_sec * 1000.0 + (double)tv.tv_usec / 1000.0;
}

// --- 辅助函数：DNS 解析 ---
int resolve_hostname(const char *hostname, char *ip_str) {
    struct hostent *he;
    struct in_addr **addr_list;

    // 如果本身就是 IP 字符串，直接复制
    if (inet_addr(hostname) != INADDR_NONE) {
        strcpy(ip_str, hostname);
        return 0;
    }

    if ((he = gethostbyname(hostname)) == NULL) {
        return -1;
    }

    addr_list = (struct in_addr **)he->h_addr_list;
    if (addr_list[0] != NULL) {
        strcpy(ip_str, inet_ntoa(*addr_list[0]));
        return 0;
    }
    return -1;
}

// --- 服务端逻辑 ---
void run_server(int port) {
    int tcp_fd, udp_fd, max_fd;
    struct sockaddr_in server_addr, client_addr;
    fd_set readfds;
    char buffer[BUF_SIZE];
    socklen_t addr_len = sizeof(client_addr);

    // 1. 创建 TCP Socket
    if ((tcp_fd = socket(AF_INET, SOCK_STREAM, 0)) < 0) { perror("TCP Socket"); exit(1); }
    int opt = 1;
    setsockopt(tcp_fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));
    
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_addr.s_addr = INADDR_ANY;
    server_addr.sin_port = htons(port);

    if (bind(tcp_fd, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        perror("TCP Bind"); exit(1);
    }
    listen(tcp_fd, 5);

    // 2. 创建 UDP Socket
    if ((udp_fd = socket(AF_INET, SOCK_DGRAM, 0)) < 0) { perror("UDP Socket"); exit(1); }
    if (bind(udp_fd, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        perror("UDP Bind"); exit(1);
    }

    printf("========================================\n");
    printf("[Server] Listening on Port %d (TCP+UDP)\n", port);
    printf("========================================\n");

    while (1) {
        FD_ZERO(&readfds);
        FD_SET(tcp_fd, &readfds);
        FD_SET(udp_fd, &readfds);
        max_fd = (tcp_fd > udp_fd) ? tcp_fd : udp_fd;

        // 等待事件
        if (select(max_fd + 1, &readfds, NULL, NULL, NULL) < 0) {
            if (errno == EINTR) continue;
            perror("Select"); break;
        }

        // 处理 UDP 数据
        if (FD_ISSET(udp_fd, &readfds)) {
            int n = recvfrom(udp_fd, buffer, BUF_SIZE, 0, (struct sockaddr *)&client_addr, &addr_len);
            if (n > 0) {
                sendto(udp_fd, buffer, n, 0, (struct sockaddr *)&client_addr, addr_len);
                printf("[UDP] Echo %d bytes -> %s\n", n, inet_ntoa(client_addr.sin_addr));
            }
        }

        // 处理 TCP 连接
        if (FD_ISSET(tcp_fd, &readfds)) {
            int new_sock = accept(tcp_fd, (struct sockaddr *)&client_addr, &addr_len);
            if (new_sock >= 0) {
                printf("[TCP] Connect -> %s\n", inet_ntoa(client_addr.sin_addr));
                // 简单的 Fork 处理，或者单线程阻塞处理（为了嵌入式简单，这里用简单的循环处理直到断开）
                int n;
                while ((n = recv(new_sock, buffer, BUF_SIZE, 0)) > 0) {
                    send(new_sock, buffer, n, 0); // Echo
                }
                close(new_sock);
                printf("[TCP] Disconnect\n");
            }
        }
    }
    close(tcp_fd);
    close(udp_fd);
}

// --- 客户端逻辑 ---
void run_client(const char *target, int port, int count) {
    char ip[100];
    
    printf("Resolving %s ... ", target);
    if (resolve_hostname(target, ip) != 0) {
        printf("FAILED\n[Error] Could not resolve hostname.\n");
        return;
    }
    printf("%s\n", ip);

    struct sockaddr_in servaddr;
    memset(&servaddr, 0, sizeof(servaddr));
    servaddr.sin_family = AF_INET;
    servaddr.sin_port = htons(port);
    servaddr.sin_addr.s_addr = inet_addr(ip);

    struct timeval tv;
    tv.tv_sec = 2; tv.tv_usec = 0; // 2秒超时

    // --- 1. TCP 测试 ---
    printf("\n>>> Testing TCP (Handshake & RTT) <<<\n");
    int tcp_sock = socket(AF_INET, SOCK_STREAM, 0);
    if (tcp_sock < 0) { perror("TCP Socket"); return; }
    setsockopt(tcp_sock, SOL_SOCKET, SO_RCVTIMEO, (const char*)&tv, sizeof tv);

    double t_start = get_timestamp();
    if (connect(tcp_sock, (struct sockaddr *)&servaddr, sizeof(servaddr)) < 0) {
        printf("TCP Connect Failed: %s\n", strerror(errno));
    } else {
        double t_connected = get_timestamp();
        printf("TCP Handshake: %.2f ms\n", t_connected - t_start);

        char buf[BUF_SIZE];
        char *msg = "ping";
        double total_rtt = 0;
        int ok_count = 0;

        for (int i = 1; i <= count; i++) {
            double t1 = get_timestamp();
            if (send(tcp_sock, msg, strlen(msg), 0) < 0) break;
            if (recv(tcp_sock, buf, BUF_SIZE, 0) <= 0) break;
            double t2 = get_timestamp();
            double rtt = t2 - t1;
            printf("Seq=%d RTT=%.2f ms\n", i, rtt);
            total_rtt += rtt;
            ok_count++;
            usleep(200000); // 200ms
        }
        if (ok_count > 0) printf("TCP Avg RTT: %.2f ms\n", total_rtt / ok_count);
        close(tcp_sock);
    }

    // --- 2. UDP 测试 ---
    printf("\n>>> Testing UDP (Packet RTT) <<<\n");
    int udp_sock = socket(AF_INET, SOCK_DGRAM, 0);
    setsockopt(udp_sock, SOL_SOCKET, SO_RCVTIMEO, (const char*)&tv, sizeof tv);

    double total_udp = 0;
    int ok_udp = 0;
    char udp_buf[BUF_SIZE];

    for (int i = 1; i <= count; i++) {
        char msg[32]; snprintf(msg, 32, "udp_seq_%d", i);
        double t1 = get_timestamp();
        
        sendto(udp_sock, msg, strlen(msg), 0, (struct sockaddr *)&servaddr, sizeof(servaddr));
        
        socklen_t len = sizeof(servaddr); // 必须初始化
        if (recvfrom(udp_sock, udp_buf, BUF_SIZE, 0, NULL, NULL) > 0) {
            double t2 = get_timestamp();
            double rtt = t2 - t1;
            printf("Seq=%d RTT=%.2f ms\n", i, rtt);
            total_udp += rtt;
            ok_udp++;
        } else {
            printf("Seq=%d Timeout\n", i);
        }
        usleep(200000);
    }
    
    if (ok_udp > 0) {
        printf("UDP Avg RTT: %.2f ms (Loss: %.0f%%)\n", 
               total_udp / ok_udp, 
               100.0 * (count - ok_udp) / count);
    } else {
        printf("UDP Loss: 100%%\n");
    }
    close(udp_sock);
}

int main(int argc, char *argv[]) {
    // 忽略 broken pipe 信号，防止 TCP 断开导致进程退出
    signal(SIGPIPE, SIG_IGN);

    if (argc < 2) {
        printf("Nettool (Dual Protocol) - Usage:\n");
        printf("  Server: %s -s <port>\n", argv[0]);
        printf("  Client: %s -c <host> <port>\n", argv[0]);
        return 1;
    }

    if (strcmp(argv[1], "-s") == 0 && argc == 3) {
        run_server(atoi(argv[2]));
    } else if (strcmp(argv[1], "-c") == 0 && argc == 4) {
        run_client(argv[2], atoi(argv[3]), 5);
    } else {
        printf("Invalid arguments.\n");
    }
    return 0;
}