samples/wsendpoint.h

/*
 * This file has been partially derived from the WebSockets++ project at
 * https://github.com/zaphoyd/websocketpp which is licensed under a BSD-license.
 */

#ifndef WSENDPOINT_H
#define WSENDPOINT_H

#include <vector>
#include <sstream>
#include <iostream>
#include <string>
#include <boost/thread.hpp>
#include "wsframe.h"

#ifndef HAVE_BOOST_LOCK_GUARD
#include <pthread.h>
class MutexHelper {
    public:
        MutexHelper(pthread_mutex_t *mutex, bool locknow = true) :
            m_pMutex(mutex), m_bLocked(false)
        {
            if (locknow)
                Lock();
        }

        ~MutexHelper()
        {
            if (m_bLocked)
                pthread_mutex_unlock(m_pMutex);
        }

        void Lock()
        {
            pthread_mutex_lock(m_pMutex);
            m_bLocked = true;
        }

        void Unlock()
        {
            m_bLocked = false;
            pthread_mutex_unlock(m_pMutex);
        }
    private:
        pthread_mutex_t *m_pMutex;
        bool m_bLocked;

        MutexHelper(const MutexHelper &);
        MutexHelper & operator=(const MutexHelper &);
};
#endif

namespace wspp {
    class wsendpoint;

    class wshandler {
        public:
            void send_text(const std::string data) {
                send(data, frame::opcode::TEXT);
            }

            void send_binary(const std::string data) {
                send(data, frame::opcode::BINARY);
            }

            wshandler() : m_endpoint(0) {}

            virtual ~wshandler() {}

        private:
            virtual void on_message(std::string header, std::string data) = 0;
            virtual void on_close() = 0;
            virtual bool on_ping(const std::string data) = 0;
            virtual void on_pong(const std::string data) = 0;
            virtual void do_response(const std::string data) = 0;

            void send(const std::string& payload, frame::opcode::value op);

            // Non-copyable
            wshandler(const wshandler&);
            wshandler& operator=(const wshandler&);

            wsendpoint *m_endpoint;
            friend class wsendpoint;
    };

    class wsendpoint {
        private:
            // Non-copyable
            wsendpoint(const wsendpoint&);
            wsendpoint& operator=(const wsendpoint&);

        public:
            wsendpoint(wshandler *h)
                : m_rng(simple_rng())
                  , m_parser(frame::parser<simple_rng>(m_rng))
                  , m_state(session::state::OPEN)
                  , m_lock()
                  , m_handler(h)
        {
#ifndef HAVE_BOOST_LOCK_GUARD
            pthread_mutexattr_t mattr;
            pthread_mutexattr_init(&mattr);
            pthread_mutexattr_settype(&mattr, PTHREAD_MUTEX_RECURSIVE);
            pthread_mutex_init(&m_lock, &mattr);
            pthread_mutexattr_destroy(&mattr);
#endif
            m_handler->m_endpoint = this;
        }

#ifndef HAVE_BOOST_LOCK_GUARD
            ~wsendpoint() { pthread_mutex_destroy(&m_lock); }
#endif

            void AddRxData(std::string data)
            {
                std::istringstream s(data);
                while (m_state != session::state::CLOSED && s.rdbuf()->in_avail()) {
                    try {
                        m_parser.consume(s);
                        if (m_parser.ready()) {
                            if (m_parser.is_control()) {
                                process_control();
                            } else {
                                process_data();
                            }                   
                            m_parser.reset();
                        }
                    } catch (const tracing::wserror & e) {
                        if (m_parser.ready()) {
                            m_parser.reset();
                        }
                        switch(e.code()) {
                            case tracing::wserror::PROTOCOL_VIOLATION:
                                send_close(close::status::PROTOCOL_ERROR,e.what());
                                break;
                            case tracing::wserror::PAYLOAD_VIOLATION:
                                send_close(close::status::INVALID_PAYLOAD,e.what());
                                break;
                            case tracing::wserror::INTERNAL_ENDPOINT_ERROR:
                                send_close(close::status::INTERNAL_ENDPOINT_ERROR,e.what());
                                break;
                            case tracing::wserror::SOFT_ERROR:
                                continue;
                            case tracing::wserror::MESSAGE_TOO_BIG:
                                send_close(close::status::MESSAGE_TOO_BIG,e.what());
                                break;
                            case tracing::wserror::OUT_OF_MESSAGES:
                                // we need to wait for a message to be returned by the
                                // client. We exit the read loop. handle_read_frame
                                // will be restarted by recycle()
                                //m_read_state = WAITING;
                                //m_endpoint.wait(type::shared_from_this());
                                return;
                            default:
                                // Fatal error, forcibly end connection immediately.
                                std::cerr
                                    << "Dropping TCP due to unrecoverable exception: " << e.code()
                                    << " (" << e.what() << ")" << std::endl;
                                shutdown();
                        }
                        break;
                    }
                }
            }

            void send(const std::string& payload, frame::opcode::value op) {
                frame::parser<simple_rng> control(m_rng);
                control.set_opcode(op);
                control.set_fin(true);
                control.set_masked(false);
                control.set_payload(payload);

                std::string tmp(control.get_header_str());
                tmp.append(control.get_payload_str());
                m_handler->do_response(tmp);
            }

        private:
            void process_data() {
                m_handler->on_message(m_parser.get_header_str(), m_parser.get_payload_str());
            }


            void shutdown() {
#ifdef HAVE_BOOST_LOCK_GUARD
                boost::lock_guard<boost::recursive_mutex> lock(m_lock);
#else
                MutexHelper((pthread_mutex_t *)&m_lock);
#endif

                if (m_state == session::state::CLOSED) {return;}

                m_state = session::state::CLOSED;
                m_handler->on_close();
            }

            void pong(const std::vector<unsigned char> & payload) {
#ifdef HAVE_BOOST_LOCK_GUARD
                boost::lock_guard<boost::recursive_mutex> lock(m_lock);
#else
                MutexHelper((pthread_mutex_t *)&m_lock);
#endif

                if (m_state != session::state::OPEN) {return;}
                // if (m_detached) {return;}

                // TODO: optimize control messages and handle case where 
                // endpoint is out of messages
                frame::parser<simple_rng> control(m_rng);
                control.set_opcode(frame::opcode::PONG);
                control.set_fin(true);
                control.set_masked(false);
                control.set_payload(payload);

                std::string tmp(control.get_header_str());
                tmp.append(control.get_payload_str());
                m_handler->do_response(tmp);
            }


            void send_close(close::status::value code, const std::string& reason) {
#ifdef HAVE_BOOST_LOCK_GUARD
                boost::lock_guard<boost::recursive_mutex> lock(m_lock);
#else
                MutexHelper((pthread_mutex_t *)&m_lock);
#endif

                // if (m_detached) {return;}

                if (m_state != session::state::OPEN) {
                    std::cerr << "Tried to disconnect a session that wasn't open" << std::endl;
                    return;
                }

                if (close::status::invalid(code)) {
                    std::cerr << "Tried to close a connection with invalid close code: " 
                        << code << std::endl;
                    return;
                } else if (close::status::reserved(code)) {
                    std::cerr << "Tried to close a connection with reserved close code: " 
                        << code << std::endl;
                    return;
                }

                m_state = session::state::CLOSING;

                frame::parser<simple_rng> control(m_rng);
                control.set_opcode(frame::opcode::CLOSE);
                control.set_fin(true);
                control.set_masked(false);
                if (code != close::status::NO_STATUS) {
                    const uint16_t payload = htons(code);
                    std::string pl(reinterpret_cast<const char*>(&payload), 2);
                    pl.append(reason);
                    control.set_payload(pl);
                }

                std::string tmp(control.get_header_str());
                tmp.append(control.get_payload_str());
                m_handler->do_response(tmp);
            }


            void send_close_ack(close::status::value remote_close_code, std::string remote_close_reason) {
                close::status::value local_close_code;
                std::string local_close_reason;
                // echo close value unless there is a good reason not to.
                if (remote_close_code == close::status::NO_STATUS) {
                    local_close_code = close::status::NORMAL;
                    local_close_reason = "";
                } else if (remote_close_code == close::status::ABNORMAL_CLOSE) {
                    // TODO: can we possibly get here? This means send_close_ack was
                    //       called after a connection ended without getting a close
                    //       frame
                    throw "shouldn't be here";
                } else if (close::status::invalid(remote_close_code)) {
                    // TODO: shouldn't be able to get here now either
                    local_close_code = close::status::PROTOCOL_ERROR;
                    local_close_reason = "Status code is invalid";
                } else if (close::status::reserved(remote_close_code)) {
                    // TODO: shouldn't be able to get here now either
                    local_close_code = close::status::PROTOCOL_ERROR;
                    local_close_reason = "Status code is reserved";
                } else {
                    local_close_code = remote_close_code;
                    local_close_reason = remote_close_reason;
                }

                // TODO: check whether we should cancel the current in flight write.
                //       if not canceled the close message will be sent as soon as the
                //       current write completes.


                frame::parser<simple_rng> control(m_rng);
                control.set_opcode(frame::opcode::CLOSE);
                control.set_fin(true);
                control.set_masked(false);
                if (local_close_code != close::status::NO_STATUS) {
                    const uint16_t payload = htons(local_close_code);
                    std::string pl(reinterpret_cast<const char*>(&payload), 2);
                    pl.append(local_close_reason);
                    control.set_payload(pl);
                }

                std::string tmp(control.get_header_str());
                tmp.append(control.get_payload_str());
                m_handler->do_response(tmp);
                shutdown();
            }

            void process_control() {
                switch (m_parser.get_opcode()) {
                    case frame::opcode::PING:
                        if (m_handler->on_ping(m_parser.get_payload_str())) {
                            pong(m_parser.get_payload());
                        }
                        break;
                    case frame::opcode::PONG:
                        m_handler->on_pong(m_parser.get_payload_str());
                        break;
                    case frame::opcode::CLOSE:
                        // check that the codes we got over the wire are valid
                        if (m_state == session::state::OPEN) {
                            // other end is initiating
                            std::cerr << "sending close ack" << std::endl;

                            // TODO:
                            send_close_ack(m_parser.get_close_code(), m_parser.get_close_reason());
                        } else if (m_state == session::state::CLOSING) {
                            // ack of our close
                            std::cerr << "got close ack" << std::endl;
                            shutdown();
                        }
                        break;
                    default:
                        throw tracing::wserror("Invalid Opcode",
                                tracing::wserror::PROTOCOL_VIOLATION);
                        break;
                }
            }

        private:
            simple_rng m_rng;
            frame::parser<simple_rng> m_parser;
            session::state::value m_state;
#ifdef HAVE_BOOST_LOCK_GUARD
            mutable boost::recursive_mutex m_lock;
#else
            mutable pthread_mutex_t m_lock;
#endif
            wshandler *m_handler;
    };

    void wshandler::send(const std::string& payload, frame::opcode::value op)
    {
        if (m_endpoint) {
            m_endpoint->send(payload, op);
        }
    }

}

#endif