Home | History | Annotate | Download | only in lib 
      1 #include "toys.h"
      2 
      3 int xsocket(int domain, int type, int protocol)
      4 {
      5   int fd = socket(domain, type, protocol);
      6 
      7   if (fd < 0) perror_exit("socket %x %x", type, protocol);
      8   fcntl(fd, F_SETFD, FD_CLOEXEC);
      9 
     10   return fd;
     11 }
     12 
     13 void xsetsockopt(int fd, int level, int opt, void *val, socklen_t len)
     14 {
     15   if (-1 == setsockopt(fd, level, opt, val, len)) perror_exit("setsockopt");
     16 }
     17 
     18 // if !host bind to all local interfaces
     19 struct addrinfo *xgetaddrinfo(char *host, char *port, int family, int socktype,
     20   int protocol, int flags)
     21 {
     22   struct addrinfo info, *ai;
     23   int rc;
     24 
     25   memset(&info, 0, sizeof(struct addrinfo));
     26   info.ai_family = family;
     27   info.ai_socktype = socktype;
     28   info.ai_protocol = protocol;
     29   info.ai_flags = flags;
     30   if (!host) info.ai_flags |= AI_PASSIVE;
     31 
     32   rc = getaddrinfo(host, port, &info, &ai);
     33   if (rc || !ai)
     34     error_exit("%s%s%s: %s", host ? host : "*", port ? ":" : "",
     35       port ? port : "", rc ? gai_strerror(rc) : "not found");
     36 
     37   return ai;
     38 }
     39 
     40 int xconnbind(struct addrinfo *ai_arg, int dobind)
     41 {
     42   struct addrinfo *ai;
     43   int fd = -1, one = 1;
     44 
     45   // Try all the returned addresses. Report errors if last entry can't connect.
     46   for (ai = ai_arg; ai; ai = ai->ai_next) {
     47     fd = (ai->ai_next ? socket : xsocket)(ai->ai_family, ai->ai_socktype,
     48       ai->ai_protocol);
     49     xsetsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one));
     50     if (!(dobind ? bind : connect)(fd, ai->ai_addr, ai->ai_addrlen)) break;
     51     else if (!ai->ai_next) perror_exit_raw(dobind ? "bind" : "connect");
     52     close(fd);
     53   }
     54   freeaddrinfo(ai_arg);
     55 
     56   return fd;
     57 }
     58 
     59 int xconnect(struct addrinfo *ai)
     60 {
     61   return xconnbind(ai, 0);
     62 }
     63 
     64 
     65 int xbind(struct addrinfo *ai)
     66 {
     67   return xconnbind(ai, 1);
     68 }
     69 
     70 int xpoll(struct pollfd *fds, int nfds, int timeout)
     71 {
     72   int i;
     73   long long now, then = timeout>0 ? millitime() : 0;
     74 
     75   for (;;) {
     76     if (0<=(i = poll(fds, nfds, timeout)) || toys.signal) return i;
     77     if (errno != EINTR && errno != ENOMEM) perror_exit("xpoll");
     78     else {
     79       now = millitime();
     80       timeout -= now-then;
     81       then = now;
     82     }
     83   }
     84 }
     85 
     86 // Loop forwarding data from in1 to out1 and in2 to out2, handling
     87 // half-connection shutdown. timeouts return if no data for X ms.
     88 // Returns 0: both closed, 1 shutdown_timeout, 2 timeout
     89 int pollinate(int in1, int in2, int out1, int out2, int timeout, int shutdown_timeout)
     90 {
     91   struct pollfd pollfds[2];
     92   int i, pollcount = 2;
     93 
     94   memset(pollfds, 0, 2*sizeof(struct pollfd));
     95   pollfds[0].events = pollfds[1].events = POLLIN;
     96   pollfds[0].fd = in1;
     97   pollfds[1].fd = in2;
     98 
     99   // Poll loop copying data from each fd to the other one.
    100   for (;;) {
    101     if (!xpoll(pollfds, pollcount, timeout)) return pollcount;
    102 
    103     for (i=0; i<pollcount; i++) {
    104       if (pollfds[i].revents & POLLIN) {
    105         int len = read(pollfds[i].fd, libbuf, sizeof(libbuf));
    106         if (len<1) pollfds[i].revents = POLLHUP;
    107         else xwrite(i ? out2 : out1, libbuf, len);
    108       }
    109       if (pollfds[i].revents & POLLHUP) {
    110         // Close half-connection.  This is needed for things like
    111         // "echo GET / | netcat landley.net 80"
    112         // Note that in1 closing triggers timeout, in2 returns now.
    113         if (i) {
    114           shutdown(pollfds[0].fd, SHUT_WR);
    115           pollcount--;
    116           timeout = shutdown_timeout;
    117         } else return 0;
    118       }
    119     }
    120   }
    121 }
    122 
    123 // Return converted ipv4/ipv6 numeric address in libbuf
    124 char *ntop(struct sockaddr *sa)
    125 {
    126   void *addr;
    127 
    128   if (sa->sa_family == AF_INET) addr = &((struct sockaddr_in *)sa)->sin_addr;
    129   else addr = &((struct sockaddr_in6 *)sa)->sin6_addr;
    130 
    131   inet_ntop(sa->sa_family, addr, libbuf, sizeof(libbuf));
    132 
    133   return libbuf;
    134 }
    135 
    136 void xsendto(int sockfd, void *buf, size_t len, struct sockaddr *dest)
    137 {
    138   int rc = sendto(sockfd, buf, len, 0, dest,
    139     dest->sa_family == AF_INET ? sizeof(struct sockaddr_in) :
    140       sizeof(struct sockaddr_in6));
    141 
    142   if (rc != len) perror_exit("sendto");
    143 }
    144 
    145 // xrecvfrom with timeout in milliseconds
    146 int xrecvwait(int fd, char *buf, int len, union socksaddr *sa, int timeout)
    147 {
    148   socklen_t sl = sizeof(*sa);
    149 
    150   if (timeout >= 0) {
    151     struct pollfd pfd;
    152 
    153     pfd.fd = fd;
    154     pfd.events = POLLIN;
    155     if (!xpoll(&pfd, 1, timeout)) return 0;
    156   }
    157 
    158   len = recvfrom(fd, buf, len, 0, (void *)sa, &sl);
    159   if (len<0) perror_exit("recvfrom");
    160 
    161   return len;
    162 }
    163