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