anet.c

   1 /* anet.c -- Basic TCP socket stuff made a bit less boring
   2  *
   3  * Copyright (c) 2006-2009, Salvatore Sanfilippo <antirez at gmail dot com>
   4  * All rights reserved.
   5  *
   6  * Redistribution and use in source and binary forms, with or without
   7  * modification, are permitted provided that the following conditions are met:
   8  *
   9  *   * Redistributions of source code must retain the above copyright notice,
  10  *     this list of conditions and the following disclaimer.
  11  *   * Redistributions in binary form must reproduce the above copyright
  12  *     notice, this list of conditions and the following disclaimer in the
  13  *     documentation and/or other materials provided with the distribution.
  14  *   * Neither the name of Redis nor the names of its contributors may be used
  15  *     to endorse or promote products derived from this software without
  16  *     specific prior written permission.
  17  *
  18  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
  19  * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
  20  * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
  21  * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
  22  * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
  23  * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
  24  * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
  25  * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
  26  * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
  27  * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
  28  * POSSIBILITY OF SUCH DAMAGE.
  29  */
  30
  31 #include <sys/types.h>
  32 #include <sys/socket.h>
  33 #include <netinet/in.h>
  34 #include <netinet/tcp.h>
  35 #include <arpa/inet.h>
  36 #include <unistd.h>
  37 #include <fcntl.h>
  38 #include <string.h>
  39 #include <netdb.h>
  40 #include <errno.h>
  41 #include <stdarg.h>
  42 #include <stdio.h>
  43
  44 #include "anet.h"
  45
  46 static void anetSetError(char *err, const char *fmt, ...)
  47 {
  48     va_list ap;
  49
  50     if (!err) return;
  51     va_start(ap, fmt);
  52     vsnprintf(err, ANET_ERR_LEN, fmt, ap);
  53     va_end(ap);
  54 }
  55
  56 int anetNonBlock(char *err, int fd)
  57 {
  58     int flags;
  59
  60     /* Set the socket nonblocking.
  61      * Note that fcntl(2) for F_GETFL and F_SETFL can't be
  62      * interrupted by a signal. */
  63     if ((flags = fcntl(fd, F_GETFL)) == -1) {
  64         anetSetError(err, "fcntl(F_GETFL): %s\n", strerror(errno));
  65         return ANET_ERR;
  66     }
  67     if (fcntl(fd, F_SETFL, flags | O_NONBLOCK) == -1) {
  68         anetSetError(err, "fcntl(F_SETFL,O_NONBLOCK): %s\n", strerror(errno));
  69         return ANET_ERR;
  70     }
  71     return ANET_OK;
  72 }
  73
  74 int anetTcpNoDelay(char *err, int fd)
  75 {
  76     int yes = 1;
  77     if (setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &yes, sizeof(yes)) == -1)
  78     {
  79         anetSetError(err, "setsockopt TCP_NODELAY: %s\n", strerror(errno));
  80         return ANET_ERR;
  81     }
  82     return ANET_OK;
  83 }
  84
  85 int anetSetSendBuffer(char *err, int fd, int buffsize)
  86 {
  87     if (setsockopt(fd, SOL_SOCKET, SO_SNDBUF, &buffsize, sizeof(buffsize)) == -1)
  88     {
  89         anetSetError(err, "setsockopt SO_SNDBUF: %s\n", strerror(errno));
  90         return ANET_ERR;
  91     }
  92     return ANET_OK;
  93 }
  94
  95 int anetTcpKeepAlive(char *err, int fd)
  96 {
  97     int yes = 1;
  98     if (setsockopt(fd, SOL_SOCKET, SO_KEEPALIVE, &yes, sizeof(yes)) == -1) {
  99         anetSetError(err, "setsockopt SO_KEEPALIVE: %s\n", strerror(errno));
 100         return ANET_ERR;
 101     }
 102     return ANET_OK;
 103 }
 104
 105 int anetResolve(char *err, char *host, char *ipbuf)
 106 {
 107     struct sockaddr_in sa;
 108
 109     sa.sin_family = AF_INET;
 110     if (inet_aton(host, &sa.sin_addr) == 0) {
 111         struct hostent *he;
 112
 113         he = gethostbyname(host);
 114         if (he == NULL) {
 115             anetSetError(err, "can't resolve: %s\n", host);
 116             return ANET_ERR;
 117         }
 118         memcpy(&sa.sin_addr, he->h_addr, sizeof(struct in_addr));
 119     }
 120     strcpy(ipbuf,inet_ntoa(sa.sin_addr));
 121     return ANET_OK;
 122 }
 123
 124 #define ANET_CONNECT_NONE 0
 125 #define ANET_CONNECT_NONBLOCK 1
 126 static int anetTcpGenericConnect(char *err, char *addr, int port, int flags)
 127 {
 128     int s, on = 1;
 129     struct sockaddr_in sa;
 130
 131     if ((s = socket(AF_INET, SOCK_STREAM, 0)) == -1) {
 132         anetSetError(err, "creating socket: %s\n", strerror(errno));
 133         return ANET_ERR;
 134     }
 135     /* Make sure connection-intensive things like the redis benckmark
 136      * will be able to close/open sockets a zillion of times */
 137     setsockopt(s, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on));
 138
 139     sa.sin_family = AF_INET;
 140     sa.sin_port = htons(port);
 141     if (inet_aton(addr, &sa.sin_addr) == 0) {
 142         struct hostent *he;
 143
 144         he = gethostbyname(addr);
 145         if (he == NULL) {
 146             anetSetError(err, "can't resolve: %s\n", addr);
 147             close(s);
 148             return ANET_ERR;
 149         }
 150         memcpy(&sa.sin_addr, he->h_addr, sizeof(struct in_addr));
 151     }
 152     if (flags & ANET_CONNECT_NONBLOCK) {
 153         if (anetNonBlock(err,s) != ANET_OK)
 154             return ANET_ERR;
 155     }
 156     if (connect(s, (struct sockaddr*)&sa, sizeof(sa)) == -1) {
 157         if (errno == EINPROGRESS &&
 158             flags & ANET_CONNECT_NONBLOCK)
 159             return s;
 160
 161         anetSetError(err, "connect: %s\n", strerror(errno));
 162         close(s);
 163         return ANET_ERR;
 164     }
 165     return s;
 166 }
 167
 168 int anetTcpConnect(char *err, char *addr, int port)
 169 {
 170     return anetTcpGenericConnect(err,addr,port,ANET_CONNECT_NONE);
 171 }
 172
 173 int anetTcpNonBlockConnect(char *err, char *addr, int port)
 174 {
 175     return anetTcpGenericConnect(err,addr,port,ANET_CONNECT_NONBLOCK);
 176 }
 177
 178 /* Like read(2) but make sure 'count' is read before to return
 179  * (unless error or EOF condition is encountered) */
 180 int anetRead(int fd, void *buf, int count)
 181 {
 182     int nread, totlen = 0;
 183     while(totlen != count) {
 184         nread = read(fd,buf,count-totlen);
 185         if (nread == 0) return totlen;
 186         if (nread == -1) return -1;
 187         totlen += nread;
 188         buf += nread;
 189     }
 190     return totlen;
 191 }
 192
 193 /* Like write(2) but make sure 'count' is read before to return
 194  * (unless error is encountered) */
 195 int anetWrite(int fd, void *buf, int count)
 196 {
 197     int nwritten, totlen = 0;
 198     while(totlen != count) {
 199         nwritten = write(fd,buf,count-totlen);
 200         if (nwritten == 0) return totlen;
 201         if (nwritten == -1) return -1;
 202         totlen += nwritten;
 203         buf += nwritten;
 204     }
 205     return totlen;
 206 }
 207
 208 int anetTcpServer(char *err, int port, char *bindaddr)
 209 {
 210     int s, on = 1;
 211     struct sockaddr_in sa;
 212
 213     if ((s = socket(AF_INET, SOCK_STREAM, 0)) == -1) {
 214         anetSetError(err, "socket: %s\n", strerror(errno));
 215         return ANET_ERR;
 216     }
 217     if (setsockopt(s, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on)) == -1) {
 218         anetSetError(err, "setsockopt SO_REUSEADDR: %s\n", strerror(errno));
 219         close(s);
 220         return ANET_ERR;
 221     }
 222     memset(&sa,0,sizeof(sa));
 223     sa.sin_family = AF_INET;
 224     sa.sin_port = htons(port);
 225     sa.sin_addr.s_addr = htonl(INADDR_ANY);
 226     if (bindaddr) {
 227         if (inet_aton(bindaddr, &sa.sin_addr) == 0) {
 228             anetSetError(err, "Invalid bind address\n");
 229             close(s);
 230             return ANET_ERR;
 231         }
 232     }
 233     if (bind(s, (struct sockaddr*)&sa, sizeof(sa)) == -1) {
 234         anetSetError(err, "bind: %s\n", strerror(errno));
 235         close(s);
 236         return ANET_ERR;
 237     }
 238     if (listen(s, 32) == -1) {
 239         anetSetError(err, "listen: %s\n", strerror(errno));
 240         close(s);
 241         return ANET_ERR;
 242     }
 243     return s;
 244 }
 245
 246 int anetAccept(char *err, int serversock, char *ip, int *port)
 247 {
 248     int fd;
 249     struct sockaddr_in sa;
 250     unsigned int saLen;
 251
 252     while(1) {
 253         saLen = sizeof(sa);
 254         fd = accept(serversock, (struct sockaddr*)&sa, &saLen);
 255         if (fd == -1) {
 256             if (errno == EINTR)
 257                 continue;
 258             else {
 259                 anetSetError(err, "accept: %s\n", strerror(errno));
 260                 return ANET_ERR;
 261             }
 262         }
 263         break;
 264     }
 265     if (ip) strcpy(ip,inet_ntoa(sa.sin_addr));
 266     if (port) *port = ntohs(sa.sin_port);
 267     return fd;
 268 }