aos/lib/aos/ump_net_client.c

#include <aos/aos.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>
#include <collections/hash_table.h>

struct udp_listen_entry {
    ump_net_udp_recv_handler_t recv_handler;
    void *recv_handler_arg;
};

static bool net_initialized = false;

static struct ump_send_chan *net_send_chan;
static struct ump_recv_chan *net_recv_chan;

static struct ump_net_call *call_queue_head = NULL;
static struct ump_net_call *call_queue_tail = NULL;

static collections_hash_table* udp_listen_table;

static void net_recv_next (void);
static void handle_net_header_recv (void *arg, size_t header_size, void *header_raw, size_t payload_size);

errval_t ump_net_init (struct waitset *ws) {
    if (net_initialized) return LIB_ERR_NET_ALREADY_INIT;
    net_initialized = true;

    errval_t err;
    struct aos_rpc *rpc = aos_rpc_get_init_channel();
    struct capref net_cap;
    for (int attempts = 0; attempts < 60; attempts++) {
        err = aos_rpc_ump_connect(rpc, UMP_SERVER_NET, &net_cap);
        if (err != LIB_ERR_UMP_NOT_REGISTERED) break;
        barrelfish_usleep(500000);
    }
    if (err_is_fail(err)) return err;

    // Create the two uni-directional channels
    err = ump_chan_init(UMP_ROLE_CLIENT, &net_send_chan, &net_recv_chan,
        sizeof(struct ump_net_in_header), net_cap, ws);
    if (err_is_fail(err)) return err;

    net_recv_next();

    collections_hash_create_with_buckets(&udp_listen_table, 100, free);

    return SYS_ERR_OK;
}

static void net_recv_next (void) {
    ump_recv_header(net_recv_chan, handle_net_header_recv, NULL);
}

static void net_recv_ignore_payload (void *arg, size_t payload_size, void *payload) {
    net_recv_next();
}

// This is called twice, once from the send callback, once from the receive
// handler. We do it this way because we don't know in which order these
// two events happen.
static void ump_net_op_callback (void *arg, struct ump_send_queue_entry *entry) {
    struct ump_net_call *call = arg;
    if (call->half_done == false) {
        call->half_done = true;
        return;
    }
    if (call->in_header.op == UMP_NET_OP_UDP_SEND) {
        ((ump_net_udp_send_callback_t)call->cb)(
            call->cb_arg, call->in_header.ret_err
        );
    } else if (call->in_header.op == UMP_NET_OP_UDP_LISTEN) {
        ((ump_net_udp_listen_callback_t)call->cb)(
            call->cb_arg, call->in_header.ret_err,
            call->in_header.d.udp_listen.dest_port
        );
    } else if (call->in_header.op == UMP_NET_OP_UDP_LISTEN_STOP) {
        ((ump_net_udp_listen_stop_callback_t)call->cb)(
            call->cb_arg, call->in_header.ret_err
        );
    }
}

static void handle_net_header_recv (void *arg, size_t header_size, void *header_raw, size_t payload_size) {
    struct ump_net_in_header *header = header_raw;

    if (header->op == UMP_NET_EV_UDP_RECV) {
        struct udp_listen_entry *entry = collections_hash_find(udp_listen_table, header->d.udp_recv.dest_port);
        if (entry != NULL) {
            entry->recv_handler(entry->recv_handler_arg, &header->d.udp_recv, payload_size);
            return;
        }
    } else {
        assert(call_queue_head != NULL);
        assert(call_queue_head->out_header.op == header->op);
        call_queue_head->in_header = *header;
        ump_net_op_callback(call_queue_head, NULL);
        call_queue_head = call_queue_head->next;
        if (call_queue_head == NULL) call_queue_tail = NULL;
    }
    ump_recv_payload(net_recv_chan, NULL, net_recv_ignore_payload, NULL);
}

static void ump_net_send_call (
    struct ump_net_call *call,
    size_t payload_size, const void *payload,
    void *cb, void *cb_arg
) {
    call->half_done = false;
    call->cb = cb;
    call->cb_arg = cb_arg;
    if (call_queue_head == NULL) {
        call_queue_head = call;
    } else {
        call_queue_tail->next = call;
    }
    call_queue_tail = call;
    ump_send(net_send_chan, &call->ump_entry,
        sizeof(struct ump_net_out_header), &call->out_header,
        payload_size, payload,
        ump_net_op_callback, call);
}

void ump_net_udp_send (
    struct ump_net_call *call,
    uint32_t dest_ip,
    uint16_t src_port,
    uint16_t dest_port,
    size_t payload_size, const void *payload,
    ump_net_udp_send_callback_t cb, void *cb_arg
) {
    call->out_header = (struct ump_net_out_header){
        .op = UMP_NET_OP_UDP_SEND,
        .d = { .udp_send = {
            .dest_ip = dest_ip,
            .src_port = src_port,
            .dest_port = dest_port,
        } }
    };
    ump_net_send_call(call, payload_size, payload, cb, cb_arg);
}

void ump_net_udp_listen (
    struct ump_net_call *call,
    uint16_t dest_port,
    ump_net_udp_recv_handler_t recv_handler, void *recv_handler_arg,
    ump_net_udp_listen_callback_t cb, void *cb_arg
) {
    if (collections_hash_find(udp_listen_table, dest_port) != NULL) {
        cb(cb_arg, LIB_ERR_NET_PORT_IN_USE, 0);
        return;
    }
    struct udp_listen_entry *entry = malloc(sizeof(struct udp_listen_entry));
    if (entry == NULL) {
        cb(cb_arg, LIB_ERR_MALLOC_FAIL, 0);
        return;
    }
    entry->recv_handler = recv_handler;
    entry->recv_handler_arg = recv_handler_arg;
    collections_hash_insert(udp_listen_table, dest_port, entry);

    call->out_header = (struct ump_net_out_header){
        .op = UMP_NET_OP_UDP_LISTEN,
        .d = { .udp_listen = { .dest_port = dest_port } }
    };
    ump_net_send_call(call, 0, NULL, cb, cb_arg);
}

void ump_net_udp_listen_stop (
    struct ump_net_call *call,
    uint16_t dest_port,
    ump_net_udp_listen_stop_callback_t cb, void *cb_arg
) {
    if (collections_hash_find(udp_listen_table, dest_port) == NULL) {
        cb(cb_arg, LIB_ERR_NET_NOT_LISTENING);
        return;
    }
    collections_hash_delete(udp_listen_table, dest_port);

    call->out_header = (struct ump_net_out_header){
        .op = UMP_NET_OP_UDP_LISTEN_STOP,
        .d = { .udp_listen_stop = { .dest_port = dest_port } }
    };
    ump_net_send_call(call, 0, NULL, cb, cb_arg);
}

static ump_recv_payload_callback_fn_t recv_payload_callback;
static void *recv_payload_callback_arg;

static void recv_payload_callback_wrap (void *arg, size_t payload_size, void *payload) {
    if (recv_payload_callback != NULL) {
        recv_payload_callback(recv_payload_callback_arg, payload_size, payload);
    }
    net_recv_next();
}

void ump_net_recv_payload (
    void *payload,
    ump_recv_payload_callback_fn_t callback, void *callback_arg
) {
    recv_payload_callback = callback;
    recv_payload_callback_arg = callback_arg;
    ump_recv_payload(net_recv_chan, payload, recv_payload_callback_wrap, NULL);
}