handler.c

/**
 *
 * Copyright (c) 2010, Zed A. Shaw and Mongrel2 Project Contributors.
 * All rights reserved.
 * 
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are
 * met:
 * 
 *     * Redistributions of source code must retain the above copyright
 *       notice, this list of conditions and the following disclaimer.
 * 
 *     * Redistributions in binary form must reproduce the above copyright
 *       notice, this list of conditions and the following disclaimer in the
 *       documentation and/or other materials provided with the distribution.
 * 
 *     * Neither the name of the Mongrel2 Project, Zed A. Shaw, nor the names
 *       of its contributors may be used to endorse or promote products
 *       derived from this software without specific prior written
 *       permission.
 * 
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
 * IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
 * THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
 * PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
 * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
 * EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
 * PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
 * PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
 * NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

#include <handler.h>
#include <handler_parser.h>
#include <task/task.h>
#include <zmq.h>
#include <dbg.h>
#include <stdlib.h>
#include <string.h>
#include <connection.h>
#include <assert.h>
#include <register.h>

#include "setting.h"

struct tagbstring LEAVE_HEADER_JSON = bsStatic("{\"METHOD\":\"JSON\"}");
struct tagbstring LEAVE_HEADER_TNET = bsStatic("16:6:METHOD,4:JSON,}");
struct tagbstring LEAVE_MSG = bsStatic("{\"type\":\"disconnect\"}");


int HANDLER_STACK;

static void cstr_free(void *data, void *hint)
{
    free(data);
}

void Handler_notify_leave(Handler *handler, int id)
{
    void *socket = handler->send_socket;
    assert(socket && "Socket can't be NULL");
    bstring payload = NULL;
    
    if(handler->protocol == HANDLER_PROTO_TNET) {
        payload = bformat("%s %d @* %s%d:%s,",
                bdata(handler->send_ident), id,
                bdata(&LEAVE_HEADER_TNET),
                blength(&LEAVE_MSG), bdata(&LEAVE_MSG));
    } else {
        payload = bformat("%s %d @* %d:%s,%d:%s,",
                bdata(handler->send_ident), id,
                blength(&LEAVE_HEADER_JSON), bdata(&LEAVE_HEADER_JSON),
                blength(&LEAVE_MSG), bdata(&LEAVE_MSG));
    }

    check(payload != NULL, "Failed to make the payload for disconnect.");

    if(Handler_deliver(socket, bdata(payload), blength(payload)) == -1) {
        log_err("Can't tell handler %d died.", id);
    }

error: //fallthrough
    if(payload) free(payload);
}


int Handler_setup(Handler *handler)
{
    taskname("Handler_task");

    handler->task = taskself();

    handler->send_socket = Handler_send_create(bdata(handler->send_spec), bdata(handler->send_ident));
    check(handler->send_socket, "Failed to create handler socket.");

    handler->recv_socket = Handler_recv_create(bdata(handler->recv_spec), bdata(handler->recv_ident));
    check(handler->recv_socket, "Failed to create listener socket.");

    return 0;

error:
    return -1;

}

static inline void handler_cap_payload(bstring payload)
{
    if(payload->data[payload->slen - 1] == '\0') {
        btrunc(payload, blength(payload)-1);
    } else if(payload->data[payload->slen] != '\0') {
        bconchar(payload, '\0');
    }
}

static inline int deliver_payload(int raw, int fd, Connection *conn, bstring payload)
{
    int rc = 0;

    if(raw) {
        debug("Sending raw message to %d length %d", fd, blength(payload));
        rc = Connection_deliver_raw(conn, payload);
        check(rc != -1, "Error sending raw message to HTTP listener on FD %d, closing them.", fd);
    } else {
        debug("Sending BASE64 message to %d length %d", fd, blength(payload));
        handler_cap_payload(payload);
        rc = Connection_deliver(conn, payload);
        check(rc != -1, "Error sending to MSG listener on FD %d, closing them.", fd);
    }

    return 0;
error:
    return -1;
}

static inline void handler_process_request(Handler *handler, int id, int fd,
        Connection *conn, bstring payload)
{
    int rc = 0;

    if(conn == NULL) {
        debug("Ident %d (fd %d) is no longer connected.", id, fd);
        Handler_notify_leave(handler, id);
    } else {
        if(blength(payload) == 0) {
            rc = Register_disconnect(fd);
            check(rc != -1, "Register disconnect failed for: %d", fd);
        } else {
            int raw = conn->type != CONN_TYPE_MSG || handler->raw;

            rc = deliver_payload(raw, fd, conn, payload);
            check(rc != -1, "Failed to deliver to connection %d on socket %d", id, fd);
        }
    }

    return;

error:
    Register_disconnect(fd);  // return ignored
    return;
}


static inline int handler_recv_parse(Handler *handler, HandlerParser *parser)
{
    check(handler->running, "Called while handler wasn't running, that's not good.");

    zmq_msg_t *inmsg = calloc(sizeof(zmq_msg_t), 1);
    int rc = 0;

    check_mem(inmsg);

    rc = zmq_msg_init(inmsg);
    check(rc == 0, "Failed to initialize message.");

    taskstate("recv");

    rc = mqrecv(handler->recv_socket, inmsg, 0);
    check(rc == 0, "Receive on handler socket failed.");
    check(handler->running, "Handler marked as not running.");

    rc = HandlerParser_execute(parser, zmq_msg_data(inmsg), zmq_msg_size(inmsg));
    check(rc == 1, "Failed to parse message from handler.");

    check(parser->target_count > 0, "Message sent had 0 targets: %.*s",
            (int)zmq_msg_size(inmsg), (char *)zmq_msg_data(inmsg));

    zmq_msg_close(inmsg);
    free(inmsg);
    return 0;

error:
    if(inmsg) {
        zmq_msg_close(inmsg);
        free(inmsg);
    }
    return -1;
}


void Handler_task(void *v)
{
    int rc = 0;
    int i = 0;
    Handler *handler = (Handler *)v;
    HandlerParser *parser = NULL;
    int max_targets = Setting_get_int("limits.handler_targets", 128);
    log_info("MAX allowing limits.handler_targets=%d", 128);

    parser = HandlerParser_create(max_targets);
    check_mem(parser);

    check(Handler_setup(handler) == 0, "Failed to initialize handler, exiting.");

    while(handler->running && !task_was_signaled()) {
        taskstate("delivering");

        rc = handler_recv_parse(handler, parser);
        if(task_was_signaled()) {
            break;
        }

        if(rc != -1 && parser->target_count > 0) {
            for(i = 0; i < (int)parser->target_count; i++) {
                int id = (int)parser->targets[i];
                int fd = Register_fd_for_id(id);
                Connection *conn = (Connection *)Register_fd_exists(fd);

                handler_process_request(handler, id, fd, conn, parser->body);
            }
        } else {
            debug("Skipped invalid message from handler: %s", bdata(handler->send_spec));
        }

        HandlerParser_reset(parser);
    }

    debug("############################### HANDLER EXITED: %p, running: %d, task: %p",
            handler, handler->running, handler->task);
    handler->running = 0;
    handler->task = NULL;
    HandlerParser_destroy(parser);
    taskexit(0);

error:
    handler->running = 0;
    handler->task = NULL;
    HandlerParser_destroy(parser);
    log_err("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! HANDLER TASK DIED");
    taskexit(1);
}


int Handler_deliver(void *handler_socket, char *buffer, size_t len)
{
    int rc = 0;
    zmq_msg_t msg;

    rc = zmq_msg_init(&msg);
    check(rc == 0, "Failed to initialize 0mq message to send.");

    rc = zmq_msg_init_data(&msg, buffer, len, cstr_free, NULL);
    check(rc == 0, "Failed to init 0mq message data.");

    rc = mqsend(handler_socket, &msg, 0);
    check(rc == 0, "Failed to deliver 0mq message to handler.");

    // zeromq owns the ram now
    return 0;

error:
    free(buffer); // we own the ram now
    return -1;
}


void *Handler_send_create(const char *send_spec, const char *identity)
{
    int bind_attempts = 10;
    void *handler_socket = mqsocket(ZMQ_PUSH);
    int rc = zmq_setsockopt(handler_socket, ZMQ_IDENTITY, identity, strlen(identity));
    check(rc == 0, "Failed to set handler socket %s identity %s", send_spec, identity);

    log_info("Binding handler PUSH socket %s with identity: %s", send_spec, identity);

    rc = zmq_bind(handler_socket, send_spec);

    while(rc != 0 && bind_attempts-- > 0) {
        taskdelay(1000);
        log_warn("Failed to bind send socket trying again for: %s", send_spec);
        rc = zmq_bind(handler_socket, send_spec);
    }

    check(bind_attempts > 0, "Too many bind attempts for handler %s", send_spec);

    return handler_socket;

error:
    return NULL;
}

void *Handler_recv_create(const char *recv_spec, const char *uuid)
{
    int bind_attempts = 10;
    void *listener_socket = mqsocket(ZMQ_SUB);
    check(listener_socket, "Can't create ZMQ_SUB socket.");

    int rc = zmq_setsockopt(listener_socket, ZMQ_SUBSCRIBE, uuid, 0);
    check(rc == 0, "Failed to subscribe listener socket: %s", recv_spec);
    log_info("Binding listener SUB socket %s subscribed to: %s", recv_spec, uuid);

    rc = zmq_bind(listener_socket, recv_spec);

    while(rc != 0 && bind_attempts-- > 0) {
        taskdelay(1000);
        debug("Failed to bind recv socket trying again.");
        rc = zmq_bind(listener_socket, recv_spec);
    }

    check(bind_attempts > 0, "Too many bind attempts for handler %s", recv_spec);

    return listener_socket;

error:
    return NULL;
}

const int DEFAULT_HANDLER_STACK = 100 * 1024;

Handler *Handler_create(bstring send_spec, bstring send_ident,
        bstring recv_spec, bstring recv_ident)
{
    debug("Creating handler %s:%s", bdata(send_spec), bdata(send_ident));

    if(!HANDLER_STACK) {
        HANDLER_STACK = Setting_get_int("limits.handler_stack", DEFAULT_HANDLER_STACK);
        log_info("MAX limits.handler_stack=%d", HANDLER_STACK);
    }

    Handler *handler = calloc(sizeof(Handler), 1);
    check_mem(handler);

    handler->send_ident = bstrcpy(send_ident);
    handler->recv_ident = bstrcpy(recv_ident);
    handler->recv_spec = bstrcpy(recv_spec);
    handler->send_spec = bstrcpy(send_spec);
    handler->running = 0;
    handler->raw = 0;
    handler->protocol = HANDLER_PROTO_JSON;

    return handler;
error:

    if(handler) free(handler);
    return NULL;
}


void Handler_destroy(Handler *handler)
{
    if(handler) {
        if(handler->recv_socket) zmq_close(handler->recv_socket);
        if(handler->send_socket) zmq_close(handler->send_socket);

        bdestroy(handler->send_ident);
        bdestroy(handler->recv_ident);
        bdestroy(handler->send_spec);
        bdestroy(handler->recv_spec);
        free(handler);
    }
}