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