Skip to content

Instantly share code, notes, and snippets.

@KruskalLin

KruskalLin/gist:bcb1465982ad32dfcf4124d6e1075ee0

Created December 8, 2024 09:30
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save KruskalLin/bcb1465982ad32dfcf4124d6e1075ee0 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save KruskalLin/bcb1465982ad32dfcf4124d6e1075ee0 to your computer and use it in GitHub Desktop.
Download ZIP

Revisions

  1. KruskalLin created this gist Dec 8, 2024.
    348 changes: 348 additions & 0 deletions gistfile1.txt
    View file
    Original file line number Diff line number Diff line change
    @@ -0,0 +1,348 @@
    #define _WINSOCK_DEPRECATED_NO_WARNINGS

    #include <WinSock2.h>
    #include <ws2tcpip.h>
    #include <iostream>
    #include <cstring>
    #include <cstdlib>
    #include <set>
    #include <string>
    #include <vector>
    #include <queue>
    #include <unordered_map>
    #include <chrono>
    #include <mutex>
    #include <thread>
    #pragma comment(lib, "ws2_32.lib")

    #define MAX_HOST_LEN 256
    #define MAX_URL_LEN 2048
    #define MAX_REQUEST_LEN 2048
    #define MAX_ROBOT_READ 16384
    #define MAX_PAGE_READ 2097152

    #define IP_HDR_SIZE 20 // RFC 791
    #define ICMP_HDR_SIZE 8 // RFC 792
    // max payload size of an ICMP message originated in the program
    #define MAX_ICMP_SIZE 65200
    // max size of an IP datagram
    // the returned ICMP message will most likely include only 8 bytes of the original message plus the IP header (as per RFC 792); however, longer replies (e.g., 68 bytes) are possible
    #define MAX_REPLY_SIZE (IP_HDR_SIZE + ICMP_HDR_SIZE + MAX_ICMP_SIZE)

    // ICMP packet types
    #define ICMP_ECHO_REPLY 0
    #define ICMP_DEST_UNREACH 3
    #define ICMP_TTL_EXPIRED 11
    #define ICMP_ECHO_REQUEST 8

    // remember the current packing state
    #pragma pack (push)
    #pragma pack (1)

    // define the IP header (20 bytes)
    struct IPHeader {
    u_char h_len : 4; /* lower 4 bits: length of the header in dwords */
    u_char version : 4; /* upper 4 bits: version of IP, i.e., 4 */
    u_char tos; /* type of service (TOS), ignore */
    u_short len; /* length of packet */
    u_short ident; /* unique identifier */
    u_short flags; /* flags together with fragment offset - 16 bits */
    u_char ttl; /* time to live */
    u_char proto; /* protocol number (6=TCP, 17=UDP, etc.) */
    u_short checksum; /* IP header checksum */
    u_long source_ip;
    u_long dest_ip;
    };

    // define the ICMP header (8 bytes)
    struct ICMPHeader {
    u_char type; /* ICMP packet type */
    u_char code; /* type subcode */
    u_short checksum; /* checksum of the ICMP */
    u_short id; /* application-specific ID */
    u_short seq; /* application-specific sequence */
    };

    // now restore the previous packing state
    #pragma pack (pop)

    struct Probe {
    int ttl;
    int seq;
    bool completed;
    bool responded;
    std::string router_ip;
    std::string router_name;
    double rtt; // in milliseconds
    int probe_num; // 1, 2, or 3
    std::chrono::steady_clock::time_point send_time;
    };

    struct TracerouteTarget {
    std::string target;
    std::string target_ip;
    bool resolved;
    struct sockaddr_in addr;
    std::vector<Probe> probes;
    int max_hops;
    bool completed;
    };


    std::unordered_map<std::string, std::string> dns_cache;
    std::mutex dns_mutex;


    u_short ip_checksum(u_short* buffer, int size)
    {
    u_long cksum = 0;

    /* sum all the words together, adding the final byte if size is odd */
    while (size > 1)
    {
    cksum += *buffer++;
    size -= sizeof(u_short);
    }
    if (size)
    {
    cksum += *(u_char*)buffer;

    }

    /* add carry bits to lower u_short word */
    cksum = (cksum >> 16) + (cksum & 0xffff);

    /* return a bitwise complement of the resulting mishmash */
    return (u_short) (~cksum);
    }

    void dns_lookup(const std::string& ip) {
    char hostname[NI_MAXHOST];
    sockaddr_in sa;
    sa.sin_family = AF_INET;
    inet_pton(AF_INET, ip.c_str(), &(sa.sin_addr));

    int res = getnameinfo((sockaddr*)&sa, sizeof(sa), hostname, sizeof(hostname), NULL, 0, 0);
    std::string name = (res == 0) ? std::string(hostname) : "";

    std::lock_guard<std::mutex> lock(dns_mutex);
    dns_cache[ip] = name;
    }


    void perform_traceroute(TracerouteTarget& target, SOCKET sock, int max_hops = 30, int probes_per_hop = 3) {
    target.max_hops = max_hops;
    target.completed = false;

    // Initialize probes
    for (int ttl = 1; ttl <= max_hops; ++ttl) {
    for (int probe_num = 1; probe_num <= probes_per_hop; ++probe_num) {
    Probe p;
    p.ttl = ttl;
    p.seq = ttl * probes_per_hop + probe_num;
    p.completed = false;
    p.responded = false;
    p.rtt = 0.0;
    p.probe_num = probe_num;
    target.probes.push_back(p);
    }
    }

    // Send all probes
    for (auto& probe : target.probes) {
    // Create ICMP Echo Request
    char send_buf[ICMP_HDR_SIZE] = { 0 };
    ICMPHeader* icmp = (ICMPHeader*)send_buf;
    icmp->type = ICMP_ECHO_REQUEST;
    icmp->code = 0;
    icmp->id = (u_short)GetCurrentProcessId();
    icmp->seq = (u_short)probe.seq;
    icmp->checksum = ip_checksum((u_short*)send_buf, sizeof(ICMPHeader));

    // Set TTL
    int ttl = probe.ttl;
    if (setsockopt(sock, IPPROTO_IP, IP_TTL, (const char*)&ttl, sizeof(ttl)) == SOCKET_ERROR) {
    printf("setsockopt failed with error: %d\n", WSAGetLastError());
    continue;
    }

    // Record send time
    probe.send_time = std::chrono::steady_clock::now();

    // Send packet
    int bytes_sent = sendto(sock, send_buf, sizeof(send_buf), 0, (struct sockaddr*)&target.addr, sizeof(target.addr));
    if (bytes_sent == SOCKET_ERROR) {
    printf("sendto failed with error: %d\n", WSAGetLastError());
    continue;
    }
    }

    // Receive responses
    int received_probes = 0;
    while (received_probes < target.probes.size()) {
    char recv_buf[MAX_REPLY_SIZE];
    sockaddr_in sender;
    int sender_len = sizeof(sender);

    // Set timeout for select
    timeval timeout;
    timeout.tv_sec = 2;
    timeout.tv_usec = 0;

    fd_set read_fds;
    FD_ZERO(&read_fds);
    FD_SET(sock, &read_fds);

    int ret = select(0, &read_fds, NULL, NULL, &timeout);
    if (ret > 0 && FD_ISSET(sock, &read_fds)) {
    int bytes_received = recvfrom(sock, recv_buf, sizeof(recv_buf), 0, (sockaddr*)&sender, &sender_len);
    if (bytes_received >= (int)(IP_HDR_SIZE + ICMP_HDR_SIZE)) {
    // Parse IP Header
    IPHeader* ip_hdr = (IPHeader*)recv_buf;
    // Parse ICMP Header
    ICMPHeader* icmp_hdr = (ICMPHeader*)(recv_buf + IP_HDR_SIZE);

    // Check if it's a TTL Expired message
    if (icmp_hdr->type == ICMP_TTL_EXPIRED || icmp_hdr->type == ICMP_ECHO_REPLY) {
    // Extract original ICMP packet
    ICMPHeader* orig_icmp = (ICMPHeader*)(recv_buf + IP_HDR_SIZE + ICMP_HDR_SIZE + IP_HDR_SIZE);
    if (orig_icmp->id != (u_short)GetCurrentProcessId()) {
    // Not our packet
    continue;
    }

    // Find corresponding probe
    int probe_seq = orig_icmp->seq;
    auto it = std::find_if(target.probes.begin(), target.probes.end(),
    [&](const Probe& p) { return p.seq == probe_seq; });

    if (it != target.probes.end()) {
    // Calculate RTT
    auto now = std::chrono::steady_clock::now();
    double rtt = std::chrono::duration_cast<std::chrono::milliseconds>(now - it->send_time).count();
    it->rtt = rtt;
    it->responded = true;
    it->completed = true;
    received_probes++;

    // Get router IP
    char sender_ip[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &(sender.sin_addr), sender_ip, INET_ADDRSTRLEN);
    it->router_ip = std::string(sender_ip);

    // Perform DNS lookup asynchronously
    std::thread(dns_lookup, it->router_ip).detach();

    // Print probe result
    std::string hostname;
    {
    std::lock_guard<std::mutex> lock(dns_mutex);
    if (dns_cache.find(it->router_ip) != dns_cache.end()) {
    hostname = dns_cache[it->router_ip];
    }
    }
    if (hostname.empty()) hostname = it->router_ip;
    std::cout << it->ttl << " " << hostname << " (" << it->router_ip << ") " << it->rtt << " ms" << std::endl;

    // If Echo Reply, traceroute is complete
    if (icmp_hdr->type == ICMP_ECHO_REPLY) {
    target.completed = true;
    break;
    }
    }
    }
    }
    }
    else {
    // Timeout, mark probes as lost
    for (auto& probe : target.probes) {
    if (!probe.completed) {
    probe.completed = true;
    std::cout << probe.ttl << " *" << std::endl;
    received_probes++;
    }
    }
    break;
    }
    }

    target.completed = true;
    }


    int processURL(const char* url) {
    clock_t t = clock();

    WSADATA wsaData;

    //Initialize WinSock; once per program run
    WORD wVersionRequested = MAKEWORD(2, 2);
    if (WSAStartup(wVersionRequested, &wsaData) != 0) {
    printf("failed with %d\n", WSAGetLastError());
    WSACleanup();
    return 1;
    }

    // structure used in DNS lookups
    struct hostent* remote;

    // structure for connecting to server
    struct sockaddr_in server;
    memset(&server, 0, sizeof(sockaddr_in));
    server.sin_family = AF_INET;
    // first assume that the string is an IP address
    DWORD IP = inet_addr(url);
    if (IP == INADDR_NONE)
    {
    // if not a valid IP, then do a DNS lookup
    if ((remote = gethostbyname(url)) == NULL)
    {
    printf("failed with %d\n", WSAGetLastError());
    return 1;
    }
    else // take the first IP address and copy into sin_addr
    memcpy((char*)&(server.sin_addr), remote->h_addr, remote->h_length);
    }
    else
    {
    // if a valid IP, directly drop its binary version into sin_addr
    server.sin_addr.S_un.S_addr = IP;
    }
    printf("Tracerouting to %s...\n", inet_ntoa(server.sin_addr));


    SOCKET sock = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP);
    if (sock == INVALID_SOCKET) {
    printf("Unable to create an ICMP socket: error %d\n", WSAGetLastError());
    WSACleanup();
    return 1;
    }

    TracerouteTarget target;
    target.target = std::string(url);
    target.target_ip = std::string(inet_ntoa(server.sin_addr));
    target.resolved = true;
    target.addr = server;
    target.completed = false;

    perform_traceroute(target, sock);

    closesocket(sock);
    WSACleanup();
    printf("\nTotal execution time: %.0f", (double)(clock() - t) / CLOCKS_PER_SEC * 1000);
    return 0;
    }


    int main(int argc, char* argv[])
    {

    if (argc == 2) {
    const char* url = argv[1];
    return processURL(url);
    }
    else {
    printf("Usage: hw4.exe [url]\n");
    return 1;
    }
    }