/** * Copyright (c) 2022 Brian Starkey * * Based on the Pico W tcp_server example: * Copyright (c) 2022 Raspberry Pi (Trading) Ltd. * * SPDX-License-Identifier: BSD-3-Clause */ #include #include #include "pico/stdlib.h" #include "pico/cyw43_arch.h" #include "lwip/pbuf.h" #include "lwip/tcp.h" extern const char *wifi_ssid; extern const char *wifi_pass; #define TCP_PORT 4242 #define DEBUG_printf printf #define POLL_TIME_S 5 #define MAX_LEN 2048 #define COMM_MAX_NARG 5 #define COMM_MAX_DATA_LEN 1024 #define COMM_RSP_OK (('O' << 0) | ('K' << 8) | ('O' << 16) | ('K' << 24)) #define COMM_RSP_ERR (('E' << 0) | ('R' << 8) | ('R' << 16) | ('!' << 24)) #define CMD_SYNC (('S' << 0) | ('Y' << 8) | ('N' << 16) | ('C' << 24)) #define RSP_SYNC (('W' << 0) | ('O' << 8) | ('T' << 16) | ('A' << 24)) struct comm_command { uint32_t opcode; uint32_t nargs; uint32_t resp_nargs; uint32_t (*size)(uint32_t *args_in, uint32_t *data_len_out, uint32_t *resp_data_len_out); uint32_t (*handle)(uint32_t *args_in, uint8_t *data_in, uint32_t *resp_args_out, uint8_t *resp_data_out); }; enum conn_state { CONN_STATE_WAIT_FOR_SYNC, CONN_STATE_READ_OPCODE, CONN_STATE_READ_ARGS, CONN_STATE_READ_DATA, CONN_STATE_HANDLE, CONN_STATE_WRITE_RESP, CONN_STATE_WRITE_ERROR, CONN_STATE_CLOSED, }; struct tcp_comm_ctx { struct tcp_pcb *serv_pcb; volatile bool serv_done; enum conn_state conn_state; struct tcp_pcb *conn_pcb; uint8_t buf[(sizeof(uint32_t) * (1 + COMM_MAX_NARG)) + COMM_MAX_DATA_LEN]; uint16_t rx_bytes_received; uint16_t rx_bytes_remaining; uint16_t tx_bytes_sent; uint16_t tx_bytes_remaining; uint32_t resp_data_len; const struct comm_command *cmd; const struct comm_command *const *cmds; unsigned int n_cmds; uint32_t sync_opcode; }; #define COMM_BUF_OPCODE(_buf) ((uint32_t *)((uint8_t *)(_buf))) #define COMM_BUF_ARGS(_buf) ((uint32_t *)((uint8_t *)(_buf) + sizeof(uint32_t))) #define COMM_BUF_BODY(_buf, _nargs) ((uint8_t *)(_buf) + (sizeof(uint32_t) * ((_nargs) + 1))) static const struct comm_command *find_command_desc(struct tcp_comm_ctx *ctx, uint32_t opcode) { unsigned int i; for (i = 0; i < ctx->n_cmds; i++) { if (ctx->cmds[i]->opcode == opcode) { return ctx->cmds[i]; } } return NULL; } static bool is_error(uint32_t status) { return status == COMM_RSP_ERR; } static int tcp_conn_sync_begin(struct tcp_comm_ctx *ctx); static int tcp_conn_sync_complete(struct tcp_comm_ctx *ctx); static int tcp_conn_opcode_begin(struct tcp_comm_ctx *ctx); static int tcp_conn_opcode_complete(struct tcp_comm_ctx *ctx); static int tcp_conn_args_begin(struct tcp_comm_ctx *ctx); static int tcp_conn_args_complete(struct tcp_comm_ctx *ctx); static int tcp_conn_data_begin(struct tcp_comm_ctx *ctx, uint32_t data_len); static int tcp_conn_data_complete(struct tcp_comm_ctx *ctx); static int tcp_conn_response_begin(struct tcp_comm_ctx *ctx); static int tcp_conn_response_complete(struct tcp_comm_ctx *ctx); static int tcp_conn_error_begin(struct tcp_comm_ctx *ctx); static int tcp_conn_sync_begin(struct tcp_comm_ctx *ctx) { ctx->conn_state = CONN_STATE_WAIT_FOR_SYNC; ctx->rx_bytes_received = 0; ctx->rx_bytes_remaining = sizeof(uint32_t); DEBUG_printf("sync_begin %d\n", ctx->rx_bytes_remaining); } static int tcp_conn_sync_complete(struct tcp_comm_ctx *ctx) { if (ctx->sync_opcode != *COMM_BUF_OPCODE(ctx->buf)) { DEBUG_printf("sync not correct: %c%c%c%c\n", ctx->buf[0], ctx->buf[1], ctx->buf[2], ctx->buf[3]); return tcp_conn_error_begin(ctx); } return tcp_conn_opcode_complete(ctx); } static int tcp_conn_opcode_begin(struct tcp_comm_ctx *ctx) { ctx->conn_state = CONN_STATE_READ_OPCODE; ctx->rx_bytes_received = 0; ctx->rx_bytes_remaining = sizeof(uint32_t); return 0; } static int tcp_conn_opcode_complete(struct tcp_comm_ctx *ctx) { ctx->cmd = find_command_desc(ctx, *COMM_BUF_OPCODE(ctx->buf)); if (!ctx->cmd) { DEBUG_printf("no command for '%c%c%c%c'\n", ctx->buf[0], ctx->buf[1], ctx->buf[2], ctx->buf[3]); return tcp_conn_error_begin(ctx); } else { DEBUG_printf("got command '%c%c%c%c'\n", ctx->buf[0], ctx->buf[1], ctx->buf[2], ctx->buf[3]); } return tcp_conn_args_begin(ctx); } static int tcp_conn_args_begin(struct tcp_comm_ctx *ctx) { ctx->conn_state = CONN_STATE_READ_ARGS; ctx->rx_bytes_received = 0; ctx->rx_bytes_remaining = ctx->cmd->nargs * sizeof(uint32_t); if (ctx->cmd->nargs == 0) { return tcp_conn_args_complete(ctx); } return 0; } static int tcp_conn_args_complete(struct tcp_comm_ctx *ctx) { const struct comm_command *cmd = ctx->cmd; uint32_t data_len = 0; if (cmd->size) { uint32_t status = cmd->size(COMM_BUF_ARGS(ctx->buf), &data_len, &ctx->resp_data_len); if (is_error(status)) { return tcp_conn_error_begin(ctx); } } return tcp_conn_data_begin(ctx, data_len); } static int tcp_conn_data_begin(struct tcp_comm_ctx *ctx, uint32_t data_len) { const struct comm_command *cmd = ctx->cmd; ctx->conn_state = CONN_STATE_READ_DATA; ctx->rx_bytes_received = 0; ctx->rx_bytes_remaining = data_len; if (data_len == 0) { return tcp_conn_data_complete(ctx); } return 0; } static int tcp_conn_data_complete(struct tcp_comm_ctx *ctx) { const struct comm_command *cmd = ctx->cmd; if (cmd->handle) { uint32_t status = cmd->handle(COMM_BUF_ARGS(ctx->buf), COMM_BUF_BODY(ctx->buf, cmd->nargs), COMM_BUF_ARGS(ctx->buf), COMM_BUF_BODY(ctx->buf, cmd->resp_nargs)); if (is_error(status)) { return tcp_conn_error_begin(ctx); } *COMM_BUF_OPCODE(ctx->buf) = status; } else { // TODO: Should we just assert(desc->handle)? *COMM_BUF_OPCODE(ctx->buf) = COMM_RSP_OK; } return tcp_conn_response_begin(ctx); } static int tcp_conn_response_begin(struct tcp_comm_ctx *ctx) { ctx->conn_state = CONN_STATE_WRITE_RESP; ctx->tx_bytes_sent = 0; ctx->tx_bytes_remaining = ctx->resp_data_len + ((ctx->cmd->resp_nargs + 1) * sizeof(uint32_t)); err_t err = tcp_write(ctx->conn_pcb, ctx->buf, ctx->tx_bytes_remaining, 0); if (err != ERR_OK) { return -1; } return 0; } static int tcp_conn_error_begin(struct tcp_comm_ctx *ctx) { ctx->conn_state = CONN_STATE_WRITE_ERROR; ctx->tx_bytes_sent = 0; ctx->tx_bytes_remaining = sizeof(uint32_t); *COMM_BUF_OPCODE(ctx->buf) = COMM_RSP_ERR; err_t err = tcp_write(ctx->conn_pcb, ctx->buf, ctx->tx_bytes_remaining, 0); if (err != ERR_OK) { return -1; } return 0; } static int tcp_conn_response_complete(struct tcp_comm_ctx *ctx) { return tcp_conn_opcode_begin(ctx); } static int tcp_conn_rx_complete(struct tcp_comm_ctx *ctx) { switch (ctx->conn_state) { case CONN_STATE_WAIT_FOR_SYNC: return tcp_conn_sync_complete(ctx); case CONN_STATE_READ_OPCODE: return tcp_conn_opcode_complete(ctx); case CONN_STATE_READ_ARGS: return tcp_conn_args_complete(ctx); case CONN_STATE_READ_DATA: return tcp_conn_data_complete(ctx); default: return -1; } } static int tcp_conn_tx_complete(struct tcp_comm_ctx *ctx) { switch (ctx->conn_state) { case CONN_STATE_WRITE_RESP: return tcp_conn_response_complete(ctx); case CONN_STATE_WRITE_ERROR: return -1; default: return -1; } } static err_t tcp_conn_close(struct tcp_comm_ctx *ctx) { err_t err = ERR_OK; cyw43_arch_gpio_put (0, false); ctx->conn_state = CONN_STATE_CLOSED; if (!ctx->conn_pcb) { return err; } tcp_arg(ctx->conn_pcb, NULL); tcp_poll(ctx->conn_pcb, NULL, 0); tcp_sent(ctx->conn_pcb, NULL); tcp_recv(ctx->conn_pcb, NULL); tcp_err(ctx->conn_pcb, NULL); err = tcp_close(ctx->conn_pcb); if (err != ERR_OK) { DEBUG_printf("close failed %d, calling abort\n", err); tcp_abort(ctx->conn_pcb); err = ERR_ABRT; } ctx->conn_pcb = NULL; return err; } static err_t tcp_server_close(void *arg) { struct tcp_comm_ctx *ctx = (struct tcp_comm_ctx *)arg; err_t err = ERR_OK; err = tcp_conn_close(ctx); if ((err != ERR_OK) && ctx->serv_pcb) { tcp_arg(ctx->serv_pcb, NULL); tcp_abort(ctx->serv_pcb); ctx->serv_pcb = NULL; return ERR_ABRT; } if (!ctx->serv_pcb) { return err; } tcp_arg(ctx->serv_pcb, NULL); err = tcp_close(ctx->serv_pcb); if (err != ERR_OK) { tcp_abort(ctx->serv_pcb); err = ERR_ABRT; } ctx->serv_pcb = NULL; return err; } static void tcp_server_complete(void *arg, int status) { struct tcp_comm_ctx *ctx = (struct tcp_comm_ctx *)arg; if (status == 0) { DEBUG_printf("server completed normally\n"); } else { DEBUG_printf("server error %d\n", status); } tcp_server_close(ctx); ctx->serv_done = true; } static err_t tcp_conn_complete(void *arg, int status) { struct tcp_comm_ctx *ctx = (struct tcp_comm_ctx *)arg; if (status == 0) { DEBUG_printf("conn completed normally\n"); } else { DEBUG_printf("conn error %d\n", status); } return tcp_conn_close(ctx); } static err_t tcp_conn_sent(void *arg, struct tcp_pcb *tpcb, u16_t len) { struct tcp_comm_ctx *ctx = (struct tcp_comm_ctx *)arg; DEBUG_printf("tcp_server_sent %u\n", len); cyw43_arch_lwip_check(); if (len > ctx->tx_bytes_remaining) { DEBUG_printf("tx len %d > remaining %d\n", len, ctx->tx_bytes_remaining); return tcp_conn_complete(ctx, ERR_ARG); } ctx->tx_bytes_remaining -= len; ctx->tx_bytes_sent += len; if (ctx->tx_bytes_remaining == 0) { int res = tcp_conn_tx_complete(ctx); if (res) { return tcp_conn_complete(ctx, ERR_ARG); } } return ERR_OK; } static err_t tcp_conn_recv(void *arg, struct tcp_pcb *tpcb, struct pbuf *p, err_t err) { struct tcp_comm_ctx *ctx = (struct tcp_comm_ctx *)arg; if (!p) { DEBUG_printf("no pbuf\n"); return tcp_conn_complete(ctx, 0); } // this method is callback from lwIP, so cyw43_arch_lwip_begin is not required, however you // can use this method to cause an assertion in debug mode, if this method is called when // cyw43_arch_lwip_begin IS needed cyw43_arch_lwip_check(); if (p->tot_len > 0) { DEBUG_printf("tcp_server_recv %d err %d\n", p->tot_len, err); size_t to_copy = p->tot_len > ctx->rx_bytes_remaining ? ctx->rx_bytes_remaining : p->tot_len; // Receive the buffer if (pbuf_copy_partial(p, ctx->buf + ctx->rx_bytes_received, to_copy, 0) != to_copy) { DEBUG_printf("wrong copy len\n"); return tcp_conn_complete(ctx, ERR_ARG); } ctx->rx_bytes_received += to_copy; ctx->rx_bytes_remaining -= to_copy; tcp_recved(tpcb, p->tot_len); if (ctx->rx_bytes_remaining == 0) { int res = tcp_conn_rx_complete(ctx); if (res) { return tcp_conn_complete(ctx, ERR_ARG); } } } pbuf_free(p); return ERR_OK; } static err_t tcp_conn_poll(void *arg, struct tcp_pcb *tpcb) { DEBUG_printf("tcp_server_poll_fn\n"); return ERR_OK; } static void tcp_conn_err(void *arg, err_t err) { struct tcp_comm_ctx *ctx = (struct tcp_comm_ctx *)arg; DEBUG_printf("tcp_conn_err %d\n", err); ctx->conn_pcb = NULL; ctx->conn_state = CONN_STATE_CLOSED; ctx->rx_bytes_remaining = 0; cyw43_arch_gpio_put (0, false); } static void tcp_conn_init(struct tcp_comm_ctx *ctx, struct tcp_pcb *pcb) { ctx->conn_pcb = pcb; tcp_arg(pcb, ctx); cyw43_arch_gpio_put (0, true); tcp_conn_sync_begin(ctx); tcp_sent(pcb, tcp_conn_sent); tcp_recv(pcb, tcp_conn_recv); tcp_poll(pcb, tcp_conn_poll, POLL_TIME_S * 2); tcp_err(pcb, tcp_conn_err); } static err_t tcp_server_accept(void *arg, struct tcp_pcb *client_pcb, err_t err) { struct tcp_comm_ctx *ctx = (struct tcp_comm_ctx *)arg; if (err != ERR_OK || client_pcb == NULL) { DEBUG_printf("Failure in accept\n"); tcp_server_complete(ctx, err); return ERR_VAL; } DEBUG_printf("Connection opened\n"); if (ctx->conn_pcb) { DEBUG_printf("Already have a connection\n"); tcp_abort(client_pcb); return ERR_ABRT; } tcp_conn_init(ctx, client_pcb); return ERR_OK; } static err_t tcp_server_listen(struct tcp_comm_ctx *ctx) { DEBUG_printf("Starting server at %s on port %u\n", ip4addr_ntoa(netif_ip4_addr(netif_list)), TCP_PORT); ctx->serv_done = false; struct tcp_pcb *pcb = tcp_new_ip_type(IPADDR_TYPE_ANY); if (!pcb) { DEBUG_printf("failed to create pcb\n"); return ERR_MEM; } err_t err = tcp_bind(pcb, NULL, TCP_PORT); if (err) { DEBUG_printf("failed to bind to port %d\n", TCP_PORT); tcp_abort(pcb); return err; } ctx->serv_pcb = tcp_listen_with_backlog_and_err(pcb, 1, &err); if (!ctx->serv_pcb) { DEBUG_printf("failed to listen: %d\n", err); return err; } tcp_arg(ctx->serv_pcb, ctx); tcp_accept(ctx->serv_pcb, tcp_server_accept); return ERR_OK; } static uint32_t handle_sync(uint32_t *args_in, uint8_t *data_in, uint32_t *resp_args_out, uint8_t *resp_data_out) { return RSP_SYNC; } const struct comm_command util_sync_cmd = { .opcode = CMD_SYNC, .nargs = 0, .resp_nargs = 0, .size = NULL, .handle = &handle_sync, }; static void tcp_comm_init(struct tcp_comm_ctx *ctx, const struct comm_command *const *cmds, unsigned int n_cmds, uint32_t sync_opcode) { unsigned int i; for (i = 0; i < n_cmds; i++) { assert(cmds[i]->nargs <= MAX_NARG); assert(cmds[i]->resp_nargs <= MAX_NARG); } memset(ctx, 0, sizeof(*ctx)); ctx->cmds = cmds; ctx->n_cmds = n_cmds; ctx->sync_opcode = sync_opcode; } int main() { stdio_init_all(); sleep_ms(1000); if (cyw43_arch_init()) { printf("failed to initialise\n"); return 1; } cyw43_arch_enable_sta_mode(); printf("Connecting to WiFi...\n"); if (cyw43_arch_wifi_connect_timeout_ms(wifi_ssid, wifi_pass, CYW43_AUTH_WPA2_AES_PSK, 30000)) { printf("failed to connect.\n"); return 1; } else { printf("Connected.\n"); } struct tcp_comm_ctx tcp; const struct comm_command *cmds[] = { &util_sync_cmd, }; tcp_comm_init(&tcp, cmds, 1, CMD_SYNC); for ( ; ; ) { err_t err = tcp_server_listen(&tcp); if (err != ERR_OK) { printf("Failed to start server: %d\n", err); sleep_ms(1000); continue; } while (!tcp.serv_done) { cyw43_arch_poll(); sleep_ms(10); } } cyw43_arch_deinit(); return 0; }