aos/usr/echoserver/main.c

#include <stdio.h>
#include <stdlib.h>
#include <aos/aos.h>
#include <aos/simpleslab.h>
#include <aos/deferred.h>
#include <aos/aos_rpc.h>
#include <aos/ump_binding.h>
#include <aos/ump_chan.h>
#include <aos/ump_net.h>
#include <aos/ump_net_client.h>

// Echo Protocol: https://datatracker.ietf.org/doc/html/rfc862
#define PORT_ECHO 7

#define ECHO_TEST_MESSAGE "Hello Client! Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum."

__attribute__((__used__))
static void handle_payload_recv(void *arg, size_t payload_size, void *payload) {
    printf("Echo server: Received payload(%lu): %.*s\n", payload_size, payload_size, payload);
}

__attribute__((__used__))
static void handle_header_recv(void *arg, size_t header_size, void *header, size_t payload_size) {
    printf("Echo server: Received header: %u\n", *(uint8_t *)header);
    printf("Echo server: Payload is %lu bytes\n", payload_size);
    void *payload = malloc(payload_size);
    assert(payload != NULL);
    ump_recv_payload(arg, payload, handle_payload_recv, NULL);
}

__attribute__((__used__))
static void handle_send(void *arg, struct ump_send_queue_entry *entry) {
    printf("Echo server: Message sent!\n");
}

static errval_t connect_server (void *arg, struct capref cap) {
    errval_t err;

    printf("Echo server: incoming connection\n");

    struct ump_send_chan *send_chan;
    struct ump_recv_chan *recv_chan;
    err = ump_chan_init(UMP_ROLE_SERVER, &send_chan, &recv_chan, 1, cap, get_default_waitset());
    if (err_is_fail(err)) return err;

    uint8_t *value_to_send = malloc(sizeof(uint8_t));
    *value_to_send = 123;
    struct ump_send_queue_entry *entry = malloc(sizeof(struct ump_send_queue_entry));
    printf("Echo server: Sending Message\n");
    ump_send(send_chan, entry, 1, value_to_send, sizeof(ECHO_TEST_MESSAGE), ECHO_TEST_MESSAGE, handle_send, NULL);

    printf("Echo server: Receiving Message\n");
    ump_recv_header(recv_chan, handle_header_recv, recv_chan);

    return SYS_ERR_OK;
}


// UDP echo server

#define BUF_SIZE 1500
struct echo_buf {
    uint32_t src_ip;
    uint16_t src_port;
    struct ump_net_call call;
    uint8_t payload[BUF_SIZE];
};
static struct simpleslab_allocator slabs;
static uint8_t slab_buf[sizeof(struct echo_buf) * 128];


static void udp_listen_callback (void *arg, errval_t err, uint16_t src_port) {
    if (err_is_fail(err)) {
        DEBUG_ERR(err, "failed to listen on UDP");
        return;
    }
    printf("Echo server: Listening on UDP port %d.\n", src_port);
}

static void net_reply_callback (void *arg, errval_t err) {
    struct echo_buf *echo_buf = arg;
    simpleslab_free(&slabs, echo_buf);
    if (err_is_fail(err)) {
        DEBUG_ERR(err, "failed to send packet");
    }
}

static struct ump_net_ev_udp_recv net_udp_recv;

static void net_recv_packet (void *arg, size_t payload_size, void *payload) {
    struct echo_buf *echo_buf = arg;
    //printf("Echo server: Received payload: '%.*s'\n", payload_size, &echo_buf->payload);

    // Send reply
    ump_net_udp_send(
        &echo_buf->call,
        net_udp_recv.src_ip,
        PORT_ECHO, net_udp_recv.src_port,
        payload_size, &echo_buf->payload,
        net_reply_callback, echo_buf
    );
}

static void handle_udp_recv (void *arg, struct ump_net_ev_udp_recv *udp_recv, size_t payload_size) {
    if (payload_size > BUF_SIZE) {
        printf("Echo server: ERROR: Packet too big\n");
    } else {
        struct echo_buf *echo_buf = simpleslab_alloc(&slabs);
        if (echo_buf == NULL) {
            printf("Echo server: WARN: All buffers full, dropping packet\n");
        } else {
            net_udp_recv = *udp_recv;
            ump_net_recv_payload(&echo_buf->payload, net_recv_packet, echo_buf);
            return;
        }
    }
    ump_net_recv_payload(NULL, NULL, NULL);
}

// TCP echo server

#define TCP_MAX_CONNECTIONS 100
struct tcp_conn {
    struct tcp_tcb tcb;
    uint8_t rcv_buf[4*4096];
    uint8_t snd_buf[4*4096];
};

static int tcp_connection_count = 0;

static void tcp_listen_callback (void *arg, errval_t err, uint16_t src_port) {
    if (err_is_fail(err)) {
        DEBUG_ERR(err, "failed to listen on TCP");
        return;
    }
    printf("Echo server: Listening on TCP port %d.\n", src_port);
}

static void tcp_receive_available_handler (void *arg) {
    struct tcp_conn *conn = arg;
    uint8_t buf[2048];
    while (true) {
        uint32_t read_size = ump_net_tcp_receive(&conn->tcb, buf, sizeof(buf));
        if (read_size == 0) break;
        debug_printf("TCP received: '%.*s'\n", read_size, buf);
    }
}

static void tcp_send_available_handler (void *arg) {
    // struct tcp_conn *conn = arg;
    debug_printf("tcp_send_available_handler\n");
}

static void tcp_destroy_handler (void *arg) {
    struct tcp_conn *conn = arg;
    debug_printf("TCP connection destroyed for %x, port %d\n", conn->tcb.remote_ip, conn->tcb.remote_port);
    free(conn);
    tcp_connection_count--;
}

static void handle_tcp_connect (
    void *arg, struct ump_net_ev_tcp_recv *tcp_recv,
    uint32_t remote_ip, uint16_t remote_port, uint16_t local_port
) {
    debug_printf("Received connection attempt from %x, port %d\n", remote_ip, remote_port);
    if (tcp_connection_count > TCP_MAX_CONNECTIONS) {
        ump_net_tcp_refuse(tcp_recv);
    } else {
        struct tcp_conn *conn = malloc(sizeof(struct tcp_conn));
        if (conn == NULL) {
            ump_net_tcp_refuse(tcp_recv);
        } else {
            tcp_connection_count++;
            ump_net_tcp_accept(
                tcp_recv, &conn->tcb,
                tcp_receive_available_handler,
                tcp_send_available_handler,
                tcp_destroy_handler,
                conn,
                conn->rcv_buf, sizeof(conn->rcv_buf),
                conn->snd_buf, sizeof(conn->snd_buf)
            );
        }
    }
}

int main (int argc, char *argv[]) {
    errval_t err;
    struct waitset *default_ws = get_default_waitset();

    printf("Echo server: registering\n");
    struct ump_binding_server server;
    err = ump_binding_register(&server, UMP_SERVER_ECHO, connect_server, NULL);
    if (err_is_fail(err)) USER_PANIC_ERR(err, "Failed to register UMP server");
    printf("Echo server: listening\n");


    simpleslab_init(&slabs, sizeof(struct echo_buf), slab_buf, sizeof(slab_buf));

    err = ump_net_init(default_ws);
    if (err_is_fail(err)) USER_PANIC_ERR(err, "Failed to connect to net server");

    struct ump_net_call call_udp_listen;
    ump_net_udp_listen(&call_udp_listen, PORT_ECHO, handle_udp_recv, NULL, udp_listen_callback, NULL);

    struct ump_net_call call_tcp_listen;
    ump_net_tcp_listen(&call_tcp_listen, PORT_ECHO, handle_tcp_connect, NULL, tcp_listen_callback, NULL);

    while (true) {
        err = event_dispatch(default_ws);
        if (err_is_fail(err)) {
            DEBUG_ERR(err, "in event_dispatch");
            abort();
        }
    }
}