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