socket_channel_handler.c 15 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434
  1. /**
  2. * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
  3. * SPDX-License-Identifier: Apache-2.0.
  4. */
  5. #include <aws/io/socket_channel_handler.h>
  6. #include <aws/common/error.h>
  7. #include <aws/common/task_scheduler.h>
  8. #include <aws/io/event_loop.h>
  9. #include <aws/io/logging.h>
  10. #include <aws/io/socket.h>
  11. #include <aws/io/statistics.h>
  12. #ifdef _MSC_VER
  13. # pragma warning(disable : 4204) /* non-constant aggregate initializer */
  14. #endif
  15. struct socket_handler {
  16. struct aws_socket *socket;
  17. struct aws_channel_slot *slot;
  18. size_t max_rw_size;
  19. struct aws_channel_task read_task_storage;
  20. struct aws_channel_task shutdown_task_storage;
  21. struct aws_crt_statistics_socket stats;
  22. int shutdown_err_code;
  23. bool shutdown_in_progress;
  24. };
  25. static int s_socket_process_read_message(
  26. struct aws_channel_handler *handler,
  27. struct aws_channel_slot *slot,
  28. struct aws_io_message *message) {
  29. (void)handler;
  30. (void)slot;
  31. (void)message;
  32. AWS_LOGF_FATAL(
  33. AWS_LS_IO_SOCKET_HANDLER,
  34. "id=%p: process_read_message called on "
  35. "socket handler. This should never happen",
  36. (void *)handler);
  37. /*since a socket handler will ALWAYS be the first handler in a channel,
  38. * this should NEVER happen, if it does it's a programmer error.*/
  39. AWS_ASSERT(0);
  40. return aws_raise_error(AWS_IO_CHANNEL_ERROR_ERROR_CANT_ACCEPT_INPUT);
  41. }
  42. /* invoked by the socket when a write has completed or failed. */
  43. static void s_on_socket_write_complete(
  44. struct aws_socket *socket,
  45. int error_code,
  46. size_t amount_written,
  47. void *user_data) {
  48. if (user_data) {
  49. struct aws_io_message *message = user_data;
  50. struct aws_channel *channel = message->owning_channel;
  51. AWS_LOGF_TRACE(
  52. AWS_LS_IO_SOCKET_HANDLER,
  53. "static: write of size %llu, completed on channel %p",
  54. (unsigned long long)amount_written,
  55. (void *)channel);
  56. if (message->on_completion) {
  57. message->on_completion(channel, message, error_code, message->user_data);
  58. }
  59. if (socket && socket->handler) {
  60. struct socket_handler *socket_handler = socket->handler->impl;
  61. socket_handler->stats.bytes_written += amount_written;
  62. }
  63. aws_mem_release(message->allocator, message);
  64. if (error_code) {
  65. aws_channel_shutdown(channel, error_code);
  66. }
  67. }
  68. }
  69. static int s_socket_process_write_message(
  70. struct aws_channel_handler *handler,
  71. struct aws_channel_slot *slot,
  72. struct aws_io_message *message) {
  73. (void)slot;
  74. struct socket_handler *socket_handler = handler->impl;
  75. AWS_LOGF_TRACE(
  76. AWS_LS_IO_SOCKET_HANDLER,
  77. "id=%p: writing message of size %llu",
  78. (void *)handler,
  79. (unsigned long long)message->message_data.len);
  80. if (!aws_socket_is_open(socket_handler->socket)) {
  81. return aws_raise_error(AWS_IO_SOCKET_CLOSED);
  82. }
  83. struct aws_byte_cursor cursor = aws_byte_cursor_from_buf(&message->message_data);
  84. if (aws_socket_write(socket_handler->socket, &cursor, s_on_socket_write_complete, message)) {
  85. return AWS_OP_ERR;
  86. }
  87. return AWS_OP_SUCCESS;
  88. }
  89. static void s_read_task(struct aws_channel_task *task, void *arg, aws_task_status status);
  90. static void s_on_readable_notification(struct aws_socket *socket, int error_code, void *user_data);
  91. /* Ok this next function is VERY important for how back pressure works. Here's what it's supposed to be doing:
  92. *
  93. * See how much data downstream is willing to accept.
  94. * See how much we're actually willing to read per event loop tick (usually 16 kb).
  95. * Take the minimum of those two.
  96. * Try and read as much as possible up to the calculated max read.
  97. * If we didn't read up to the max_read, we go back to waiting on the event loop to tell us we can read more.
  98. * If we did read up to the max_read, we stop reading immediately and wait for either for a window update,
  99. * or schedule a task to enforce fairness for other sockets in the event loop if we read up to the max
  100. * read per event loop tick.
  101. */
  102. static void s_do_read(struct socket_handler *socket_handler) {
  103. size_t downstream_window = aws_channel_slot_downstream_read_window(socket_handler->slot);
  104. size_t max_to_read =
  105. downstream_window > socket_handler->max_rw_size ? socket_handler->max_rw_size : downstream_window;
  106. AWS_LOGF_TRACE(
  107. AWS_LS_IO_SOCKET_HANDLER,
  108. "id=%p: invoking read. Downstream window %llu, max_to_read %llu",
  109. (void *)socket_handler->slot->handler,
  110. (unsigned long long)downstream_window,
  111. (unsigned long long)max_to_read);
  112. if (max_to_read == 0) {
  113. return;
  114. }
  115. size_t total_read = 0;
  116. size_t read = 0;
  117. while (total_read < max_to_read && !socket_handler->shutdown_in_progress) {
  118. size_t iter_max_read = max_to_read - total_read;
  119. struct aws_io_message *message = aws_channel_acquire_message_from_pool(
  120. socket_handler->slot->channel, AWS_IO_MESSAGE_APPLICATION_DATA, iter_max_read);
  121. if (!message) {
  122. break;
  123. }
  124. if (aws_socket_read(socket_handler->socket, &message->message_data, &read)) {
  125. aws_mem_release(message->allocator, message);
  126. break;
  127. }
  128. total_read += read;
  129. AWS_LOGF_TRACE(
  130. AWS_LS_IO_SOCKET_HANDLER,
  131. "id=%p: read %llu from socket",
  132. (void *)socket_handler->slot->handler,
  133. (unsigned long long)read);
  134. if (aws_channel_slot_send_message(socket_handler->slot, message, AWS_CHANNEL_DIR_READ)) {
  135. aws_mem_release(message->allocator, message);
  136. break;
  137. }
  138. }
  139. AWS_LOGF_TRACE(
  140. AWS_LS_IO_SOCKET_HANDLER,
  141. "id=%p: total read on this tick %llu",
  142. (void *)&socket_handler->slot->handler,
  143. (unsigned long long)total_read);
  144. socket_handler->stats.bytes_read += total_read;
  145. /* resubscribe as long as there's no error, just return if we're in a would block scenario. */
  146. if (total_read < max_to_read) {
  147. int last_error = aws_last_error();
  148. if (last_error != AWS_IO_READ_WOULD_BLOCK && !socket_handler->shutdown_in_progress) {
  149. aws_channel_shutdown(socket_handler->slot->channel, last_error);
  150. }
  151. AWS_LOGF_TRACE(
  152. AWS_LS_IO_SOCKET_HANDLER,
  153. "id=%p: out of data to read on socket. "
  154. "Waiting on event-loop notification.",
  155. (void *)socket_handler->slot->handler);
  156. return;
  157. }
  158. /* in this case, everything was fine, but there's still pending reads. We need to schedule a task to do the read
  159. * again. */
  160. if (!socket_handler->shutdown_in_progress && total_read == socket_handler->max_rw_size &&
  161. !socket_handler->read_task_storage.task_fn) {
  162. AWS_LOGF_TRACE(
  163. AWS_LS_IO_SOCKET_HANDLER,
  164. "id=%p: more data is pending read, but we've exceeded "
  165. "the max read on this tick. Scheduling a task to read on next tick.",
  166. (void *)socket_handler->slot->handler);
  167. aws_channel_task_init(
  168. &socket_handler->read_task_storage, s_read_task, socket_handler, "socket_handler_re_read");
  169. aws_channel_schedule_task_now(socket_handler->slot->channel, &socket_handler->read_task_storage);
  170. }
  171. }
  172. /* the socket is either readable or errored out. If it's readable, kick off s_do_read() to do its thing.
  173. * If an error, start the channel shutdown process. */
  174. static void s_on_readable_notification(struct aws_socket *socket, int error_code, void *user_data) {
  175. (void)socket;
  176. struct socket_handler *socket_handler = user_data;
  177. AWS_LOGF_TRACE(AWS_LS_IO_SOCKET_HANDLER, "id=%p: socket is now readable", (void *)socket_handler->slot->handler);
  178. /* read regardless so we can pick up data that was sent prior to the close. For example, peer sends a TLS ALERT
  179. * then immediately closes the socket. On some platforms, we'll never see the readable flag. So we want to make
  180. * sure we read the ALERT, otherwise, we'll end up telling the user that the channel shutdown because of a socket
  181. * closure, when in reality it was a TLS error */
  182. s_do_read(socket_handler);
  183. if (error_code && !socket_handler->shutdown_in_progress) {
  184. aws_channel_shutdown(socket_handler->slot->channel, error_code);
  185. }
  186. }
  187. /* Either the result of a context switch (for fairness in the event loop), or a window update. */
  188. static void s_read_task(struct aws_channel_task *task, void *arg, aws_task_status status) {
  189. task->task_fn = NULL;
  190. task->arg = NULL;
  191. if (status == AWS_TASK_STATUS_RUN_READY) {
  192. struct socket_handler *socket_handler = arg;
  193. s_do_read(socket_handler);
  194. }
  195. }
  196. static int s_socket_increment_read_window(
  197. struct aws_channel_handler *handler,
  198. struct aws_channel_slot *slot,
  199. size_t size) {
  200. (void)size;
  201. struct socket_handler *socket_handler = handler->impl;
  202. if (!socket_handler->shutdown_in_progress && !socket_handler->read_task_storage.task_fn) {
  203. AWS_LOGF_TRACE(
  204. AWS_LS_IO_SOCKET_HANDLER,
  205. "id=%p: increment read window message received, scheduling"
  206. " task for another read operation.",
  207. (void *)handler);
  208. aws_channel_task_init(
  209. &socket_handler->read_task_storage, s_read_task, socket_handler, "socket_handler_read_on_window_increment");
  210. aws_channel_schedule_task_now(slot->channel, &socket_handler->read_task_storage);
  211. }
  212. return AWS_OP_SUCCESS;
  213. }
  214. static void s_close_task(struct aws_channel_task *task, void *arg, aws_task_status status) {
  215. (void)task;
  216. (void)status;
  217. struct aws_channel_handler *handler = arg;
  218. struct socket_handler *socket_handler = handler->impl;
  219. /*
  220. * Run this unconditionally regardless of status, otherwise channel will not
  221. * finish shutting down properly
  222. */
  223. /* this only happens in write direction. */
  224. /* we also don't care about the free_scarce_resource_immediately
  225. * code since we're always the last one in the shutdown sequence. */
  226. aws_channel_slot_on_handler_shutdown_complete(
  227. socket_handler->slot, AWS_CHANNEL_DIR_WRITE, socket_handler->shutdown_err_code, false);
  228. }
  229. static int s_socket_shutdown(
  230. struct aws_channel_handler *handler,
  231. struct aws_channel_slot *slot,
  232. enum aws_channel_direction dir,
  233. int error_code,
  234. bool free_scarce_resource_immediately) {
  235. struct socket_handler *socket_handler = (struct socket_handler *)handler->impl;
  236. socket_handler->shutdown_in_progress = true;
  237. if (dir == AWS_CHANNEL_DIR_READ) {
  238. AWS_LOGF_TRACE(
  239. AWS_LS_IO_SOCKET_HANDLER,
  240. "id=%p: shutting down read direction with error_code %d",
  241. (void *)handler,
  242. error_code);
  243. if (free_scarce_resource_immediately && aws_socket_is_open(socket_handler->socket)) {
  244. if (aws_socket_close(socket_handler->socket)) {
  245. return AWS_OP_ERR;
  246. }
  247. }
  248. return aws_channel_slot_on_handler_shutdown_complete(slot, dir, error_code, free_scarce_resource_immediately);
  249. }
  250. AWS_LOGF_TRACE(
  251. AWS_LS_IO_SOCKET_HANDLER,
  252. "id=%p: shutting down write direction with error_code %d",
  253. (void *)handler,
  254. error_code);
  255. if (aws_socket_is_open(socket_handler->socket)) {
  256. aws_socket_close(socket_handler->socket);
  257. }
  258. /* Schedule a task to complete the shutdown, in case a do_read task is currently pending.
  259. * It's OK to delay the shutdown, even when free_scarce_resources_immediately is true,
  260. * because the socket has been closed: mitigating the risk that the socket is still being abused by
  261. * a hostile peer. */
  262. aws_channel_task_init(&socket_handler->shutdown_task_storage, s_close_task, handler, "socket_handler_close");
  263. socket_handler->shutdown_err_code = error_code;
  264. aws_channel_schedule_task_now(slot->channel, &socket_handler->shutdown_task_storage);
  265. return AWS_OP_SUCCESS;
  266. }
  267. static size_t s_message_overhead(struct aws_channel_handler *handler) {
  268. (void)handler;
  269. return 0;
  270. }
  271. static size_t s_socket_initial_window_size(struct aws_channel_handler *handler) {
  272. (void)handler;
  273. return SIZE_MAX;
  274. }
  275. static void s_socket_destroy(struct aws_channel_handler *handler) {
  276. if (handler != NULL) {
  277. struct socket_handler *socket_handler = (struct socket_handler *)handler->impl;
  278. if (socket_handler != NULL) {
  279. aws_crt_statistics_socket_cleanup(&socket_handler->stats);
  280. }
  281. aws_mem_release(handler->alloc, handler);
  282. }
  283. }
  284. static void s_reset_statistics(struct aws_channel_handler *handler) {
  285. struct socket_handler *socket_handler = (struct socket_handler *)handler->impl;
  286. aws_crt_statistics_socket_reset(&socket_handler->stats);
  287. }
  288. static void s_gather_statistics(struct aws_channel_handler *handler, struct aws_array_list *stats_list) {
  289. struct socket_handler *socket_handler = (struct socket_handler *)handler->impl;
  290. void *stats_base = &socket_handler->stats;
  291. aws_array_list_push_back(stats_list, &stats_base);
  292. }
  293. static void s_trigger_read(struct aws_channel_handler *handler) {
  294. struct socket_handler *socket_handler = (struct socket_handler *)handler->impl;
  295. s_do_read(socket_handler);
  296. }
  297. static struct aws_channel_handler_vtable s_vtable = {
  298. .process_read_message = s_socket_process_read_message,
  299. .destroy = s_socket_destroy,
  300. .process_write_message = s_socket_process_write_message,
  301. .initial_window_size = s_socket_initial_window_size,
  302. .increment_read_window = s_socket_increment_read_window,
  303. .shutdown = s_socket_shutdown,
  304. .message_overhead = s_message_overhead,
  305. .reset_statistics = s_reset_statistics,
  306. .gather_statistics = s_gather_statistics,
  307. .trigger_read = s_trigger_read,
  308. };
  309. struct aws_channel_handler *aws_socket_handler_new(
  310. struct aws_allocator *allocator,
  311. struct aws_socket *socket,
  312. struct aws_channel_slot *slot,
  313. size_t max_read_size) {
  314. /* make sure something has assigned this socket to an event loop, in client mode this will already have occurred.
  315. In server mode, someone should have assigned it before calling us.*/
  316. AWS_ASSERT(aws_socket_get_event_loop(socket));
  317. struct aws_channel_handler *handler = NULL;
  318. struct socket_handler *impl = NULL;
  319. if (!aws_mem_acquire_many(
  320. allocator, 2, &handler, sizeof(struct aws_channel_handler), &impl, sizeof(struct socket_handler))) {
  321. return NULL;
  322. }
  323. impl->socket = socket;
  324. impl->slot = slot;
  325. impl->max_rw_size = max_read_size;
  326. AWS_ZERO_STRUCT(impl->read_task_storage);
  327. AWS_ZERO_STRUCT(impl->shutdown_task_storage);
  328. impl->shutdown_in_progress = false;
  329. if (aws_crt_statistics_socket_init(&impl->stats)) {
  330. goto cleanup_handler;
  331. }
  332. AWS_LOGF_DEBUG(
  333. AWS_LS_IO_SOCKET_HANDLER,
  334. "id=%p: Socket handler created with max_read_size of %llu",
  335. (void *)handler,
  336. (unsigned long long)max_read_size);
  337. handler->alloc = allocator;
  338. handler->impl = impl;
  339. handler->vtable = &s_vtable;
  340. handler->slot = slot;
  341. if (aws_socket_subscribe_to_readable_events(socket, s_on_readable_notification, impl)) {
  342. goto cleanup_handler;
  343. }
  344. socket->handler = handler;
  345. return handler;
  346. cleanup_handler:
  347. aws_mem_release(allocator, handler);
  348. return NULL;
  349. }
  350. const struct aws_socket *aws_socket_handler_get_socket(const struct aws_channel_handler *handler) {
  351. AWS_PRECONDITION(handler);
  352. const struct socket_handler *socket_handler = handler->impl;
  353. return socket_handler->socket;
  354. }