boost/beast/websocket/detail/stream_base.hpp
// // Copyright (c) 2016-2017 Vinnie Falco (vinnie dot falco at gmail dot com) // // Distributed under the Boost Software License, Version 1.0. (See accompanying // file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt) // // Official repository: https://github.com/boostorg/beast // #ifndef BOOST_BEAST_WEBSOCKET_STREAM_BASE_HPP #define BOOST_BEAST_WEBSOCKET_STREAM_BASE_HPP #include <boost/beast/websocket/option.hpp> #include <boost/beast/websocket/detail/pmd_extension.hpp> #include <boost/beast/zlib/deflate_stream.hpp> #include <boost/beast/zlib/inflate_stream.hpp> #include <boost/beast/core/buffers_suffix.hpp> #include <boost/beast/core/error.hpp> #include <boost/beast/core/detail/chacha.hpp> #include <boost/beast/core/detail/integer_sequence.hpp> #include <boost/align/aligned_alloc.hpp> #include <boost/asio/buffer.hpp> #include <boost/core/exchange.hpp> #include <atomic> #include <cstdint> #include <memory> #include <new> #include <random> // Turn this on to avoid using thread_local //#define BOOST_BEAST_NO_THREAD_LOCAL 1 #ifdef BOOST_BEAST_NO_THREAD_LOCAL #include <atomic> #include <mutex> #endif namespace boost { namespace beast { namespace websocket { namespace detail { // used to order reads and writes class soft_mutex { int id_ = 0; public: soft_mutex() = default; soft_mutex(soft_mutex const&) = delete; soft_mutex& operator=(soft_mutex const&) = delete; soft_mutex(soft_mutex&& other) noexcept : id_(boost::exchange(other.id_, 0)) { } soft_mutex& operator=(soft_mutex&& other) noexcept { id_ = other.id_; other.id_ = 0; return *this; } // VFALCO I'm not too happy that this function is needed void reset() { id_ = 0; } bool is_locked() const { return id_ != 0; } template<class T> bool is_locked(T const*) const { return id_ == T::id; } template<class T> void lock(T const*) { BOOST_ASSERT(id_ == 0); id_ = T::id; } template<class T> void unlock(T const*) { BOOST_ASSERT(id_ == T::id); id_ = 0; } template<class T> bool try_lock(T const*) { // If this assert goes off it means you are attempting to // simultaneously initiate more than one of same asynchronous // operation, which is not allowed. For example, you must wait // for an async_read to complete before performing another // async_read. // BOOST_ASSERT(id_ != T::id); if(id_ != 0) return false; id_ = T::id; return true; } template<class T> bool try_unlock(T const*) { if(id_ != T::id) return false; id_ = 0; return true; } }; //------------------------------------------------------------------------------ struct stream_prng { bool secure_prng_ = true; struct prng_type { std::minstd_rand fast; beast::detail::chacha<20> secure; #if BOOST_BEAST_NO_THREAD_LOCAL prng_type* next = nullptr; #endif prng_type(std::uint32_t const* v, std::uint64_t stream) : fast(static_cast<typename decltype(fast)::result_type>( v[0] + v[1] + v[2] + v[3] + v[4] + v[5] + v[6] + v[7] + stream)) , secure(v, stream) { } }; class prng_ref { prng_type* p_; public: prng_ref& operator=(prng_ref&&) = delete; explicit prng_ref(prng_type& p) : p_(&p) { } prng_ref(prng_ref&& other) : p_(boost::exchange(other.p_, nullptr)) { } #ifdef BOOST_BEAST_NO_THREAD_LOCAL ~prng_ref() { if(p_) pool::impl().release(*p_); } #endif prng_type* operator->() const { return p_; } }; #ifndef BOOST_BEAST_NO_THREAD_LOCAL static prng_ref prng() { static std::atomic<std::uint64_t> stream{0}; thread_local prng_type p{seed(), stream++}; return prng_ref(p); } #else static prng_ref prng() { return prng_ref(pool::impl().acquire()); } #endif static std::uint32_t const* seed(std::seed_seq* ss = nullptr) { static seed_data d(ss); return d.v; } std::uint32_t create_mask() { auto p = prng(); if(secure_prng_) for(;;) if(auto key = p->secure()) return key; for(;;) if(auto key = p->fast()) return key; } private: struct seed_data { std::uint32_t v[8]; explicit seed_data(std::seed_seq* pss) { if(! pss) { std::random_device g; std::seed_seq ss{ g(), g(), g(), g(), g(), g(), g(), g()}; ss.generate(v, v+8); } else { pss->generate(v, v+8); } } }; #ifdef BOOST_BEAST_NO_THREAD_LOCAL class pool { prng_type* head_ = nullptr; std::atomic<std::uint64_t> n_{0}; std::mutex m_; public: ~pool() { for(auto p = head_; p;) { auto next = p->next; p->~prng_type(); boost::alignment::aligned_free(p); p = next; } } prng_type& acquire() { for(;;) { std::lock_guard<std::mutex> lock(m_); if(! head_) break; auto p = head_; head_ = head_->next; return *p; } auto p = boost::alignment::aligned_alloc( 16, sizeof(prng_type)); if(! p) BOOST_THROW_EXCEPTION(std::bad_alloc{}); return *(new(p) prng_type(seed(), n_++)); } void release(prng_type& p) { std::lock_guard<std::mutex> lock(m_); p.next = head_; head_ = &p; } static pool& impl() { static pool instance; return instance; } }; #endif }; //------------------------------------------------------------------------------ template<bool deflateSupported> struct stream_base : stream_prng { // State information for the permessage-deflate extension struct pmd_type { // `true` if current read message is compressed bool rd_set = false; zlib::deflate_stream zo; zlib::inflate_stream zi; }; std::unique_ptr<pmd_type> pmd_; // pmd settings or nullptr permessage_deflate pmd_opts_; // local pmd options detail::pmd_offer pmd_config_; // offer (client) or negotiation (server) // return `true` if current message is deflated bool rd_deflated() const { return pmd_ && pmd_->rd_set; } // set whether current message is deflated // returns `false` on protocol violation bool rd_deflated(bool rsv1) { if(pmd_) { pmd_->rd_set = rsv1; return true; } return ! rsv1; // pmd not negotiated } template<class ConstBufferSequence> bool deflate( boost::asio::mutable_buffer& out, buffers_suffix<ConstBufferSequence>& cb, bool fin, std::size_t& total_in, error_code& ec); void do_context_takeover_write(role_type role); void inflate( zlib::z_params& zs, zlib::Flush flush, error_code& ec); void do_context_takeover_read(role_type role); }; template<> struct stream_base<false> : stream_prng { // These stubs are for avoiding linking in the zlib // code when permessage-deflate is not enabled. bool rd_deflated() const { return false; } bool rd_deflated(bool rsv1) { return ! rsv1; } template<class ConstBufferSequence> bool deflate( boost::asio::mutable_buffer&, buffers_suffix<ConstBufferSequence>&, bool, std::size_t&, error_code&) { return false; } void do_context_takeover_write(role_type) { } void inflate( zlib::z_params&, zlib::Flush, error_code&) { } void do_context_takeover_read(role_type) { } }; } // detail } // websocket } // beast } // boost #endif