simplestore/server.c

   1 /* A very simple AWS-like server, without any authentication.  This is meant
   2  * performance testing in a local environment.  The server offers weak
   3  * guarantees on data durability--data is stored directly to the file system
   4  * without synchronization, so data might be lost in a crash.  This should
   5  * offer good performance though for benchmarking.
   6  *
   7  * Protocol: Each request is a whitespace-separated (typically, a single space)
   8  * list of parameters terminated by a newline character.  The response is a
   9  * line containing a response code and a body length (again white-separated and
  10  * newline-terminated) followed by the response body.  Requests:
  11  *   GET filename offset length
  12  *   PUT filename length
  13  *   DELETE filename
  14  *   LIST prefix
  15  */
  16
  17 #include <stdio.h>
  18 #include <stdlib.h>
  19 #include <assert.h>
  20 #include <errno.h>
  21 #include <fcntl.h>
  22 #include <signal.h>
  23 #include <string.h>
  24 #include <time.h>
  25 #include <unistd.h>
  26 #include <netinet/in.h>
  27 #include <sys/socket.h>
  28 #include <sys/stat.h>
  29 #include <sys/types.h>
  30 #include <sys/wait.h>
  31
  32 #define SIMPLESTORE_PORT 9541
  33
  34 /* Maximum number of connections the server will queue waiting for us to call
  35  * accept(). */
  36 #define TCP_BACKLOG 32
  37
  38 #define WHITESPACE " \t\r\n"
  39 #define MAX_ARGS 4
  40
  41 /* Maximum length of a command that we will accept. */
  42 #define MAX_CMD_LEN 4096
  43
  44 #define MAX_OBJECT_SIZE (256 << 20)
  45
  46 enum command { GET, PUT, LIST, DELETE };
  47
  48 void write_data(int fd, const char *buf, size_t len)
  49 {
  50     while (len > 0) {
  51         ssize_t bytes = write(fd, buf, len);
  52         if (bytes < 0) {
  53             if (errno == EINTR)
  54                 continue;
  55             exit(1);
  56         }
  57         buf += bytes;
  58         len -= bytes;
  59     }
  60 }
  61
  62 /* SIGCHLD handler, used to clean up processes that handle the connections. */
  63 static void sigchld_handler(int signal)
  64 {
  65     int pid;
  66     int reaped = 0;
  67
  68     while ((pid = waitpid(WAIT_ANY, NULL, WNOHANG)) > 0) {
  69         reaped++;
  70     }
  71
  72     if (pid < 0) {
  73         if (errno == ECHILD && reaped) {
  74             /* Don't print an error for the caes that we successfully cleaned
  75              * up after all children. */
  76         } else {
  77             perror("waitpid");
  78         }
  79     }
  80 }
  81
  82 /* Return a text representation of a socket address.  Returns a pointer to a
  83  * static buffer so it is non-reentrant. */
  84 const char *sockname(struct sockaddr_in *addr, socklen_t len)
  85 {
  86     static char buf[128];
  87     if (len < sizeof(struct sockaddr_in))
  88         return NULL;
  89     if (addr->sin_family != AF_INET)
  90         return NULL;
  91
  92     uint32_t ip = ntohl(addr->sin_addr.s_addr);
  93     sprintf(buf, "%d.%d.%d.%d:%d",
  94             (int)((ip >> 24) & 0xff),
  95             (int)((ip >> 16) & 0xff),
  96             (int)((ip >> 8) & 0xff),
  97             (int)(ip & 0xff),
  98             ntohs(addr->sin_port));
  99
 100     return buf;
 101 }
 102
 103 /* Convert a path from a client to the actual filename used.  Returns a string
 104  * that must be freed. */
 105 char *normalize_path(const char *in)
 106 {
 107     char *out = malloc(2 * strlen(in) + 1);
 108     int i, j;
 109     for (i = 0, j = 0; in[i] != '\0'; i++) {
 110         if (in[i] == '/') {
 111             out[j++] = '_';
 112             out[j++] = 's';
 113         } else if (in[i] == '_') {
 114             out[j++] = '_';
 115             out[j++] = 'u';
 116         } else {
 117             out[j++] = in[i];
 118         }
 119     }
 120     out[j++] = '\0';
 121     return out;
 122 }
 123
 124 void cmd_get(int fd, char *path, size_t start, ssize_t len)
 125 {
 126     char buf[65536];
 127
 128     char *response = "-1\n";
 129     int file = open(path, O_RDONLY);
 130     if (file < 0) {
 131         write_data(fd, response, strlen(response));
 132         return;
 133     }
 134
 135     struct stat statbuf;
 136     if (fstat(file, &statbuf) < 0) {
 137         write_data(fd, response, strlen(response));
 138         return;
 139     }
 140
 141     size_t datalen = statbuf.st_size;
 142     if (start > 0) {
 143         if (start >= datalen) {
 144             datalen = 0;
 145         } else {
 146             lseek(file, start, SEEK_SET);
 147             datalen -= start;
 148         }
 149     }
 150     if (len > 0 && len < datalen) {
 151         datalen = len;
 152     }
 153
 154     sprintf(buf, "%zd\n", datalen);
 155     write_data(fd, buf, strlen(buf));
 156
 157     while (datalen > 0) {
 158         size_t needed = datalen > sizeof(buf) ? sizeof(buf) : datalen;
 159         ssize_t bytes = read(file, buf, needed);
 160         if (bytes < 0 && errno == EINTR)
 161             continue;
 162         if (bytes <= 0) {
 163             /* Error reading necessary data, but we already told the client the
 164              * file size so pad the data to the client with null bytes. */
 165             memset(buf, 0, needed);
 166             bytes = needed;
 167         }
 168         write_data(fd, buf, bytes);
 169         datalen -= bytes;
 170     }
 171
 172     close(file);
 173 }
 174
 175 void cmd_put(int fd, char *path, char *buf,
 176              size_t object_size, size_t buf_used)
 177 {
 178     while (buf_used < object_size) {
 179         ssize_t bytes;
 180
 181         bytes = read(fd, buf + buf_used, object_size - buf_used);
 182         if (bytes < 0) {
 183             if (errno == EINTR)
 184                 continue;
 185             else
 186                 exit(1);
 187         }
 188
 189         if (bytes == 0)
 190             exit(1);
 191
 192         assert(bytes <= object_size - buf_used);
 193         buf_used += bytes;
 194
 195         continue;
 196     }
 197
 198     /* printf("Got %zd bytes for object '%s'\n", buf_used, path); */
 199
 200     char *response = "-1\n";
 201     int file = open(path, O_WRONLY|O_CREAT|O_TRUNC, 0600);
 202     if (file >= 0) {
 203         write_data(file, buf, object_size);
 204         response = "0\n";
 205         close(file);
 206     }
 207
 208     write_data(fd, response, strlen(response));
 209 }
 210
 211 /* The core handler for processing requests from the client.  This can be
 212  * single-threaded since each connection is handled in a separate process. */
 213 void handle_connection(int fd)
 214 {
 215     char cmdbuf[MAX_CMD_LEN];
 216     size_t buflen = 0;
 217
 218     while (1) {
 219         /* Keep reading data until reaching a newline, so that a complete
 220          * command can be parsed. */
 221         if (buflen == 0 || memchr(cmdbuf, '\n', buflen) == NULL) {
 222             ssize_t bytes;
 223
 224             if (buflen == MAX_CMD_LEN) {
 225                 /* Command is too long and thus malformed; close the
 226                  * connection. */
 227                 return;
 228             }
 229
 230             bytes = read(fd, cmdbuf + buflen, MAX_CMD_LEN - buflen);
 231             if (bytes < 0) {
 232                 if (errno == EINTR)
 233                     continue;
 234                 else
 235                     return;
 236             }
 237             if (bytes == 0)
 238                 return;
 239
 240             assert(bytes <= MAX_CMD_LEN - buflen);
 241             buflen += bytes;
 242
 243             continue;
 244         }
 245
 246         size_t cmdlen = (char *)memchr(cmdbuf, '\n', buflen) - cmdbuf + 1;
 247         cmdbuf[cmdlen - 1] = '\0';
 248         char *token;
 249         int arg_count;
 250         char *args[MAX_ARGS];
 251         int arg_int[MAX_ARGS];
 252         for (token = strtok(cmdbuf, WHITESPACE), arg_count = 0;
 253              token != NULL;
 254              token = strtok(NULL, WHITESPACE), arg_count++)
 255         {
 256             args[arg_count] = token;
 257             arg_int[arg_count] = atoi(token);
 258         }
 259
 260         if (arg_count < 2) {
 261             return;
 262         }
 263         char *path = normalize_path(args[1]);
 264         enum command cmd;
 265         if (strcmp(args[0], "GET") == 0 && arg_count == 4) {
 266             cmd = GET;
 267         } else if (strcmp(args[0], "PUT") == 0 && arg_count == 3) {
 268             cmd = PUT;
 269         } else if (strcmp(args[0], "DELETE") == 0 && arg_count == 2) {
 270             cmd = DELETE;
 271         } else {
 272             return;
 273         }
 274
 275         if (cmdlen < buflen)
 276             memmove(cmdbuf, cmdbuf + cmdlen, buflen - cmdlen);
 277         buflen -= cmdlen;
 278
 279         switch (cmd) {
 280         case GET:
 281             cmd_get(fd, path, arg_int[2], arg_int[3]);
 282             break;
 283         case PUT: {
 284             size_t object_size = arg_int[2];
 285             if (object_size > MAX_OBJECT_SIZE)
 286                 return;
 287             char *data_buf = malloc(object_size);
 288             if (data_buf == NULL)
 289                 return;
 290             size_t data_buflen = buflen > object_size ? object_size : buflen;
 291             if (data_buflen > 0)
 292                 memcpy(data_buf, cmdbuf, data_buflen);
 293             if (data_buflen < buflen) {
 294                 memmove(cmdbuf, cmdbuf + data_buflen, buflen - data_buflen);
 295                 buflen -= cmdlen;
 296             } else {
 297                 buflen = 0;
 298             }
 299             cmd_put(fd, path, data_buf, object_size, data_buflen);
 300             break;
 301         }
 302         case DELETE:
 303             //cmd_delete(fd, path);
 304             break;
 305         default:
 306             return;
 307         }
 308
 309         free(path);
 310     }
 311 }
 312
 313 /* Create a listening TCP socket on a new address (we do not use a fixed port).
 314  * Return the file descriptor of the listening socket. */
 315 int server_init()
 316 {
 317     int fd;
 318     struct sockaddr_in server_addr;
 319     socklen_t addr_len = sizeof(server_addr);
 320
 321     fd = socket(PF_INET, SOCK_STREAM, 0);
 322     if (fd < 0) {
 323         perror("socket");
 324         exit(1);
 325     }
 326
 327     int val = 1;
 328     setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &val, sizeof(val));
 329
 330     server_addr.sin_family = AF_INET;
 331     server_addr.sin_addr.s_addr = htonl(INADDR_ANY);
 332     server_addr.sin_port = htons(SIMPLESTORE_PORT);
 333     if (bind(fd, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
 334         perror("bind");
 335         exit(1);
 336     }
 337
 338     if (listen(fd, TCP_BACKLOG) < 0) {
 339         perror("listen");
 340         exit(1);
 341     }
 342
 343     if (getsockname(fd, (struct sockaddr *)&server_addr, &addr_len) < 0) {
 344         perror("getsockname");
 345         exit(1);
 346     }
 347
 348     printf("Server listening on %s ...\n", sockname(&server_addr, addr_len));
 349     fflush(stdout);
 350
 351     return fd;
 352 }
 353
 354 /* Process-based server main loop.  Wait for a connection, accept it, fork a
 355  * child process to handle the connection, and repeat.  Child processes are
 356  * reaped in the SIGCHLD handler. */
 357 void server_main(int listen_fd)
 358 {
 359     struct sigaction handler;
 360
 361     /* Install signal handler for SIGCHLD. */
 362     handler.sa_handler = sigchld_handler;
 363     sigemptyset(&handler.sa_mask);
 364     handler.sa_flags = SA_RESTART;
 365     if (sigaction(SIGCHLD, &handler, NULL) < 0) {
 366         perror("sigaction");
 367         exit(1);
 368     }
 369
 370     while (1) {
 371         struct sockaddr_in client_addr;
 372         socklen_t addr_len = sizeof(client_addr);
 373         int client_fd = accept(listen_fd, (struct sockaddr *)&client_addr,
 374                                &addr_len);
 375         int pid;
 376
 377         /* Very simple error handling.  Exit if something goes wrong.  Later,
 378          * might want to look into not killing off current connections abruptly
 379          * if we encounter an error in the accept(). */
 380         if (client_fd < 0) {
 381             if (errno == EINTR)
 382                 continue;
 383
 384             perror("accept");
 385             exit(1);
 386         }
 387
 388         printf("Accepted connection from %s ...\n",
 389                sockname(&client_addr, addr_len));
 390         fflush(stdout);
 391
 392         pid = fork();
 393         if (pid < 0) {
 394             perror("fork");
 395         } else if (pid == 0) {
 396             handle_connection(client_fd);
 397             printf("Closing connection %s ...\n",
 398                    sockname(&client_addr, addr_len));
 399             close(client_fd);
 400             exit(0);
 401         }
 402
 403         close(client_fd);
 404     }
 405 }
 406
 407 /* Print a help message describing command-line options to stdout. */
 408 static void usage(char *argv0)
 409 {
 410     printf("Usage: %s [options...]\n"
 411            "A simple key-value storage server.\n", argv0);
 412 }
 413
 414 int main(int argc, char *argv[])
 415 {
 416     int fd;
 417     //int i;
 418     int display_help = 0, cmdline_error = 0;
 419
 420 #if 0
 421     for (i = 1; i < argc; i++) {
 422         if (strcmp(argv[i], "-help") == 0) {
 423             display_help = 1;
 424         } else if (strcmp(argv[i], "-document_root") == 0) {
 425             i++;
 426             if (i >= argc) {
 427                 fprintf(stderr,
 428                         "Error: Argument to -document_root expected.\n");
 429                 cmdline_error = 1;
 430             } else {
 431                 document_root = argv[i];
 432             }
 433         } else if (strcmp(argv[i], "-port") == 0) {
 434             i++;
 435             if (i >= argc) {
 436                 fprintf(stderr,
 437                         "Error: Expected port number after -port.\n");
 438                 cmdline_error = 1;
 439             } else {
 440                 server_port = atoi(argv[i]);
 441                 if (server_port < 1 || server_port > 65535) {
 442                     fprintf(stderr,
 443                             "Error: Port must be between 1 and 65535.\n");
 444                     cmdline_error = 1;
 445                 }
 446             }
 447         } else if (strcmp(argv[i], "-mime_types") == 0) {
 448             i++;
 449             if (i >= argc) {
 450                 fprintf(stderr,
 451                         "Error: Argument to -mime_types expected.\n");
 452                 cmdline_error = 1;
 453             } else {
 454                 mime_types_file = argv[i];
 455             }
 456         } else if (strcmp(argv[i], "-log") == 0) {
 457             i++;
 458             if (i >= argc) {
 459                 fprintf(stderr,
 460                         "Error: Argument to -log expected.\n");
 461                 cmdline_error = 1;
 462             } else {
 463                 log_fname = argv[i];
 464             }
 465         } else {
 466             fprintf(stderr, "Error: Unrecognized option '%s'\n", argv[i]);
 467             cmdline_error = 1;
 468         }
 469     }
 470 #endif
 471
 472     /* Display a help message if requested, or let the user know to look at the
 473      * help message if there was an error parsing the command line.  In either
 474      * case, the server never starts. */
 475     if (display_help) {
 476         usage(argv[0]);
 477         exit(0);
 478     } else if (cmdline_error) {
 479         fprintf(stderr, "Run '%s -help' for a summary of options.\n", argv[0]);
 480         exit(2);
 481     }
 482
 483     fd = server_init();
 484     if (fd >= 0) {
 485         server_main(fd);
 486     }
 487
 488     return 0;
 489 }