|
| 1 | +#include "sockets.h" |
| 2 | + |
| 3 | +#include <../wasi-0.2.0/handles.h> |
| 4 | + |
| 5 | +template <> struct HandleOps<host_api::TCPSocket> { |
| 6 | + using owned = wasi_sockets_tcp_own_tcp_socket_t; |
| 7 | + using borrowed = wasi_sockets_tcp_borrow_tcp_socket_t; |
| 8 | +}; |
| 9 | + |
| 10 | +class TCPSocketHandle final : public WASIHandle<host_api::TCPSocket> { |
| 11 | + wasi_sockets_instance_network_own_network_t network_; |
| 12 | + PollableHandle pollable_handle_; |
| 13 | + wasi_io_streams_own_input_stream_t input_; |
| 14 | + wasi_io_streams_own_output_stream_t output_; |
| 15 | + |
| 16 | + friend host_api::TCPSocket; |
| 17 | + |
| 18 | +public: |
| 19 | + explicit TCPSocketHandle(HandleOps<host_api::TCPSocket>::owned handle) |
| 20 | + : WASIHandle(handle), pollable_handle_(INVALID_POLLABLE_HANDLE) { |
| 21 | + network_ = wasi_sockets_instance_network_instance_network(); |
| 22 | + } |
| 23 | + |
| 24 | + static TCPSocketHandle *cast(HandleState *handle) { |
| 25 | + return reinterpret_cast<TCPSocketHandle *>(handle); |
| 26 | + } |
| 27 | + |
| 28 | + wasi_sockets_tcp_borrow_network_t network() const { |
| 29 | + return wasi_sockets_tcp_borrow_network_t(network_.__handle); |
| 30 | + } |
| 31 | + PollableHandle pollable_handle() { |
| 32 | + if (pollable_handle_ == INVALID_POLLABLE_HANDLE) { |
| 33 | + pollable_handle_ = wasi_sockets_tcp_method_tcp_socket_subscribe(borrow()).__handle; |
| 34 | + } |
| 35 | + return pollable_handle_; |
| 36 | + } |
| 37 | +}; |
| 38 | + |
| 39 | +namespace host_api { |
| 40 | + |
| 41 | +TCPSocket::TCPSocket(std::unique_ptr<HandleState> state) { |
| 42 | + this->handle_state_ = std::move(state); |
| 43 | +} |
| 44 | + |
| 45 | +TCPSocket *TCPSocket::make(IPAddressFamily address_family) { |
| 46 | + wasi_sockets_tcp_create_socket_ip_address_family_t family = |
| 47 | + address_family == IPV4 ? WASI_SOCKETS_NETWORK_IP_ADDRESS_FAMILY_IPV4 |
| 48 | + : WASI_SOCKETS_NETWORK_IP_ADDRESS_FAMILY_IPV6; |
| 49 | + wasi_sockets_tcp_create_socket_own_tcp_socket_t ret; |
| 50 | + wasi_sockets_tcp_create_socket_error_code_t err; |
| 51 | + if (!wasi_sockets_tcp_create_socket_create_tcp_socket(family, &ret, &err)) { |
| 52 | + return nullptr; |
| 53 | + } |
| 54 | + return new TCPSocket(std::unique_ptr<HandleState>(new TCPSocketHandle(ret))); |
| 55 | +} |
| 56 | + |
| 57 | +bool TCPSocket::connect(AddressIPV4 address, Port port) { |
| 58 | + auto state = TCPSocketHandle::cast(handle_state_.get()); |
| 59 | + auto handle = state->borrow(); |
| 60 | + auto addr = wasi_sockets_network_ipv4_address_t(get<0>(address), get<1>(address), get<2>(address), |
| 61 | + get<3>(address)); |
| 62 | + wasi_sockets_tcp_ip_socket_address_t socket_address = { |
| 63 | + WASI_SOCKETS_NETWORK_IP_SOCKET_ADDRESS_IPV4, {{port, addr}}}; |
| 64 | + wasi_sockets_tcp_error_code_t err; |
| 65 | + if (!wasi_sockets_tcp_method_tcp_socket_start_connect(handle, state->network(), &socket_address, |
| 66 | + &err)) { |
| 67 | + // TODO: handle error |
| 68 | + return false; |
| 69 | + } |
| 70 | + |
| 71 | + wasi_sockets_tcp_tuple2_own_input_stream_own_output_stream_t streams; |
| 72 | + while (true) { |
| 73 | + if (!wasi_sockets_tcp_method_tcp_socket_finish_connect(handle, &streams, &err)) { |
| 74 | + if (err == WASI_SOCKETS_NETWORK_ERROR_CODE_WOULD_BLOCK) { |
| 75 | + block_on_pollable_handle(state->pollable_handle()); |
| 76 | + continue; |
| 77 | + } |
| 78 | + // TODO: handle error |
| 79 | + return false; |
| 80 | + } |
| 81 | + state->input_ = streams.f0; |
| 82 | + state->output_ = streams.f1; |
| 83 | + break; |
| 84 | + } |
| 85 | + |
| 86 | + return true; |
| 87 | +} |
| 88 | +void TCPSocket::close() { |
| 89 | + auto state = TCPSocketHandle::cast(handle_state_.get()); |
| 90 | + if (!state->valid()) { |
| 91 | + return; |
| 92 | + } |
| 93 | + wasi_sockets_tcp_error_code_t err; |
| 94 | + wasi_sockets_tcp_method_tcp_socket_shutdown(state->borrow(), |
| 95 | + WASI_SOCKETS_TCP_SHUTDOWN_TYPE_BOTH, &err); |
| 96 | + wasi_io_streams_output_stream_drop_own(state->output_); |
| 97 | + wasi_io_streams_input_stream_drop_own(state->input_); |
| 98 | + if (state->pollable_handle_ != INVALID_POLLABLE_HANDLE) { |
| 99 | + wasi_io_poll_pollable_drop_own(own_pollable_t{state->pollable_handle_}); |
| 100 | + } |
| 101 | + wasi_sockets_tcp_tcp_socket_drop_own(state->take()); |
| 102 | +} |
| 103 | + |
| 104 | +bool TCPSocket::send(HostString chunk) { |
| 105 | + auto state = TCPSocketHandle::cast(handle_state_.get()); |
| 106 | + auto borrow = wasi_io_streams_borrow_output_stream(state->output_); |
| 107 | + bindings_list_u8_t list{reinterpret_cast<uint8_t *>(chunk.ptr.get()), chunk.len}; |
| 108 | + uint64_t capacity = 0; |
| 109 | + wasi_io_streams_stream_error_t err; |
| 110 | + if (!wasi_io_streams_method_output_stream_check_write(borrow, &capacity, &err)) { |
| 111 | + // TODO: proper error handling. |
| 112 | + } |
| 113 | + // TODO: proper error handling. |
| 114 | + MOZ_ASSERT(chunk.len <= capacity); |
| 115 | + return wasi_io_streams_method_output_stream_write(borrow, &list, &err); |
| 116 | +} |
| 117 | + |
| 118 | +HostString TCPSocket::receive(uint32_t chunk_size) { |
| 119 | + auto state = TCPSocketHandle::cast(handle_state_.get()); |
| 120 | + auto borrow = wasi_io_streams_borrow_input_stream(state->input_); |
| 121 | + bindings_list_u8_t ret{}; |
| 122 | + wasi_io_streams_stream_error_t err{}; |
| 123 | + mozilla::DebugOnly<bool> success; |
| 124 | + success = wasi_io_streams_method_input_stream_blocking_read(borrow, chunk_size, &ret, &err); |
| 125 | + MOZ_ASSERT(success, "Why you not handle errors"); |
| 126 | + UniqueChars chars((char*)ret.ptr); |
| 127 | + return HostString(std::move(chars), ret.len); |
| 128 | +} |
| 129 | + |
| 130 | +} // namespace host_api |
0 commit comments