1 

     2 // splice

     3 #define _GNU_SOURCE

     4 #include <fcntl.h>

     5 

     6 #include <unistd.h>

     7 #include <stdlib.h>

     8 #include <signal.h>

     9 #include <errno.h>

    10 

    11 #include <event2/util.h>

    12 #include <event2/event.h>

    13 #include <event2/event_compat.h>

    14 #include <event2/event_struct.h>

    15 #include <pcl.h>

    16 

    17 #include "common.h"

    18 #include "config.h"

    19 #include "socket.h"

    20 

    21 // what port to listen on

    22 #define LISTEN_PORT 10559

    23 

    24 // stack size to use for coroutines

    25 #define CLIENT_STACK_SIZE 128 * 1024

    26 #define PIPELINE_STACK_SIZE 128 * 1024

    27 

    28 // this should be equal to the system page size

    29 #define PIPELINE_CHUNKSIZE 4096

    30 

    31 struct co_connect_info {

    32     // what co this is running as

    33     coroutine_t co;

    34 

    35     // set to the value of errno in case of an error, zero otherwise

    36     int error;

    37 };

    38 

    39 static void _co_connect_write (evutil_socket_t fd, short what, void *arg) {

    40     struct co_connect_info *ctx = arg;

    41     

    42     // retreive the error code

    43     // XXX: this function was lost in a hg-fuckup

    44     if (socket_connect_error(fd, &ctx->error))

    45         ctx->error = -1;   // mark it as failed

    46     

    47     // resume the co

    48     co_call(ctx->co);

    49 }

    50 

    51 // coroutine-connect

    52 int co_connect (struct config_endpoint *endpoint, int sock_type, int *sock) {

    53     // our state

    54     struct co_connect_info ctx;

    55 

    56     // the event for the connect

    57     struct event ev;

    58 

    59     // clean it

    60     memset(&ctx, 0, sizeof(ctx));

    61     

    62     // initiate the connect()

    63     if ((*sock = socket_connect_async(endpoint, sock_type)) == -1)

    64         goto error;

    65     

    66     // set up the event

    67     event_set(&ev, *sock, EV_WRITE, &_co_connect_write, &ctx);

    68 

    69     if (event_add(&ev, NULL))

    70         goto error;

    71     

    72     // store our co handle in there

    73     ctx.co = co_current();

    74 

    75     // wait for the result

    76     co_resume();

    77     

    78     // return the error code

    79     return ctx.error;

    80 

    81 error:

    82     return -1;

    83 }

    84 

    85 struct co_splice_info {

    86     int fd_in, fd_out;

    87     size_t len;

    88 

    89     short wait;

    90 

    91     struct event ev_read, ev_write;

    92 

    93     int error;

    94 

    95     coroutine_t co;

    96 };

    97 

    98 void _co_splice_ev (int fd, short what, void *arg) {

    99     struct co_splice_info *ctx = arg;

   100     

   101     ctx->wait &= ~what;

   102 

   103     if (!ctx->wait) {

   104         // both ends are readable/writeable!

   105 

   106         long ret;

   107 

   108         if ((ret = splice(ctx->fd_in, NULL, ctx->fd_out, NULL, ctx->len, SPLICE_F_MOVE | SPLICE_F_NONBLOCK)) == -1) {

   109             if (errno == EAGAIN) {

   110                 ctx->error = 0;

   111                 ctx->len = 0;

   112 

   113             } else {

   114                 PERROR("splice");

   115             }

   116         } else {

   117             ctx->error = 0;

   118             ctx->len = ret;

   119         }

   120 

   121         // let the co do it's thing

   122         co_call(ctx->co);

   123 

   124     }

   125     

   126 

   127     // ok, done succesfully

   128     return;

   129 

   130 error:

   131     // let the co handle this error

   132     ctx->error = errno;

   133     co_call(ctx->co);

   134 }

   135 

   136 // coroutine-splice

   137 int co_splice (int fd_in, int fd_out, size_t *len, int *error) {

   138     // first, try and splice directly

   139     // XXX: ensure all sockets are O_NONBLOCK!

   140     long ret;

   141 

   142     if ((ret = splice(fd_in, NULL, fd_out, NULL, *len, SPLICE_F_MOVE | SPLICE_F_NONBLOCK)) == -1) {

   143         if (errno == EAGAIN) {

   144             // wait for in/out to become readable/writeable

   145             

   146             // our state

   147             ZINIT(struct co_splice_info, ctx);

   148 

   149             ctx.fd_in = fd_in;

   150             ctx.fd_out = fd_out;

   151             ctx.len = *len;

   152             ctx.wait = EV_READ | EV_WRITE;

   153             ctx.co = co_current();

   154             

   155             // two events

   156             event_set(&ctx.ev_read, fd_in, EV_READ, &_co_splice_ev, &ctx);

   157             event_set(&ctx.ev_write, fd_out, EV_WRITE, &_co_splice_ev, &ctx);

   158 

   159             // add both

   160             if (

   161                     event_add(&ctx.ev_read, NULL)

   162                 ||  event_add(&ctx.ev_write, NULL)

   163             )

   164                 ERROR("event_add"); // XXX: cleanup

   165             

   166             // and wait...

   167             co_resume();

   168             

   169             // error

   170             if (ctx.error) {

   171                 *error = ctx.error;

   172 

   173                 return -1;

   174 

   175             } else {

   176                 *len = ctx.len;

   177 

   178                 return 0;

   179             }

   180         } else

   181             PERROR("splice");

   182     } else {

   183         // success!

   184         *len = ret;

   185         

   186         return 0;

   187     }

   188 

   189 error:

   190     return -1;

   191 }

   192 

   193 /*

   194 struct co_pipeline_info {

   195     int sock_read, sock_write;

   196 

   197     int pipe_read, pipe_write;

   198     

   199     / *

   200      * Two coroutines:

   201      *  *   co_write waits for co_read to indicate that there is enough data in the pipe to send, whereupon it will

   202      *      attempt to splice data from pipe_read to sock_write.

   203      *  *   co_read waits for data to become available on sock_src, splices it into pipe_write, and notifies co_write

   204      *      that there is data avilable for sending.

   205      * /

   206     coroutine_t co_read, co_write;

   207 

   208     char co_read_stack[PIPELINE_STACK_SIZE], co_write_stack[PIPELINE_STACK_SIZE];

   209 };

   210 

   211 void _co_pipeline_read (void *arg) {

   212     

   213 }

   214 

   215 void _co_pipeline_write (void *arg) {

   216 

   217 }

   218 */

   219 

   220 int co_pipeline (int sock_read, int sock_write) {

   221 //    ZINIT(struct co_pipeline_info, ctx);

   222     

   223     // create our pipe

   224     int pipefds[2];

   225 

   226     if (pipe(pipefds))

   227         PERROR("pipe");

   228     

   229     // copy over the fds into ctx

   230     int pipe_read = pipefds[0];

   231     int pipe_write = pipefds[1];

   232 

   233     // start shuffling data around

   234     size_t pipe_size = 0, read_len, write_len;

   235     int error = 0;

   236 

   237     while (1) {

   238         read_len = PIPELINE_CHUNKSIZE;

   239 

   240         if (co_splice(sock_read, pipe_write, &read_len, &error))

   241             PERROR("co_splice(sock_read, pipe_write, %zu)", read_len);

   242         

   243         if (!read_len) {

   244             // EOF

   245             break;

   246         }

   247 

   248         pipe_size += read_len;

   249 

   250         if (!pipe_size)

   251             continue;

   252 

   253         while (pipe_size > 0) {

   254             write_len = pipe_size;

   255 

   256             if (co_splice(pipe_read, sock_write, &write_len, &error))

   257                 PERROR("co_splice(pipe_read, sock_write, %zu", write_len);

   258 

   259             pipe_size -= write_len;

   260         }

   261     }

   262 

   263 /*

   264     // create the two coros

   265     if (

   266             (ctx.co_read = co_create(&_co_pipeline_read, &ctx, ctx.co_read_stack, PIPELINE_STACK_SIZE)) == NULL

   267         ||  (ctx.co_write = co_create(&_co_pipeline_write, &ctx, ctx.co_write_stack, PIPELINE_STACK_SIZE)) == NULL

   268     ) 

   269         ERROR("co_create");

   270     

   271     // first start off the read one, which then eventually calls the write one

   272     co_call(ctx.co_read);

   273 */

   274 

   275     // success

   276     return 0;

   277 

   278 error:

   279 /*

   280     if (ctx.co_read)

   281         co_delete(ctx.co_read);

   282 

   283     if (ctx.co_write)

   284         co_delete(ctx.co_write);

   285 */

   286     return -1;

   287 }

   288 

   289 /*

   290  * State needed to handle a client

   291  */

   292 struct client_info {

   293     // the client socket

   294     evutil_socket_t socket;

   295 

   296     // the connect target endpoint

   297     struct config_endpoint *connect_target;

   298 

   299     // the coroutine that handles it

   300     coroutine_t co;

   301 };

   302 

   303 void client_co (void *arg) {

   304     struct client_info *ctx = arg;

   305     

   306     // the outbound socket

   307     int conn_sock = -1;

   308 

   309     // connect

   310     if (co_connect(ctx->connect_target, SOCK_STREAM, &conn_sock))

   311         goto error;

   312     

   313     // pipe data around \o/

   314     if (co_pipeline(ctx->socket, conn_sock))

   315         goto error;

   316 

   317 error:

   318     // close the sockets

   319     if (conn_sock != -1)

   320         if (close(conn_sock))

   321             PWARNING("close connect_socket");

   322     

   323     if (close(ctx->socket))

   324         PWARNING("close listen_socket");

   325     

   326     // free the ctx

   327     free(ctx);

   328 

   329     // exit

   330     co_exit();

   331 }

   332 

   333 static void handle_accept (evutil_socket_t fd, short event, void *arg) {

   334     // the new client's client_info

   335     struct client_info *ctx = NULL;

   336     

   337     // the socket from accept()

   338     evutil_socket_t socket = -1;

   339 

   340     // not used

   341     struct sockaddr_storage addr;

   342     socklen_t addr_len;

   343     

   344     // arg is NULL and unused

   345     (void) arg;

   346     

   347     // accept the connection

   348     addr_len = sizeof(struct sockaddr_storage);

   349 

   350     if ((socket = accept(fd, (struct sockaddr *) &addr, &addr_len)) == -1)

   351         PERROR("accept");

   352     

   353     // alloc a new client_info

   354     if (!(ctx = calloc(1, sizeof(*ctx))))

   355         ERROR("calloc");

   356     

   357     // store the socket

   358     ctx->socket = socket;

   359 

   360     // and the endpoint

   361     ctx->connect_target = arg;

   362 

   363     // create the coroutine

   364     if ((ctx->co = co_create(&client_co, ctx, NULL, CLIENT_STACK_SIZE)) == NULL)

   365         ERROR("co_create");

   366 

   367     // we can start up the coroutine right away

   368     co_call(ctx->co);

   369     

   370     // done handling this accept

   371     return;

   372 

   373 error:

   374     if (ctx) {

   375         if (ctx->co)

   376             co_delete(ctx->co);

   377 

   378         free(ctx);

   379     }

   380 

   381     else if (socket != -1)

   382         close(socket);

   383 }

   384 

   385 

   386 int main (int argc, char **argv) {

   387      // libevent init

   388     struct event_base *ev_base = event_init();

   389 

   390     if (!ev_base)

   391         FATAL("event_init");

   392     

   393     // process command-line arguments

   394     int opt;

   395     const char *listen_spec = NULL, *connect_spec = NULL;

   396 

   397     while ((opt = getopt(argc, argv, "hl:c:")) != -1) {

   398         switch (opt) {

   399             case 'l':

   400                 if (listen_spec)

   401                     FATAL("only specify -l once");

   402 

   403                 listen_spec = optarg;

   404 

   405                 break;

   406 

   407             case 'c':

   408                 if (connect_spec)

   409                     FATAL("only specify -c once");

   410                 

   411                 connect_spec = optarg;

   412 

   413                 break;

   414 

   415             case 'h':

   416             default:

   417                 err_exit("Usage: %s [-h] [-l listen] -c <connect>", argv[0]);

   418         

   419         }

   420     }

   421 

   422     if (!connect_spec)

   423         err_exit("Must specify -c; see `%s -h`", argv[0]);

   424     

   425     // ignore SIGPIPE

   426     struct sigaction sigpipe;

   427     memset(&sigpipe, 0, sizeof(sigpipe));

   428 

   429     sigpipe.sa_handler = SIG_IGN;

   430 

   431     sigaction(SIGPIPE, &sigpipe, NULL);

   432     

   433     // endpoints

   434     struct config_endpoint listen_endpoint, connect_endpoint;

   435     

   436     endpoint_init(&listen_endpoint, LISTEN_PORT);

   437     endpoint_init(&connect_endpoint, 0);

   438 

   439     if (

   440             endpoint_parse(&listen_endpoint, listen_spec)

   441         ||  endpoint_parse(&connect_endpoint, connect_spec)

   442     )

   443         goto error;

   444     

   445     // the listen socket

   446     int listen_sock;

   447 

   448     if ((listen_sock = socket_listen(&listen_endpoint, SOCK_STREAM)) == -1)

   449         goto error;

   450 

   451     // create the listen event

   452     struct event listen_ev;

   453 

   454     event_set(&listen_ev, listen_sock, EV_READ | EV_PERSIST, &handle_accept, &connect_endpoint);

   455 

   456     if (event_add(&listen_ev, NULL))

   457         PERROR("event_add");

   458 

   459     // we shall now run

   460     INFO("run");

   461     

   462     // run the libevent mainloop

   463     if (event_base_dispatch(ev_base))

   464         WARNING("event_dispatch");

   465 

   466     INFO("shutdown");

   467 

   468     // cleanup

   469     event_base_free(ev_base);

   470     close(listen_sock);

   471 

   472     return EXIT_SUCCESS;

   473 

   474 error:

   475     return EXIT_FAILURE;

   476 }

   477