libs/beast/example/websocket/server/stackless-ssl/websocket_server_stackless_ssl.cpp

   1 //
   2 // Copyright (c) 2016-2019 Vinnie Falco (vinnie dot falco at gmail dot com)
   3 //
   4 // Distributed under the Boost Software License, Version 1.0. (See accompanying
   5 // file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
   6 //
   7 // Official repository: https://github.com/boostorg/beast
   8 //
   9
  10 //------------------------------------------------------------------------------
  11 //
  12 // Example: WebSocket SSL server, stackless coroutine
  13 //
  14 //------------------------------------------------------------------------------
  15
  16 #include "example/common/server_certificate.hpp"
  17
  18 #include <boost/beast/core.hpp>
  19 #include <boost/beast/ssl.hpp>
  20 #include <boost/beast/websocket.hpp>
  21 #include <boost/beast/websocket/ssl.hpp>
  22 #include <boost/asio/coroutine.hpp>
  23 #include <boost/asio/strand.hpp>
  24 #include <boost/asio/dispatch.hpp>
  25 #include <algorithm>
  26 #include <cstdlib>
  27 #include <functional>
  28 #include <iostream>
  29 #include <memory>
  30 #include <string>
  31 #include <thread>
  32 #include <vector>
  33
  34 namespace beast = boost::beast;         // from <boost/beast.hpp>
  35 namespace http = beast::http;           // from <boost/beast/http.hpp>
  36 namespace websocket = beast::websocket; // from <boost/beast/websocket.hpp>
  37 namespace net = boost::asio;            // from <boost/asio.hpp>
  38 namespace ssl = boost::asio::ssl;       // from <boost/asio/ssl.hpp>
  39 using tcp = boost::asio::ip::tcp;       // from <boost/asio/ip/tcp.hpp>
  40
  41 //------------------------------------------------------------------------------
  42
  43 // Report a failure
  44 void
  45 fail(beast::error_code ec, char const* what)
  46 {
  47     std::cerr << what << ": " << ec.message() << "\n";
  48 }
  49
  50 // Echoes back all received WebSocket messages
  51 class session
  52     : public boost::asio::coroutine
  53     , public std::enable_shared_from_this<session>
  54 {
  55     websocket::stream<beast::ssl_stream<beast::tcp_stream>> ws_;
  56     beast::flat_buffer buffer_;
  57
  58 public:
  59     // Take ownership of the socket
  60     session(tcp::socket&& socket, ssl::context& ctx)
  61         : ws_(std::move(socket), ctx)
  62     {
  63     }
  64
  65     // Start the asynchronous operation
  66     void
  67     run()
  68     {
  69         // We need to be executing within a strand to perform async operations
  70         // on the I/O objects in this session. Although not strictly necessary
  71         // for single-threaded contexts, this example code is written to be
  72         // thread-safe by default.
  73         net::dispatch(ws_.get_executor(),
  74                       beast::bind_front_handler(&session::loop,
  75                                                 shared_from_this(),
  76                                                 beast::error_code{},
  77                                                 0));
  78     }
  79
  80     #include <boost/asio/yield.hpp>
  81
  82     void
  83     loop(
  84         beast::error_code ec,
  85         std::size_t bytes_transferred)
  86     {
  87         boost::ignore_unused(bytes_transferred);
  88
  89         reenter(*this)
  90         {
  91             // Set the timeout.
  92             beast::get_lowest_layer(ws_).expires_after(std::chrono::seconds(30));
  93
  94             // Perform the SSL handshake
  95             yield ws_.next_layer().async_handshake(
  96                 ssl::stream_base::server,
  97                 std::bind(
  98                     &session::loop,
  99                     shared_from_this(),
 100                     std::placeholders::_1,
 101                     0));
 102             if(ec)
 103                 return fail(ec, "handshake");
 104
 105             // Turn off the timeout on the tcp_stream, because
 106             // the websocket stream has its own timeout system.
 107             beast::get_lowest_layer(ws_).expires_never();
 108
 109             // Set suggested timeout settings for the websocket
 110             ws_.set_option(
 111                 websocket::stream_base::timeout::suggested(
 112                     beast::role_type::server));
 113
 114             // Set a decorator to change the Server of the handshake
 115             ws_.set_option(websocket::stream_base::decorator(
 116                 [](websocket::response_type& res)
 117                 {
 118                     res.set(http::field::server,
 119                         std::string(BOOST_BEAST_VERSION_STRING) +
 120                             " websocket-server-stackless-ssl");
 121                 }));
 122
 123             // Accept the websocket handshake
 124             yield ws_.async_accept(
 125                 std::bind(
 126                     &session::loop,
 127                     shared_from_this(),
 128                     std::placeholders::_1,
 129                     0));
 130             if(ec)
 131                 return fail(ec, "accept");
 132
 133             for(;;)
 134             {
 135                 // Read a message into our buffer
 136                 yield ws_.async_read(
 137                     buffer_,
 138                     std::bind(
 139                         &session::loop,
 140                         shared_from_this(),
 141                         std::placeholders::_1,
 142                         std::placeholders::_2));
 143                 if(ec == websocket::error::closed)
 144                 {
 145                     // This indicates that the session was closed
 146                     return;
 147                 }
 148                 if(ec)
 149                     fail(ec, "read");
 150
 151                 // Echo the message
 152                 ws_.text(ws_.got_text());
 153                 yield ws_.async_write(
 154                     buffer_.data(),
 155                     std::bind(
 156                         &session::loop,
 157                         shared_from_this(),
 158                         std::placeholders::_1,
 159                         std::placeholders::_2));
 160                 if(ec)
 161                     return fail(ec, "write");
 162
 163                 // Clear the buffer
 164                 buffer_.consume(buffer_.size());
 165             }
 166         }
 167     }
 168
 169     #include <boost/asio/unyield.hpp>
 170 };
 171
 172 //------------------------------------------------------------------------------
 173
 174 // Accepts incoming connections and launches the sessions
 175 class listener
 176     : public boost::asio::coroutine
 177     , public std::enable_shared_from_this<listener>
 178 {
 179     net::io_context& ioc_;
 180     ssl::context& ctx_;
 181     tcp::acceptor acceptor_;
 182     tcp::socket socket_;
 183
 184 public:
 185     listener(
 186         net::io_context& ioc,
 187         ssl::context& ctx,
 188         tcp::endpoint endpoint)
 189         : ioc_(ioc)
 190         , ctx_(ctx)
 191         , acceptor_(ioc)
 192         , socket_(ioc)
 193     {
 194         beast::error_code ec;
 195
 196         // Open the acceptor
 197         acceptor_.open(endpoint.protocol(), ec);
 198         if(ec)
 199         {
 200             fail(ec, "open");
 201             return;
 202         }
 203
 204         // Allow address reuse
 205         acceptor_.set_option(net::socket_base::reuse_address(true), ec);
 206         if(ec)
 207         {
 208             fail(ec, "set_option");
 209             return;
 210         }
 211
 212         // Bind to the server address
 213         acceptor_.bind(endpoint, ec);
 214         if(ec)
 215         {
 216             fail(ec, "bind");
 217             return;
 218         }
 219
 220         // Start listening for connections
 221         acceptor_.listen(
 222             net::socket_base::max_listen_connections, ec);
 223         if(ec)
 224         {
 225             fail(ec, "listen");
 226             return;
 227         }
 228     }
 229
 230     // Start accepting incoming connections
 231     void
 232     run()
 233     {
 234         loop();
 235     }
 236
 237 private:
 238
 239     #include <boost/asio/yield.hpp>
 240
 241     void
 242     loop(beast::error_code ec = {})
 243     {
 244         reenter(*this)
 245         {
 246             for(;;)
 247             {
 248                 yield acceptor_.async_accept(
 249                     socket_,
 250                     std::bind(
 251                         &listener::loop,
 252                         shared_from_this(),
 253                         std::placeholders::_1));
 254                 if(ec)
 255                 {
 256                     fail(ec, "accept");
 257                 }
 258                 else
 259                 {
 260                     // Create the session and run it
 261                     std::make_shared<session>(std::move(socket_), ctx_)->run();
 262                 }
 263
 264                 // Make sure each session gets its own strand
 265                 socket_ = tcp::socket(net::make_strand(ioc_));
 266             }
 267         }
 268     }
 269
 270     #include <boost/asio/unyield.hpp>
 271 };
 272
 273 //------------------------------------------------------------------------------
 274
 275 int main(int argc, char* argv[])
 276 {
 277     // Check command line arguments.
 278     if (argc != 4)
 279     {
 280         std::cerr <<
 281             "Usage: websocket-server-async-ssl <address> <port> <threads>\n" <<
 282             "Example:\n" <<
 283             "    websocket-server-async-ssl 0.0.0.0 8080 1\n";
 284         return EXIT_FAILURE;
 285     }
 286     auto const address = net::ip::make_address(argv[1]);
 287     auto const port = static_cast<unsigned short>(std::atoi(argv[2]));
 288     auto const threads = std::max<int>(1, std::atoi(argv[3]));
 289
 290     // The io_context is required for all I/O
 291     net::io_context ioc{threads};
 292
 293     // The SSL context is required, and holds certificates
 294     ssl::context ctx{ssl::context::tlsv12};
 295
 296     // This holds the self-signed certificate used by the server
 297     load_server_certificate(ctx);
 298
 299     // Create and launch a listening port
 300     std::make_shared<listener>(ioc, ctx, tcp::endpoint{address, port})->run();
 301
 302     // Run the I/O service on the requested number of threads
 303     std::vector<std::thread> v;
 304     v.reserve(threads - 1);
 305     for(auto i = threads - 1; i > 0; --i)
 306         v.emplace_back(
 307         [&ioc]
 308         {
 309             ioc.run();
 310         });
 311     ioc.run();
 312
 313     return EXIT_SUCCESS;
 314 }