websock in simple c++

ws-c.cpp

#include <unistd.h>
#include <assert.h>
#include <iostream>
#include <string>

#include <wslay/wslay.h>
#include <sys/socket.h>
#include <netdb.h>
#include <cstring>
#include <poll.h>

#define FD (*(int*)user_data)
static ssize_t recv_callback(wslay_event_context_ptr ctx,
uint8_t *buf, size_t len,
int /*flags*/, void *user_data){
    auto r = recv(FD, buf, len, MSG_DONTWAIT);
    if(errno == EAGAIN || errno == EWOULDBLOCK) {
        wslay_event_set_error(ctx, WSLAY_ERR_WOULDBLOCK);
        return -1;
    }
    if(r < 0){
        wslay_event_set_error(ctx, WSLAY_ERR_CALLBACK_FAILURE);
        return -1;
    }
    return r;
}

ssize_t send_callback(wslay_event_context_ptr ctx,
                      const uint8_t *data, size_t len,
                      int flags, void *user_data){
    auto r = send(FD, data, len, MSG_DONTWAIT);
    if(errno == EAGAIN || errno == EWOULDBLOCK){
        wslay_event_set_error(ctx, WSLAY_ERR_WOULDBLOCK);
        return -1;
    }
    if(r < 0){
        wslay_event_set_error(ctx, WSLAY_ERR_CALLBACK_FAILURE);
        return -1;
    }
    return r;
}

int genmask_callback(wslay_event_context_ptr ctx,
                     uint8_t *buf, size_t len,
                     void *user_data){
    strncpy(reinterpret_cast<char*>(buf), "iflyos-embedded-client", len);
    return 0;
}

void on_frame_recv_start_callback(wslay_event_context_ptr ctx,
const struct wslay_event_on_frame_recv_start_arg *arg, void *user_data){

}

void on_frame_recv_chunk_callback(wslay_event_context_ptr ctx,
const struct wslay_event_on_frame_recv_chunk_arg *arg, void *user_data){

}

void on_frame_recv_end_callback(wslay_event_context_ptr ctx, void *user_data){

}

void on_msg_recv_callback(wslay_event_context_ptr ctx,
                          const struct wslay_event_on_msg_recv_arg *arg, void *user_data){
    if(arg->opcode == 1){
        //text message
        std::string msg(reinterpret_cast<const char*>(arg->msg), arg->msg_length);
        std::cout << "<" << msg << std::endl;
    }else if(arg->opcode == 2){
        std::cout << "<[binary message " << arg->msg_length << "]" << std::endl;
    }
}

static struct wslay_event_callbacks callbacks = {
    /*wslay_event_recv_callback*/ recv_callback,
    /*wslay_event_send_callback*/ send_callback,
    /*wslay_event_genmask_callback*/ genmask_callback,
    /*wslay_event_on_frame_recv_start_callback*/ on_frame_recv_start_callback,
    /*wslay_event_on_frame_recv_chunk_callback*/ on_frame_recv_chunk_callback,
    /*wslay_event_on_frame_recv_end_callback*/ on_frame_recv_end_callback,
    /*wslay_event_on_msg_recv_callback*/ on_msg_recv_callback,
};

int connect_to(const char *host, const char *service)
{
    struct addrinfo hints;
    int fd = -1;
    int r;
    memset(&hints, 0, sizeof(struct addrinfo));
    hints.ai_family = AF_UNSPEC;
    hints.ai_socktype = SOCK_STREAM;
    struct addrinfo *res;
    r = getaddrinfo(host, "http", &hints, &res);
    if(r != 0) {
        std::cerr << "getaddrinfo: " << gai_strerror(r) << std::endl;
        return -1;
    }
    for(struct addrinfo *rp = res; rp; rp = rp->ai_next) {
        fd = socket(rp->ai_family, rp->ai_socktype, rp->ai_protocol);
        if(fd == -1) {
            continue;
        }
        while((r = connect(fd, rp->ai_addr, rp->ai_addrlen)) == -1 &&
              errno == EINTR);
        if(r == 0) {
            break;
        }
        close(fd);
        fd = -1;
    }
    freeaddrinfo(res);
    return fd;
}

#define SEND(x) do{ \
auto r = send(fd, x, sizeof(x) - 1, 0); \
assert(r == sizeof(x) - 1);\
}while(0) \

int main() {
    auto fd = connect_to("demos.kaazing.com", nullptr);
    assert(fd > 0);
    SEND("GET /echo HTTP/1.1\r\n");
    SEND("Host: demos.kaazing.com\r\n");
    SEND("Upgrade: websocket\r\n");
    SEND("Connection: Upgrade\r\n");
    SEND("Sec-WebSocket-Key: fyxfhR8QIm3BSb0q/Tw5w==\r\n");
    SEND("Sec-WebSocket-Version: 13\r\n");
    SEND("\r\n");

    char headers[256];

    //blocking read all headers, server make sure all header data less than 256 bytes
    int rd = recv(fd, headers, sizeof(headers), 0);
    if(rd <= 0) {
        std::cerr << "http failed!" << std::endl;
        close(fd);
        exit(0);
    }
    std::cout << std::string(&headers[0], rd);

    wslay_event_context_ptr ws;
    auto ret = wslay_event_context_client_init(&ws, &callbacks, &fd);
    assert(ret == 0);
    wslay_event_msg msg = {1, reinterpret_cast<const uint8_t *>("hello world"), sizeof("hello world") - 1};
    //send a message wslay_event_queue_msg will copy message
    rd = wslay_event_queue_msg(ws, &msg);
    assert(rd == 0);
    rd = wslay_event_send(ws);
    assert(rd == 0);
    std::cout << ">hello world" << std::endl;

    //receive a message with timeout 1000
    struct pollfd poll_set[1];
    poll_set[0].fd = fd;
    poll_set[0].events = POLLIN | POLLERR;
    poll_set[0].revents = 0;
    poll(poll_set, 1, 1000);
    if(poll_set[0].revents & POLLIN) {
        rd = wslay_event_recv(ws);
        assert(rd == 0);
    } else if(poll_set[0].revents & POLLERR){
        std::cerr << "socket error" << std::endl;
    } else {
        std::cout << "recv timeout" << std::endl;
    }

    assert(rd == 0);

    //close and free
    wslay_event_queue_close(ws, 0, nullptr, 0);
    wslay_event_context_free(ws);
    close(fd);
    return 0;
}