stream_base.hpp 8.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414
  1. //
  2. // Copyright (c) 2016-2017 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. #ifndef BOOST_BEAST_WEBSOCKET_STREAM_BASE_HPP
  10. #define BOOST_BEAST_WEBSOCKET_STREAM_BASE_HPP
  11. #include <boost/beast/websocket/option.hpp>
  12. #include <boost/beast/websocket/detail/pmd_extension.hpp>
  13. #include <boost/beast/zlib/deflate_stream.hpp>
  14. #include <boost/beast/zlib/inflate_stream.hpp>
  15. #include <boost/beast/core/buffers_suffix.hpp>
  16. #include <boost/beast/core/error.hpp>
  17. #include <boost/beast/core/detail/chacha.hpp>
  18. #include <boost/beast/core/detail/integer_sequence.hpp>
  19. #include <boost/align/aligned_alloc.hpp>
  20. #include <boost/asio/buffer.hpp>
  21. #include <boost/core/exchange.hpp>
  22. #include <atomic>
  23. #include <cstdint>
  24. #include <memory>
  25. #include <new>
  26. #include <random>
  27. // Turn this on to avoid using thread_local
  28. //#define BOOST_BEAST_NO_THREAD_LOCAL 1
  29. #ifdef BOOST_BEAST_NO_THREAD_LOCAL
  30. #include <atomic>
  31. #include <mutex>
  32. #endif
  33. namespace boost {
  34. namespace beast {
  35. namespace websocket {
  36. namespace detail {
  37. // used to order reads and writes
  38. class soft_mutex
  39. {
  40. int id_ = 0;
  41. public:
  42. soft_mutex() = default;
  43. soft_mutex(soft_mutex const&) = delete;
  44. soft_mutex& operator=(soft_mutex const&) = delete;
  45. soft_mutex(soft_mutex&& other) noexcept
  46. : id_(boost::exchange(other.id_, 0))
  47. {
  48. }
  49. soft_mutex& operator=(soft_mutex&& other) noexcept
  50. {
  51. id_ = other.id_;
  52. other.id_ = 0;
  53. return *this;
  54. }
  55. // VFALCO I'm not too happy that this function is needed
  56. void reset()
  57. {
  58. id_ = 0;
  59. }
  60. bool is_locked() const
  61. {
  62. return id_ != 0;
  63. }
  64. template<class T>
  65. bool is_locked(T const*) const
  66. {
  67. return id_ == T::id;
  68. }
  69. template<class T>
  70. void lock(T const*)
  71. {
  72. BOOST_ASSERT(id_ == 0);
  73. id_ = T::id;
  74. }
  75. template<class T>
  76. void unlock(T const*)
  77. {
  78. BOOST_ASSERT(id_ == T::id);
  79. id_ = 0;
  80. }
  81. template<class T>
  82. bool try_lock(T const*)
  83. {
  84. // If this assert goes off it means you are attempting to
  85. // simultaneously initiate more than one of same asynchronous
  86. // operation, which is not allowed. For example, you must wait
  87. // for an async_read to complete before performing another
  88. // async_read.
  89. //
  90. BOOST_ASSERT(id_ != T::id);
  91. if(id_ != 0)
  92. return false;
  93. id_ = T::id;
  94. return true;
  95. }
  96. template<class T>
  97. bool try_unlock(T const*)
  98. {
  99. if(id_ != T::id)
  100. return false;
  101. id_ = 0;
  102. return true;
  103. }
  104. };
  105. //------------------------------------------------------------------------------
  106. struct stream_prng
  107. {
  108. bool secure_prng_ = true;
  109. struct prng_type
  110. {
  111. std::minstd_rand fast;
  112. beast::detail::chacha<20> secure;
  113. #if BOOST_BEAST_NO_THREAD_LOCAL
  114. prng_type* next = nullptr;
  115. #endif
  116. prng_type(std::uint32_t const* v, std::uint64_t stream)
  117. : fast(static_cast<typename decltype(fast)::result_type>(
  118. v[0] + v[1] + v[2] + v[3] + v[4] + v[5] + v[6] + v[7] + stream))
  119. , secure(v, stream)
  120. {
  121. }
  122. };
  123. class prng_ref
  124. {
  125. prng_type* p_;
  126. public:
  127. prng_ref& operator=(prng_ref&&) = delete;
  128. explicit
  129. prng_ref(prng_type& p)
  130. : p_(&p)
  131. {
  132. }
  133. prng_ref(prng_ref&& other)
  134. : p_(boost::exchange(other.p_, nullptr))
  135. {
  136. }
  137. #ifdef BOOST_BEAST_NO_THREAD_LOCAL
  138. ~prng_ref()
  139. {
  140. if(p_)
  141. pool::impl().release(*p_);
  142. }
  143. #endif
  144. prng_type*
  145. operator->() const
  146. {
  147. return p_;
  148. }
  149. };
  150. #ifndef BOOST_BEAST_NO_THREAD_LOCAL
  151. static
  152. prng_ref
  153. prng()
  154. {
  155. static std::atomic<std::uint64_t> stream{0};
  156. thread_local prng_type p{seed(), stream++};
  157. return prng_ref(p);
  158. }
  159. #else
  160. static
  161. prng_ref
  162. prng()
  163. {
  164. return prng_ref(pool::impl().acquire());
  165. }
  166. #endif
  167. static
  168. std::uint32_t const*
  169. seed(std::seed_seq* ss = nullptr)
  170. {
  171. static seed_data d(ss);
  172. return d.v;
  173. }
  174. std::uint32_t
  175. create_mask()
  176. {
  177. auto p = prng();
  178. if(secure_prng_)
  179. for(;;)
  180. if(auto key = p->secure())
  181. return key;
  182. for(;;)
  183. if(auto key = p->fast())
  184. return key;
  185. }
  186. private:
  187. struct seed_data
  188. {
  189. std::uint32_t v[8];
  190. explicit
  191. seed_data(std::seed_seq* pss)
  192. {
  193. if(! pss)
  194. {
  195. std::random_device g;
  196. std::seed_seq ss{
  197. g(), g(), g(), g(), g(), g(), g(), g()};
  198. ss.generate(v, v+8);
  199. }
  200. else
  201. {
  202. pss->generate(v, v+8);
  203. }
  204. }
  205. };
  206. #ifdef BOOST_BEAST_NO_THREAD_LOCAL
  207. class pool
  208. {
  209. prng_type* head_ = nullptr;
  210. std::atomic<std::uint64_t> n_{0};
  211. std::mutex m_;
  212. public:
  213. ~pool()
  214. {
  215. for(auto p = head_; p;)
  216. {
  217. auto next = p->next;
  218. p->~prng_type();
  219. boost::alignment::aligned_free(p);
  220. p = next;
  221. }
  222. }
  223. prng_type&
  224. acquire()
  225. {
  226. for(;;)
  227. {
  228. std::lock_guard<std::mutex> lock(m_);
  229. if(! head_)
  230. break;
  231. auto p = head_;
  232. head_ = head_->next;
  233. return *p;
  234. }
  235. auto p = boost::alignment::aligned_alloc(
  236. 16, sizeof(prng_type));
  237. if(! p)
  238. BOOST_THROW_EXCEPTION(std::bad_alloc{});
  239. return *(new(p) prng_type(seed(), n_++));
  240. }
  241. void
  242. release(prng_type& p)
  243. {
  244. std::lock_guard<std::mutex> lock(m_);
  245. p.next = head_;
  246. head_ = &p;
  247. }
  248. static
  249. pool&
  250. impl()
  251. {
  252. static pool instance;
  253. return instance;
  254. }
  255. };
  256. #endif
  257. };
  258. //------------------------------------------------------------------------------
  259. template<bool deflateSupported>
  260. struct stream_base : stream_prng
  261. {
  262. // State information for the permessage-deflate extension
  263. struct pmd_type
  264. {
  265. // `true` if current read message is compressed
  266. bool rd_set = false;
  267. zlib::deflate_stream zo;
  268. zlib::inflate_stream zi;
  269. };
  270. std::unique_ptr<pmd_type> pmd_; // pmd settings or nullptr
  271. permessage_deflate pmd_opts_; // local pmd options
  272. detail::pmd_offer pmd_config_; // offer (client) or negotiation (server)
  273. // return `true` if current message is deflated
  274. bool
  275. rd_deflated() const
  276. {
  277. return pmd_ && pmd_->rd_set;
  278. }
  279. // set whether current message is deflated
  280. // returns `false` on protocol violation
  281. bool
  282. rd_deflated(bool rsv1)
  283. {
  284. if(pmd_)
  285. {
  286. pmd_->rd_set = rsv1;
  287. return true;
  288. }
  289. return ! rsv1; // pmd not negotiated
  290. }
  291. template<class ConstBufferSequence>
  292. bool
  293. deflate(
  294. boost::asio::mutable_buffer& out,
  295. buffers_suffix<ConstBufferSequence>& cb,
  296. bool fin,
  297. std::size_t& total_in,
  298. error_code& ec);
  299. void
  300. do_context_takeover_write(role_type role);
  301. void
  302. inflate(
  303. zlib::z_params& zs,
  304. zlib::Flush flush,
  305. error_code& ec);
  306. void
  307. do_context_takeover_read(role_type role);
  308. };
  309. template<>
  310. struct stream_base<false> : stream_prng
  311. {
  312. // These stubs are for avoiding linking in the zlib
  313. // code when permessage-deflate is not enabled.
  314. bool
  315. rd_deflated() const
  316. {
  317. return false;
  318. }
  319. bool
  320. rd_deflated(bool rsv1)
  321. {
  322. return ! rsv1;
  323. }
  324. template<class ConstBufferSequence>
  325. bool
  326. deflate(
  327. boost::asio::mutable_buffer&,
  328. buffers_suffix<ConstBufferSequence>&,
  329. bool,
  330. std::size_t&,
  331. error_code&)
  332. {
  333. return false;
  334. }
  335. void
  336. do_context_takeover_write(role_type)
  337. {
  338. }
  339. void
  340. inflate(
  341. zlib::z_params&,
  342. zlib::Flush,
  343. error_code&)
  344. {
  345. }
  346. void
  347. do_context_takeover_read(role_type)
  348. {
  349. }
  350. };
  351. } // detail
  352. } // websocket
  353. } // beast
  354. } // boost
  355. #endif