summaryrefslogtreecommitdiffstats
path: root/tools/COOLWebSocket.hpp
blob: 65ba1acf9d9d5b10f393d34f75df41cdf8ccc1d9 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
/* -*- Mode: C++; tab-width: 4; indent-tabs-mode: nil; c-basic-offset: 4; fill-column: 100 -*- */
/*
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
 */

#pragma once

#include <test/testlog.hpp>

#include <cstdlib>
#include <mutex>
#include <thread>

#include <Poco/Net/WebSocket.h>

#include <Common.hpp>
#include <Protocol.hpp>
#include <Log.hpp>

/// Deprecated: do not use ... replaced by net/Socket.hpp
///
/// WebSocket that is thread safe, and handles large frames transparently.
/// Careful - sendFrame and receiveFrame are _not_ virtual,
/// we need to make sure that we use COOLWebSocket all over the place.
/// It would be a kind of more natural to encapsulate Poco::Net::WebSocket
/// instead of inheriting (from that reason,) but that would requite much
/// larger code changes.
class COOLWebSocket : public Poco::Net::WebSocket
{
public:
    COOLWebSocket(const Socket& socket) :
        Poco::Net::WebSocket(socket)
    {
    }

    COOLWebSocket(Poco::Net::HTTPClientSession& cs,
                  Poco::Net::HTTPRequest& request,
                  Poco::Net::HTTPResponse& response) :
        Poco::Net::WebSocket(cs, request, response)
    {
    }

    /// Wrapper for Poco::Net::WebSocket::receiveFrame() that handles PING frames
    /// (by replying with a PONG frame) and PONG frames. PONG frames are ignored.

    /// Returns number of bytes received, or 0 if the Poco receiveFrame() returned 0,
    /// or -1 if no "interesting" (not PING or PONG) frame was actually received).

    /// Should we also factor out the handling of non-final and continuation frames into this?
    int receiveFrame(char* buffer, const int length, int& flags,
                     const std::string& testname = std::string())
    {
        // Timeout is in microseconds. We don't need this, except to yield the cpu.
        static const Poco::Timespan waitTime(POLL_TIMEOUT_MICRO_S / 10);

        while (poll(waitTime, Poco::Net::Socket::SELECT_READ))
        {
            const int n = Poco::Net::WebSocket::receiveFrame(buffer, length, flags);

            if (n <= 0)
                TST_LOG("Got nothing (" << n << ')');
            else
                TST_LOG("Got frame: " << getAbbreviatedFrameDump(buffer, n, flags));

            if ((flags & WebSocket::FRAME_OP_BITMASK) == WebSocket::FRAME_OP_CLOSE)
            {
                // Nothing to do.
                return n;
            }

            if ((flags & WebSocket::FRAME_OP_BITMASK) == WebSocket::FRAME_OP_PING)
            {
                // Echo back the ping message.
                if (Poco::Net::WebSocket::sendFrame(buffer, n,
                                                    static_cast<int>(WebSocket::FRAME_FLAG_FIN)
                                                        | WebSocket::FRAME_OP_PONG)
                    != n)
                {
                    TST_LOG("WARN: Sending Pong failed.");
                    return -1;
                }
            }
            else if ((flags & WebSocket::FRAME_OP_BITMASK) == WebSocket::FRAME_OP_PONG)
            {
                // In case we do send pings in the future.
            }
            else
            {
                return n;
            }
        }

        // Not ready for read.
        return -1;
    }

    /// Wrapper for Poco::Net::WebSocket::sendFrame() that handles large frames.
    int sendFrame(const char* buffer, const int length, const int flags = FRAME_TEXT,
                  const std::string& testname = std::string())
    {
        const int result = Poco::Net::WebSocket::sendFrame(buffer, length, flags);

        if (result != length)
        {
            TST_LOG("ERROR: Sent incomplete message, expected "
                    << length << " bytes but sent " << result
                    << " for: " << getAbbreviatedFrameDump(buffer, length, flags));
        }
        else
        {
            TST_LOG("Sent frame: " << getAbbreviatedFrameDump(buffer, length, flags));
        }

        return result;
    }

    /// Safe shutdown by sending a normal close frame, if socket is not in error,
    /// or, otherwise, close the socket without sending close frame, if it is.
    void shutdown(const std::string& testname = std::string())
    {
        shutdown(Poco::Net::WebSocket::StatusCodes::WS_NORMAL_CLOSE, testname);
    }

    /// Safe shutdown by sending a specific close frame, if socket is not in error,
    /// or, otherwise, close the socket without sending close frame, if it is.
    void shutdown(Poco::UInt16 statusCode, const std::string& statusMessage = std::string(),
                  const std::string& testname = std::string())
    {
        try
        {
            // Calling shutdown, in case of error, would try to send a 'close' frame
            // which won't work in case of broken pipe or timeout from peer. Just close the
            // socket in that case preventing 'close' frame from being sent.
            if (Poco::Net::WebSocket::poll(Poco::Timespan(0), Socket::SelectMode::SELECT_ERROR))
            {
                Poco::Net::WebSocket::close();
            }
            else
            {
                Poco::Net::WebSocket::shutdown(statusCode, statusMessage);
            }
        }
        catch (const Poco::Exception& exc)
        {
            TST_LOG("WARN: COOLWebSocket::shutdown: Exception: "
                    << exc.displayText()
                    << (exc.nested() ? " (" + exc.nested()->displayText() + ')' : ""));

            // Just close it.
            try
            {
                Poco::Net::WebSocket::close();
            }
            catch (const std::exception&)
            {
                // Nothing we can do.
            }
        }
    }

    // Return a string dump of a WebSocket frame: Its opcode, length, first line (if present),
    // flags. For human-readable logging purposes. Format not guaranteed to be stable. Not to be
    // inspected programmatically.
    static inline
    std::string getAbbreviatedFrameDump(const char *message, const int length, const int flags)
    {
        std::ostringstream result;
        switch (flags & Poco::Net::WebSocket::FRAME_OP_BITMASK)
        {
#define CASE(x)                                                                                    \
    case Poco::Net::WebSocket::FRAME_OP_##x:                                                       \
        result << #x << " (0x" << std::hex << flags << ')';                                        \
        break
            CASE(CONT);
            CASE(TEXT);
            CASE(BINARY);
            CASE(CLOSE);
            CASE(PING);
            CASE(PONG);
#undef CASE
            default:
                result << "UNKNOWN (0x" << std::hex << flags << ')';
                break;
        }

        result << ' ' << std::setw(3) << std::dec << length << " bytes"
               << (flags & Poco::Net::WebSocket::FRAME_FLAG_FIN ? " (FIN)" : "");

        if (length > 0 &&
            ((flags & Poco::Net::WebSocket::FRAME_OP_BITMASK) == Poco::Net::WebSocket::FRAME_OP_TEXT ||
             (flags & Poco::Net::WebSocket::FRAME_OP_BITMASK) == Poco::Net::WebSocket::FRAME_OP_BINARY ||
             (flags & Poco::Net::WebSocket::FRAME_OP_BITMASK) == Poco::Net::WebSocket::FRAME_OP_PING ||
             (flags & Poco::Net::WebSocket::FRAME_OP_BITMASK) == Poco::Net::WebSocket::FRAME_OP_PONG))
            result << ": '" << COOLProtocol::getAbbreviatedMessage(message, length) << '\'';
        return result.str();
    }
};

/* vim:set shiftwidth=4 softtabstop=4 expandtab: */