diff --git a/src/ssh_transport.c b/src/ssh_transport.c index 9a1a0cf..d9ef097 100644 --- a/src/ssh_transport.c +++ b/src/ssh_transport.c @@ -36,6 +36,7 @@ #define SSH_TRANSPORT_MAX_PENDING_HANDSHAKES SSH_TRANSPORT_MAX_SESSIONS #define SSH_TRANSPORT_MAX_AUTH_ATTEMPTS 3U #define SSH_TRANSPORT_GENERATION_MAX 0x3fffffffU +#define SSH_TRANSPORT_WOLFSSH_READ_BUDGET 2048U static const char *TAG = "ssh_transport"; @@ -49,6 +50,7 @@ typedef struct { bool authenticated; bool shell_requested; uint8_t authentication_attempts; + word32 io_read_budget; bool writer; bool close_requested; int64_t handshake_deadline_us; @@ -168,6 +170,45 @@ static int ssh_seed(OS_Seed *seed, byte *output, word32 size) return secure_random_fill(output, size) == ESP_OK ? 0 : -1; } +static int bounded_ssh_receive(WOLFSSH *ssh, void *data, word32 size, + void *context) +{ + (void)ssh; + ssh_slot_t *slot = (ssh_slot_t *)context; + if (slot == NULL || data == NULL || slot->socket_fd < 0) { + return WS_CBIO_ERR_GENERAL; + } + if (slot->io_read_budget == 0U) { + return WS_CBIO_ERR_WANT_READ; + } + + word32 request_size = size; + if (request_size > slot->io_read_budget) { + request_size = slot->io_read_budget; + } + int received = recv(slot->socket_fd, data, request_size, 0); + if (received > 0) { + slot->io_read_budget -= (word32)received; + return received; + } + if (received == 0) { + return WS_CBIO_ERR_CONN_CLOSE; + } + if (errno == EAGAIN || errno == EWOULDBLOCK) { + return WS_CBIO_ERR_WANT_READ; + } + if (errno == EINTR) { + return WS_CBIO_ERR_ISR; + } + if (errno == ECONNRESET) { + return WS_CBIO_ERR_CONN_RST; + } + if (errno == ECONNABORTED) { + return WS_CBIO_ERR_CONN_CLOSE; + } + return WS_CBIO_ERR_GENERAL; +} + static int allowed_auth_types(WOLFSSH *ssh, void *context) { (void)ssh; @@ -206,6 +247,9 @@ static int authenticate_user(byte authentication_type, add_counter(&s_counters.authentication_failures, 1U); if (slot->authentication_attempts >= SSH_TRANSPORT_MAX_AUTH_ATTEMPTS) { slot->close_requested = true; + if (slot->socket_fd >= 0) { + (void)shutdown(slot->socket_fd, SHUT_RDWR); + } return WOLFSSH_USERAUTH_REJECTED; } return error == ESP_OK ? WOLFSSH_USERAUTH_INVALID_PASSWORD @@ -325,6 +369,7 @@ static esp_err_t create_context(void) return ESP_FAIL; } + wolfSSH_SetIORecv(context, bounded_ssh_receive); wolfSSH_SetUserAuth(context, authenticate_user); wolfSSH_SetUserAuthTypes(context, allowed_auth_types); (void)wolfSSH_CTX_SetChannelReqShellCb(context, accept_shell); @@ -583,6 +628,7 @@ static void accept_connections(void) publish_slot(slot, slot_index); continue; } + wolfSSH_SetIOReadCtx(slot->ssh, slot); wolfSSH_SetUserAuthCtx(slot->ssh, slot); wolfSSH_SetChannelReqCtx(slot->ssh, slot); publish_slot(slot, slot_index); @@ -637,6 +683,7 @@ static void process_handshake(ssh_slot_t *slot, size_t slot_index) return; } + slot->io_read_budget = SSH_TRANSPORT_WOLFSSH_READ_BUDGET; int result = wolfSSH_accept(slot->ssh); if (result == WS_SUCCESS) { if (!slot->authenticated || !slot->shell_requested || @@ -714,6 +761,7 @@ static bool drain_broker_events(ssh_slot_t *slot) static bool service_wolfssh_io(ssh_slot_t *slot) { word32 channel_id = 0U; + slot->io_read_budget = SSH_TRANSPORT_WOLFSSH_READ_BUDGET; int result = wolfSSH_worker(slot->ssh, &channel_id); (void)channel_id; if (result == WS_SUCCESS || wolfssh_would_block(slot->ssh, result)) { @@ -768,6 +816,7 @@ static bool receive_client_input(ssh_slot_t *slot) if (slot->rx_length != 0U) { return true; } + slot->io_read_budget = SSH_TRANSPORT_WOLFSSH_READ_BUDGET; int result = wolfSSH_stream_read(slot->ssh, slot->rx_buffer, sizeof(slot->rx_buffer)); if (result > 0) { @@ -790,6 +839,7 @@ static bool flush_client_output(ssh_slot_t *slot) return true; } + slot->io_read_budget = SSH_TRANSPORT_WOLFSSH_READ_BUDGET; int result = wolfSSH_stream_send( slot->ssh, slot->tx_buffer + slot->tx_offset, (word32)(slot->tx_length - slot->tx_offset));