civetweb/examples/ws_server/ws_server.c

// Copyright (c) 2004-2012 Sergey Lyubka
// This file is a part of civetweb project, http://github.com/sunsetbrew/civetweb
//
// v 0.1 Contributed by William Greathouse    9-Sep-2013

#include <stdio.h>
#include <string.h>
#include <time.h>
#include <unistd.h>

#include "civetweb.h"

// simple structure for keeping track of websocket connection
struct ws_connection {
    struct mg_connection    *conn;
    int     update;
    int     closing;
    long    counter;
};

// time base and structure periodic updates to client for demo
#define BASETIME 100000 /* 0.1 seconds */
struct progress {
    int     limit;
    int     increment;
    int     period;
    int     value;
};

// up to 16 independent client connections
#define CONNECTIONS 16
static struct ws_connection ws_conn[CONNECTIONS];

#define PING_ACTIVE

#define PING_THREAD

#if defined(PING_ACTIVE) && defined(PING_THREAD)
// ws_ping_thread()
// Send periodic PING to assure websocket remains connected, except if we are closing
static void *ws_ping_thread(void *parm)
{
    int wsd = (long)parm;
    struct mg_connection    *conn = ws_conn[wsd].conn;

    while(!ws_conn[wsd].closing)
    {
        usleep(8000);   /* 8 ms */
        if (!ws_conn[wsd].closing)
            mg_websocket_write(conn, WEBSOCKET_OPCODE_PING, NULL, 0);
    }
    fprintf(stderr, "ws_ping_thread %d exiting\n", wsd);
}
#endif


// ws_server_thread()
// Simple demo server thread. Sends periodic updates to connected clients
static void *ws_server_thread(void *parm)
{
    int wsd = (long)parm;
    struct mg_connection    *conn = ws_conn[wsd].conn;
    int timer = 0;
    char tstr[32];
    int i;
    struct progress meter[] = {
        /* first meter 0 to 1000, by 5 every 0.1 second */
        { 1000, 5, 1, 0 },
        /* second meter 0 to 500, by 10 every 0.5 second */
        { 500, 10, 5, 0 },
        /* third meter 0 to 100, by 10 every 1.0 second */
        { 100, 10, 10, 0},
        /* end of list */
        { 0, 0, 0, 0}
    };

    fprintf(stderr, "ws_server_thread %d\n", wsd);

    /* Send initial meter updates */
    for (i=0; meter[i].period != 0; i++) {
        if (meter[i].value >= meter[i].limit)
            meter[i].value = 0;
        if (meter[i].value >= meter[i].limit)
            meter[i].value = meter[i].limit;
        sprintf(tstr, "meter%d:%d,%d", i+1,
                meter[i].value, meter[i].limit);
        mg_websocket_write(conn, WEBSOCKET_OPCODE_TEXT, tstr, strlen(tstr));
    }

#if defined(PING_ACTIVE) && defined(PING_THREAD)
    mg_start_thread(ws_ping_thread, (void *)(long)wsd);
#endif

    /* While the connection is open, send periodic updates */
    while(!ws_conn[wsd].closing) {
        usleep(10000); /* 0.01 second */
        timer++;

        /* Send meter updates */
        if (ws_conn[wsd].update) {
            for (i=0; meter[i].period != 0; i++) {
                if (timer%meter[i].period == 0) {
                    if (meter[i].value >= meter[i].limit)
                        meter[i].value = 0;
                    else
                        meter[i].value += meter[i].increment;
                    if (meter[i].value >= meter[i].limit)
                        meter[i].value = meter[i].limit;
                    // if we are closing, server should not send new data
                    if (!ws_conn[wsd].closing) {
                        sprintf(tstr, "meter%d:%d,%d", i+1,
                                meter[i].value, meter[i].limit);
                        mg_websocket_write(conn, WEBSOCKET_OPCODE_TEXT, tstr, strlen(tstr));
                    }
                }
            }
        }

#if defined(PING_ACTIVE) && !defined(PING_THREAD)
        /* Send periodic PING to assure websocket remains connected, except if we are closing */
        if (timer%100 == 0 && !ws_conn[wsd].closing)
            mg_websocket_write(conn, WEBSOCKET_OPCODE_PING, NULL, 0);
#endif
    }

    fprintf(stderr, "ws_server_thread %d exiting\n", wsd);

    // reset connection information to allow reuse by new client
    ws_conn[wsd].conn = NULL;
    ws_conn[wsd].update = 0;
    ws_conn[wsd].closing = 2;

    return NULL;
}

// websocket_connect_handler()
// On new client connection, find next available server connection and store
// new connection information. If no more server connections are available
// tell civetweb to not accept the client request.
static int websocket_connect_handler(const struct mg_connection *conn)
{
    int i;

    fprintf(stderr, "connect handler\n");

    for(i=0; i < CONNECTIONS; ++i) {
        if (ws_conn[i].conn == NULL) {
            fprintf(stderr, "...prep for server %d\n", i);
            ws_conn[i].conn = (struct mg_connection *)conn;
            ws_conn[i].closing = 0;
            ws_conn[i].update = 0;
            ws_conn[i].counter = -1;
            break;
        }
    }
    if (i >= CONNECTIONS) {
        fprintf(stderr, "Refused connection: Max connections exceeded\n");
        return 1;
    }

    return 0;
}

// websocket_ready_handler()
// Once websocket negotiation is complete, start a server for the connection
static void websocket_ready_handler(struct mg_connection *conn)
{
    int i;

    fprintf(stderr, "ready handler\n");

    for(i=0; i < CONNECTIONS; ++i) {
        if (ws_conn[i].conn == conn) {
            fprintf(stderr, "...start server %d\n", i);
            mg_start_thread(ws_server_thread, (void *)(long)i);
            break;
        }
    }
}

// websocket_close_handler()
// When websocket is closed, tell the associated server to shut down
static void websocket_close_handler(struct mg_connection *conn)
{
    int i;

    //fprintf(stderr, "close handler\n");   /* called for every close, not just websockets */

    for(i=0; i < CONNECTIONS; ++i) {
        if (ws_conn[i].conn == conn) {
            fprintf(stderr, "...close server %d\n", i);
            ws_conn[i].closing = 1;
        }
    }
}

// Arguments:
//   flags: first byte of websocket frame, see websocket RFC,
//          http://tools.ietf.org/html/rfc6455, section 5.2
//   data, data_len: payload data. Mask, if any, is already applied.
static int websocket_data_handler(struct mg_connection *conn, int flags,
                                  char *data, size_t data_len)
{
    int i;
    int wsd;

    for(i=0; i < CONNECTIONS; ++i) {
        if (ws_conn[i].conn == conn) {
            wsd = i;
            break;
        }
    }
    if (i >= CONNECTIONS) {
        fprintf(stderr, "Received websocket data from unknown connection\n");
        return 1;
    }

    if (flags & 0x80) {
        flags &= 0x7f;
        switch (flags) {
        case WEBSOCKET_OPCODE_CONTINUATION:
            fprintf(stderr, "CONTINUATION...\n");
            break;
        case WEBSOCKET_OPCODE_TEXT:
            //fprintf(stderr, "TEXT: %-.*s\n", (int)data_len, data);
            /*** interpret data as commands here ***/
            if (strncmp("update on", data, data_len)== 0) {
                fprintf(stderr, "TEXT: %-.*s\n", (int)data_len, data);
                /* turn on updates */
                ws_conn[wsd].update = 1;
                /* echo back */
                mg_websocket_write(conn, WEBSOCKET_OPCODE_TEXT, data, data_len);
            } else if (strncmp("update off", data, data_len)== 0) {
                fprintf(stderr, "TEXT: %-.*s\n", (int)data_len, data);
                /* turn off updates */
                ws_conn[wsd].update = 0;
                /* echo back */
                mg_websocket_write(conn, WEBSOCKET_OPCODE_TEXT, data, data_len);
            } else if (strncmp("counter ", data, 8)== 0) {
                char buffer[16];
                long newval;
                strncpy(buffer, &data[8], data_len-8);
                buffer[data_len-8] = '\0';
                newval = strtol(buffer, NULL, 0);
                if (ws_conn[wsd].counter == -1)
                    ws_conn[wsd].counter = newval;
                else
                    ws_conn[wsd].counter++;
                if (ws_conn[wsd].counter != newval)
                {
                    fprintf(stderr, "Counter: %ld, received %ld\n", ws_conn[wsd].counter, newval);
                    ws_conn[wsd].counter = newval;
                }
                if (ws_conn[wsd].counter % 3000 == 0)
                    fprintf(stderr, "Counter: %ld\n", ws_conn[wsd].counter);
            }
            else
                fprintf(stderr, "TEXT: %-.*s\n", (int)data_len, data);
            break;
        case WEBSOCKET_OPCODE_BINARY:
            fprintf(stderr, "BINARY...\n");
            break;
        case WEBSOCKET_OPCODE_CONNECTION_CLOSE:
            fprintf(stderr, "CLOSE...\n");
            /* If client initiated close, respond with close message in acknowlegement */
            if (!ws_conn[wsd].closing) {
                mg_websocket_write(conn, WEBSOCKET_OPCODE_CONNECTION_CLOSE, data, data_len);
                ws_conn[wsd].closing = 1; /* we should not send addional messages when close requested/acknowledged */
            }
            return 0; /* time to close the connection */
            break;
        case WEBSOCKET_OPCODE_PING:
            /* client sent PING, respond with PONG */
            mg_websocket_write(conn, WEBSOCKET_OPCODE_PONG, data, data_len);
            break;
        case WEBSOCKET_OPCODE_PONG:
            /* received PONG to our PING, no action */
            break;
        default:
            fprintf(stderr, "Unknown flags: %02x\n", flags);
            break;
        }
    }

    return 1;   /* keep connection open */
}


int main(void)
{
    char server_name[40];
    struct mg_context *ctx;
    struct mg_callbacks callbacks;
    const char *options[] = {
        "listening_ports", "8080",
        "document_root", "docroot",
        NULL
    };

    /* get simple greeting for the web server */
    snprintf(server_name, sizeof(server_name),
             "Civetweb websocket server v. %s",
             mg_version());

    memset(&callbacks, 0, sizeof(callbacks));
    callbacks.websocket_connect = websocket_connect_handler;
    callbacks.websocket_ready = websocket_ready_handler;
    callbacks.websocket_data = websocket_data_handler;
    callbacks.connection_close = websocket_close_handler;

    ctx = mg_start(&callbacks, NULL, options);

    /* show the greeting and some basic information */
    printf("%s started on port(s) %s with web root [%s]\n",
           server_name, mg_get_option(ctx, "listening_ports"),
           mg_get_option(ctx, "document_root"));

    getchar();  // Wait until user hits "enter"
    mg_stop(ctx);

    return 0;
}