ithewei 6 years ago
parent
commit
351ab90ba1
1 changed files with 57 additions and 70 deletions
  1. 57 70
      event/nmap.cpp

+ 57 - 70
event/nmap.cpp

@@ -4,59 +4,23 @@
 #include "hstring.h"
 
 #define MAX_RECVFROM_TIMEOUT    5000 // ms
+#define MAX_SENDTO_PERSOCKET    1024
 
 typedef struct nmap_udata_s {
     Nmap*   nmap;
-    hio_t*  io;
     int     send_cnt;
     int     recv_cnt;
     int     up_cnt;
     int     idle_cnt;
 } nmap_udata_t;
 
-static int sendto_icmp(hio_t* io, Nmap* nmap) {
-    // icmp
-    char sendbuf[44]; // 20IP + 44ICMP = 64
-    icmp_t* icmp_req = (icmp_t*)sendbuf;
-    icmp_req->icmp_type = ICMP_ECHO;
-    icmp_req->icmp_code = 0;
-    icmp_req->icmp_id = getpid();
-    for (int i = 0; i < sizeof(sendbuf) - sizeof(icmphdr_t); ++i) {
-        icmp_req->icmp_data[i] = i;
-    }
-    struct sockaddr_in peeraddr;
-    int send_cnt = 0;
-    for (auto iter = nmap->begin(); iter != nmap->end(); ++iter) {
-        if (iter->second == 1) continue;
-        icmp_req->icmp_seq = iter->first;
-        icmp_req->icmp_cksum = 0;
-        icmp_req->icmp_cksum = checksum((uint8_t*)icmp_req, sizeof(sendbuf));
-        socklen_t addrlen = sizeof(peeraddr);
-        memset(&peeraddr, 0, addrlen);
-        peeraddr.sin_family = AF_INET;
-        peeraddr.sin_addr.s_addr = iter->first;
-        hio_setpeeraddr(io, (struct sockaddr*)&peeraddr, addrlen);
-        hsendto(io->loop, io->fd, sendbuf, sizeof(sendbuf), NULL);
-        ++send_cnt;
-        // NOTE: avoid SNDBUF full.
-        if (send_cnt % 1000 == 0) {
-            msleep(10);
-        }
-    }
-    return send_cnt;
-}
-
-
 static void on_idle(hidle_t* idle) {
-    nmap_udata_t* udata = (nmap_udata_t*)idle->userdata;
+    nmap_udata_t* udata = (nmap_udata_t*)idle->loop->userdata;
     udata->idle_cnt++;
     if (udata->idle_cnt == 1) {
-        // try again
-        sendto_icmp(udata->io, udata->nmap);
-    }
-    else {
-        hloop_stop(idle->loop);
+        // try again?
     }
+    hloop_stop(idle->loop);
 }
 
 static void on_timer(htimer_t* timer) {
@@ -70,7 +34,7 @@ static void on_recvfrom(hio_t* io, void* buf, int readbytes) {
     //printf("[%s] <=> [%s]\n",
             //sockaddr_snprintf(io->localaddr, localaddrstr, sizeof(localaddrstr)),
             //sockaddr_snprintf(io->peeraddr, peeraddrstr, sizeof(peeraddrstr)));
-    nmap_udata_t* udata = (nmap_udata_t*)io->userdata;
+    nmap_udata_t* udata = (nmap_udata_t*)io->loop->userdata;
     if (++udata->recv_cnt == udata->send_cnt) {
         //hloop_stop(io->loop);
     }
@@ -88,48 +52,71 @@ static void on_recvfrom(hio_t* io, void* buf, int readbytes) {
 }
 
 int nmap_discovery(Nmap* nmap) {
-    // socket
-    int sockfd = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP);
-    if (sockfd < 0) {
-        perror("socket");
-        if (errno == EPERM) {
-            fprintf(stderr, "please use root or sudo to create a raw socket.\n");
-        }
-        return -socket_errno();
-    }
-    nonblocking(sockfd);
-    int len = 425984; // 416K
-    socklen_t optlen = sizeof(len);
-    setsockopt(sockfd, SOL_SOCKET, SO_SNDBUF, (const char*)&len, optlen);
-
     hloop_t loop;
     hloop_init(&loop);
     uint64_t start_hrtime = hloop_now_hrtime(&loop);
 
-    hio_t* io = hio_get(&loop, sockfd);
-    if (io == NULL) return -1;
-    io->io_type = HIO_TYPE_IP;
-    struct sockaddr_in localaddr;
-    socklen_t addrlen = sizeof(localaddr);
-    memset(&localaddr, 0, addrlen);
-    localaddr.sin_addr.s_addr = htonl(INADDR_ANY);
-    hio_setlocaladdr(io, (struct sockaddr*)&localaddr, addrlen);
-
     nmap_udata_t udata;
     udata.nmap = nmap;
-    udata.io = io;
-    udata.send_cnt = sendto_icmp(io, nmap);
+    udata.send_cnt = 0;
     udata.recv_cnt = 0;
     udata.up_cnt = 0;
     udata.idle_cnt = 0;
-    io->userdata = &udata;
+    loop.userdata = &udata;
 
     char recvbuf[128];
-    hrecvfrom(&loop, sockfd, recvbuf, sizeof(recvbuf), on_recvfrom);
+    // icmp
+    char sendbuf[44]; // 20IP + 44ICMP = 64
+    icmp_t* icmp_req = (icmp_t*)sendbuf;
+    icmp_req->icmp_type = ICMP_ECHO;
+    icmp_req->icmp_code = 0;
+    icmp_req->icmp_id = getpid();
+    for (int i = 0; i < sizeof(sendbuf) - sizeof(icmphdr_t); ++i) {
+        icmp_req->icmp_data[i] = i;
+    }
+    struct sockaddr_in peeraddr;
+    hio_t* io = NULL;
+    for (auto iter = nmap->begin(); iter != nmap->end(); ++iter) {
+        if (iter->second == 1) continue;
+        if (udata.send_cnt % MAX_SENDTO_PERSOCKET == 0) {
+            // socket
+            int sockfd = socket(AF_INET, SOCK_RAW, IPPROTO_ICMP);
+            if (sockfd < 0) {
+                perror("socket");
+                if (errno == EPERM) {
+                    fprintf(stderr, "please use root or sudo to create a raw socket.\n");
+                }
+                return -socket_errno();
+            }
+            nonblocking(sockfd);
+            int len = 425984; // 416K
+            socklen_t optlen = sizeof(len);
+            setsockopt(sockfd, SOL_SOCKET, SO_SNDBUF, (const char*)&len, optlen);
+
+            io = hio_get(&loop, sockfd);
+            if (io == NULL) return -1;
+            io->io_type = HIO_TYPE_IP;
+            struct sockaddr_in localaddr;
+            socklen_t addrlen = sizeof(localaddr);
+            memset(&localaddr, 0, addrlen);
+            localaddr.sin_addr.s_addr = htonl(INADDR_ANY);
+            hio_setlocaladdr(io, (struct sockaddr*)&localaddr, addrlen);
+            hrecvfrom(&loop, sockfd, recvbuf, sizeof(recvbuf), on_recvfrom);
+        }
+        icmp_req->icmp_seq = iter->first;
+        icmp_req->icmp_cksum = 0;
+        icmp_req->icmp_cksum = checksum((uint8_t*)icmp_req, sizeof(sendbuf));
+        socklen_t addrlen = sizeof(peeraddr);
+        memset(&peeraddr, 0, addrlen);
+        peeraddr.sin_family = AF_INET;
+        peeraddr.sin_addr.s_addr = iter->first;
+        hio_setpeeraddr(io, (struct sockaddr*)&peeraddr, addrlen);
+        hsendto(io->loop, io->fd, sendbuf, sizeof(sendbuf), NULL);
+        ++udata.send_cnt;
+    }
 
     htimer_add(&loop, on_timer, MAX_RECVFROM_TIMEOUT, 1);
-    hidle_t* idle = hidle_add(&loop, on_idle, 3);
-    idle->userdata = &udata;
+    hidle_add(&loop, on_idle, 3);
 
     hloop_run(&loop);
     uint64_t end_hrtime = hloop_now_hrtime(&loop);