libnmdb/tcp.c

   1
   2 #if ENABLE_TCP
   3
   4 #include <sys/types.h>          /* socket defines */
   5 #include <sys/socket.h>         /* socket functions */
   6 #include <stdlib.h>             /* malloc() */
   7 #include <stdint.h>             /* uint32_t and friends */
   8 #include <arpa/inet.h>          /* htonls() and friends */
   9 #include <string.h>             /* memcpy() */
  10 #include <unistd.h>             /* close() */
  11
  12 #include <netinet/tcp.h>        /* TCP stuff */
  13 #include <netdb.h>              /* gethostbyname() */
  14
  15 #include "nmdb.h"
  16 #include "net-const.h"
  17 #include "internal.h"
  18 #include "tcp.h"
  19
  20
  21 /* Used internally to really add the server once we have an IP address. */
  22 static int add_tcp_server_addr(nmdb_t *db, in_addr_t *inetaddr, int port)
  23 {
  24         int rv, fd;
  25         struct nmdb_srv *newsrv, *newarray;
  26
  27         fd = socket(AF_INET, SOCK_STREAM, 0);
  28         if (fd < 0)
  29                 return 0;
  30
  31         newarray = realloc(db->servers,
  32                         sizeof(struct nmdb_srv) * (db->nservers + 1));
  33         if (newarray == NULL) {
  34                 close(fd);
  35                 return 0;
  36         }
  37
  38         db->servers = newarray;
  39         db->nservers++;
  40
  41         if (port < 0)
  42                 port = TCP_SERVER_PORT;
  43
  44         newsrv = &(db->servers[db->nservers - 1]);
  45
  46         newsrv->fd = fd;
  47         newsrv->info.in.srvsa.sin_family = AF_INET;
  48         newsrv->info.in.srvsa.sin_port = htons(port);
  49         newsrv->info.in.srvsa.sin_addr.s_addr = *inetaddr;
  50
  51         rv = connect(fd, (struct sockaddr *) &(newsrv->info.in.srvsa),
  52                         sizeof(newsrv->info.in.srvsa));
  53         if (rv < 0)
  54                 goto error_exit;
  55
  56         /* Disable Nagle algorithm because we often send small packets. Huge
  57          * gain in performance. */
  58         rv = 1;
  59         if (setsockopt(fd, IPPROTO_TCP, TCP_NODELAY, &rv, sizeof(rv)) < 0 )
  60                 goto error_exit;
  61
  62         newsrv->type = TCP_CONN;
  63
  64         /* keep the list sorted by port, so we can do a reliable selection */
  65         qsort(db->servers, db->nservers, sizeof(struct nmdb_srv),
  66                         compare_servers);
  67
  68         return 1;
  69
  70 error_exit:
  71         close(fd);
  72         newarray = realloc(db->servers,
  73                         sizeof(struct nmdb_srv) * (db->nservers - 1));
  74         if (newarray == NULL) {
  75                 db->servers = NULL;
  76                 db->nservers = 0;
  77                 return 0;
  78         }
  79
  80         db->servers = newarray;
  81         db->nservers -= 1;
  82
  83         return 0;
  84 }
  85
  86 /* Same as nmdb_add_tipc_server() but for TCP connections. */
  87 int nmdb_add_tcp_server(nmdb_t *db, const char *addr, int port)
  88 {
  89         int rv;
  90         struct hostent *he;
  91         struct in_addr ia;
  92
  93         /* We try to resolve and then pass it to add_tcp_server_addr(). */
  94         rv = inet_pton(AF_INET, addr, &ia);
  95         if (rv <= 0) {
  96                 he = gethostbyname(addr);
  97                 if (he == NULL)
  98                         return 0;
  99
 100                 ia.s_addr = *( (in_addr_t *) (he->h_addr_list[0]) );
 101         }
 102
 103         return add_tcp_server_addr(db, &(ia.s_addr), port);
 104 }
 105
 106 int tcp_srv_send(struct nmdb_srv *srv, unsigned char *buf, size_t bsize)
 107 {
 108         ssize_t rv;
 109         uint32_t len;
 110
 111         len = htonl(bsize);
 112         memcpy(buf, (const void *) &len, 4);
 113
 114         rv = ssend(srv->fd, buf, bsize, 0);
 115         if (rv != bsize)
 116                 return 0;
 117         return 1;
 118 }
 119
 120 static ssize_t recv_msg(int fd, unsigned char *buf, size_t bsize)
 121 {
 122         ssize_t rv, t;
 123         uint32_t msgsize;
 124
 125         rv = recv(fd, buf, bsize, 0);
 126         if (rv <= 0)
 127                 return rv;
 128
 129         if (rv < 4) {
 130                 t = srecv(fd, buf + rv, 4 - rv, 0);
 131                 if (t <= 0) {
 132                         return t;
 133                 }
 134
 135                 rv = rv + t;
 136         }
 137
 138         msgsize = * ((uint32_t *) buf);
 139         msgsize = ntohl(msgsize);
 140
 141         if (msgsize > bsize)
 142                 return -1;
 143
 144         if (rv < msgsize) {
 145                 t = srecv(fd, buf + rv, msgsize - rv, 0);
 146                 if (t <= 0) {
 147                         return t;
 148                 }
 149
 150                 rv = rv + t;
 151         }
 152
 153         return rv;
 154 }
 155
 156
 157 /* Used internally to get and parse replies from the server. */
 158 uint32_t tcp_get_rep(struct nmdb_srv *srv,
 159                 unsigned char *buf, size_t bsize,
 160                 unsigned char **payload, size_t *psize)
 161 {
 162         ssize_t rv;
 163         uint32_t id, reply;
 164
 165         rv = recv_msg(srv->fd, buf, bsize);
 166         if (rv <= 0)
 167                 return -1;
 168
 169         id = * ((uint32_t *) buf + 1);
 170         id = ntohl(id);
 171         reply = * ((uint32_t *) buf + 2);
 172         reply = ntohl(reply);
 173
 174         if (id != ID_CODE) {
 175                 return -1;
 176         }
 177
 178         if (payload != NULL) {
 179                 *payload = buf + 4 + 4 + 4;
 180                 *psize = rv - 4 - 4 - 4;
 181         }
 182         return reply;
 183 }
 184
 185 #else
 186 /* Stubs to use when TCP is not enabled. */
 187
 188 #include <stdint.h>
 189 #include "nmdb.h"
 190 #include "tcp.h"
 191
 192 int nmdb_add_tcp_server(nmdb_t *db, const char *addr, int port)
 193 {
 194         return 0;
 195 }
 196
 197 int tcp_srv_send(struct nmdb_srv *srv, unsigned char *buf, size_t bsize)
 198 {
 199         return 0;
 200 }
 201
 202 uint32_t tcp_get_rep(struct nmdb_srv *srv,
 203                 unsigned char *buf, size_t bsize,
 204                 unsigned char **payload, size_t *psize)
 205 {
 206         return -1;
 207 }
 208
 209 #endif /* ENABLE_TCP */
 210