tests/tcp_nuke_addr/tcp_nuke_addr_test.cpp

   1 #include <arpa/inet.h>
   2 #include <linux/if.h>
   3 #include <netinet/in.h>
   4 #include <stdlib.h>
   5 #include <sys/ioctl.h>
   6 #include <sys/socket.h>
   7 #include <unistd.h>
   8
   9 #include <atomic>
  10 #include <mutex>
  11 #include <thread>
  12
  13 #include "utils/RWLock.h"
  14
  15 // Defined only in ifc_utils.c, in the kernel, and in the NDK.
  16 #ifndef SIOCKILLADDR
  17 #define SIOCKILLADDR 0x8939
  18 #endif
  19
  20 #ifndef TCP_LINGER2
  21 #define TCP_LINGER2 8
  22 #endif
  23
  24 #define KILL_INTERVAL_MS 10
  25 #define CONNECT_THREADS 1
  26
  27 #define PERROR_EXIT(msg) { do { perror((msg)); exit(1); } while (0); };
  28
  29
  30 // Ensures that sockets don't stay in TIME_WAIT state.
  31 void setSoLinger(int s) {
  32     const struct linger l = {
  33         0,  // off
  34         0,  // 0 seconds
  35     };
  36     if (setsockopt(s, SOL_SOCKET, SO_LINGER, &l, sizeof(l)) == -1) {
  37         PERROR_EXIT("SO_LINGER");
  38     }
  39     const int nolinger = -1;
  40     if (setsockopt(s, SOL_TCP, TCP_LINGER2, &nolinger, sizeof(nolinger)) == -1) {
  41         PERROR_EXIT("TCP_LINGER2");
  42     }
  43 }
  44
  45
  46 // Binds to a random port on a random loopback address. We don't just use 127.0.0.1 because we don't
  47 // want this test to kill unrelated connections on loopback.
  48 int bindRandomAddr() {
  49     sockaddr_in sin;
  50     sin.sin_family = AF_INET;
  51     sin.sin_port = 0;
  52     sin.sin_addr.s_addr = htonl(INADDR_LOOPBACK);
  53
  54     while (sin.sin_addr.s_addr == htonl(INADDR_LOOPBACK)) {
  55         arc4random_buf(
  56              ((uint8_t *) &sin.sin_addr.s_addr) + 1,
  57              sizeof(sin.sin_addr.s_addr) - 1);
  58     }
  59
  60     int listensock;
  61     if ((listensock = socket(AF_INET, SOCK_STREAM, 0)) == -1) PERROR_EXIT("listensock");
  62     if (bind(listensock, (sockaddr *) &sin, sizeof(sin)) == -1) PERROR_EXIT("bind");
  63     if (listen(listensock, 10) == -1) PERROR_EXIT("listen");
  64
  65     return listensock;
  66 }
  67
  68
  69 // Thread that calls SIOCKILLADDR in a loop.
  70 void killSockets(sockaddr_in listenaddr, int intervalMs, android::RWLock *lock) {
  71     ifreq ifr;
  72     memset(&ifr, 0, sizeof(ifr));
  73     listenaddr.sin_port = 0;
  74     strncpy(ifr.ifr_name, "lo", strlen("lo"));
  75     memcpy(&ifr.ifr_addr, &listenaddr, sizeof(listenaddr));
  76
  77     int ioctlsock = socket(AF_INET, SOCK_DGRAM, 0);
  78     if (ioctlsock == -1) PERROR_EXIT("ioctlsock");
  79     while(true) {
  80         lock->writeLock();
  81         if (ioctl(ioctlsock, SIOCKILLADDR, &ifr) != 0) {
  82             PERROR_EXIT("SIOCKILLADDR failed, did you run 32-bit userspace on a 64-bit kernel?");
  83         }
  84         lock->unlock();
  85         std::this_thread::sleep_for(std::chrono::milliseconds(intervalMs));
  86     }
  87 }
  88
  89
  90 // Thread that calls connect() in a loop.
  91 void connectLoop(sockaddr_in listenaddr, int listensock,
  92         android::RWLock *lock, std::atomic<unsigned int> *attempts) {
  93     while(true) {
  94         int s = socket(AF_INET, SOCK_STREAM, 0);
  95         setSoLinger(s);
  96
  97         // Don't call SIOCKILLADDR while connect() is running, or we'll end up with lots of
  98         // connections in state FIN_WAITx or TIME_WAIT, which will then slow down future
  99         // due to SYN retransmits.
 100         lock->readLock();
 101         if (connect(s, (sockaddr *) &listenaddr, sizeof(listenaddr)) == -1) PERROR_EXIT("connect");
 102         lock->unlock();
 103
 104         send(s, "foo", 3, 0);
 105         int acceptedsock = accept(listensock, NULL, 0);
 106         if (close(acceptedsock) == -1) PERROR_EXIT("close");
 107         if (close(s) == -1) PERROR_EXIT("close");
 108
 109         attempts->fetch_add(1);
 110     }
 111 }
 112
 113
 114 // Thread that prints progress every second.
 115 void progressThread(std::atomic<unsigned int> *attempts) {
 116     uint32_t elapsed = 0;
 117     uint32_t total, previous = 0;
 118     while (true) {
 119         std::this_thread::sleep_for(std::chrono::seconds(1));
 120         elapsed++;
 121         total = attempts->load();
 122         printf("%ds: %u cps, total %u\n", elapsed, total-previous, total);
 123         fflush(stdout);
 124         previous = total;
 125     }
 126 }
 127
 128
 129 int main() {
 130     int listensock = bindRandomAddr();
 131     struct sockaddr_in sin;
 132     socklen_t len = sizeof(sin);
 133     if (getsockname(listensock, (sockaddr *) &sin, &len) == -1) PERROR_EXIT("getsockname");
 134
 135     printf("Using address %s:%d\n", inet_ntoa(sin.sin_addr), ntohs(sin.sin_port));
 136
 137     android::RWLock lock;
 138     std::atomic<unsigned int> attempts;
 139     attempts.store(0);
 140
 141     std::thread t0(killSockets, sin, KILL_INTERVAL_MS, &lock);
 142     std::thread *connectThreads[CONNECT_THREADS];
 143     for (size_t i = 0; i < CONNECT_THREADS; i++) {
 144         connectThreads[i] = new std::thread(connectLoop, sin, listensock, &lock, &attempts);
 145     }
 146     std::thread t1(progressThread, &attempts);
 147     t1.join();
 148
 149     return 0;
 150 }