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 #ifndef BOOST_BEAST_WEBSOCKET_IMPL_TEARDOWN_HPP
11 #define BOOST_BEAST_WEBSOCKET_IMPL_TEARDOWN_HPP
12 
13 #include <boost/beast/core/async_base.hpp>
14 #include <boost/beast/core/bind_handler.hpp>
15 #include <boost/beast/core/stream_traits.hpp>
16 #include <boost/beast/core/detail/bind_continuation.hpp>
17 #include <boost/beast/core/detail/is_invocable.hpp>
18 #include <boost/asio/coroutine.hpp>
19 #include <boost/asio/post.hpp>
20 #include <memory>
21 
22 namespace boost {
23 namespace beast {
24 namespace websocket {
25 
26 namespace detail {
27 
28 template<
29     class Protocol, class Executor,
30     class Handler>
31 class teardown_tcp_op
32     : public beast::async_base<
33         Handler, beast::executor_type<
34             net::basic_stream_socket<
35                 Protocol, Executor>>>
36     , public asio::coroutine
37 {
38     using socket_type =
39         net::basic_stream_socket<Protocol, Executor>;
40 
41     socket_type& s_;
42     role_type role_;
43     bool nb_;
44 
45 public:
46     template<class Handler_>
teardown_tcp_op(Handler_ && h,socket_type & s,role_type role)47     teardown_tcp_op(
48         Handler_&& h,
49         socket_type& s,
50         role_type role)
51         : async_base<Handler,
52             beast::executor_type<
53                 net::basic_stream_socket<
54                     Protocol, Executor>>>(
55             std::forward<Handler_>(h),
56             s.get_executor())
57         , s_(s)
58         , role_(role)
59         , nb_(false)
60     {
61         (*this)({}, 0, false);
62     }
63 
64     void
operator ()(error_code ec={},std::size_t bytes_transferred=0,bool cont=true)65     operator()(
66         error_code ec = {},
67         std::size_t bytes_transferred = 0,
68         bool cont = true)
69     {
70         BOOST_ASIO_CORO_REENTER(*this)
71         {
72             nb_ = s_.non_blocking();
73             s_.non_blocking(true, ec);
74             if(ec)
75                 goto upcall;
76             if(role_ == role_type::server)
77                 s_.shutdown(net::socket_base::shutdown_send, ec);
78             if(ec)
79                 goto upcall;
80             for(;;)
81             {
82                 {
83                     char buf[2048];
84                     s_.read_some(net::buffer(buf), ec);
85                 }
86                 if(ec == net::error::would_block)
87                 {
88                     BOOST_ASIO_CORO_YIELD
89                     {
90                         BOOST_ASIO_HANDLER_LOCATION((
91                             __FILE__, __LINE__,
92                             "websocket::tcp::async_teardown"
93                         ));
94 
95                         s_.async_wait(
96                             net::socket_base::wait_read,
97                                 beast::detail::bind_continuation(std::move(*this)));
98                     }
99                     continue;
100                 }
101                 if(ec)
102                 {
103                     if(ec != net::error::eof)
104                         goto upcall;
105                     ec = {};
106                     break;
107                 }
108                 if(bytes_transferred == 0)
109                 {
110                     // happens sometimes
111                     // https://github.com/boostorg/beast/issues/1373
112                     break;
113                 }
114             }
115             if(role_ == role_type::client)
116                 s_.shutdown(net::socket_base::shutdown_send, ec);
117             if(ec)
118                 goto upcall;
119             s_.close(ec);
120         upcall:
121             if(! cont)
122             {
123                 BOOST_ASIO_CORO_YIELD
124                 {
125                     BOOST_ASIO_HANDLER_LOCATION((
126                         __FILE__, __LINE__,
127                         "websocket::tcp::async_teardown"
128                         ));
129 
130                     net::post(bind_front_handler(
131                         std::move(*this), ec));
132                 }
133             }
134             {
135                 error_code ignored;
136                 s_.non_blocking(nb_, ignored);
137             }
138             this->complete_now(ec);
139         }
140     }
141 };
142 
143 } // detail
144 
145 //------------------------------------------------------------------------------
146 
147 template<class Protocol, class Executor>
148 void
teardown(role_type role,net::basic_stream_socket<Protocol,Executor> & socket,error_code & ec)149 teardown(
150     role_type role,
151     net::basic_stream_socket<
152         Protocol, Executor>& socket,
153     error_code& ec)
154 {
155     if(role == role_type::server)
156         socket.shutdown(
157             net::socket_base::shutdown_send, ec);
158     if(ec)
159         return;
160     for(;;)
161     {
162         char buf[2048];
163         auto const bytes_transferred =
164             socket.read_some(net::buffer(buf), ec);
165         if(ec)
166         {
167             if(ec != net::error::eof)
168                 return;
169             ec = {};
170             break;
171         }
172         if(bytes_transferred == 0)
173         {
174             // happens sometimes
175             // https://github.com/boostorg/beast/issues/1373
176             break;
177         }
178     }
179     if(role == role_type::client)
180         socket.shutdown(
181             net::socket_base::shutdown_send, ec);
182     if(ec)
183         return;
184     socket.close(ec);
185 }
186 
187 template<
188     class Protocol, class Executor,
189     class TeardownHandler>
190 void
async_teardown(role_type role,net::basic_stream_socket<Protocol,Executor> & socket,TeardownHandler && handler)191 async_teardown(
192     role_type role,
193     net::basic_stream_socket<
194         Protocol, Executor>& socket,
195     TeardownHandler&& handler)
196 {
197     static_assert(beast::detail::is_invocable<
198         TeardownHandler, void(error_code)>::value,
199             "TeardownHandler type requirements not met");
200     detail::teardown_tcp_op<
201         Protocol,
202         Executor,
203         typename std::decay<TeardownHandler>::type>(
204             std::forward<TeardownHandler>(handler),
205             socket,
206             role);
207 }
208 
209 } // websocket
210 } // beast
211 } // boost
212 
213 #endif
214