echo-server.c

  /*
 * A minimal io_uring-based TCP echo server.
 *
 * This server demonstrates the basic principles of using io_uring for network
 * programming. It's designed for simplicity and clarity over performance or
 * robustness.
 *
 * Core Logic:
 * 1. Setup a listening TCP socket.
 * 2. Initialize an io_uring instance.
 * 3. Submit an initial `accept` operation to the uring.
 * 4. Enter the main event loop:
 * a. Wait for a completion event (CQE) from the uring.
 * b. Process the completed event based on its type (accept, read, or write).
 * c. Chain operations:
 * - An `accept` completion triggers a `read` on the new socket and
 * submits a new `accept` operation.
 * - A `read` completion triggers a `write` to echo the data back.
 * - A `write` completion triggers a new `read` to wait for more data.
 * 5. A zero-byte read indicates the client has closed the connection, so the
 * server closes the socket.
 *
 * How to compile:
 * gcc -Wall -O2 -o echo-server echo-server.c -luring
 *
 * Requires liburing to be installed.
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <netinet/in.h>
#include <sys/socket.h>
#include <unistd.h>
#include <liburing.h>
#include <errno.h>

#define PORT 8080
#define BUFFER_SIZE 1024
#define QUEUE_DEPTH 256 // How many operations can be in-flight

// Enum to identify the type of operation associated with a submission
enum {
    EVENT_TYPE_ACCEPT,
    EVENT_TYPE_READ,
    EVENT_TYPE_WRITE,
};

// A request structure to hold data for each in-flight operation.
// We pass a pointer to this struct in the `user_data` field of the SQE.
typedef struct request {
    int event_type;
    int client_fd;
    struct iovec iov;
    char buffer[BUFFER_SIZE];
} request_t;

// Function Prototypes
static void submit_accept_request(struct io_uring *ring, int listen_fd);
static void submit_read_request(struct io_uring *ring, int client_fd);
static void submit_write_request(struct io_uring *ring, request_t *req, int bytes_read);
static int setup_listen_socket(int port);

int main() {
    // 1. Set up the listening socket
    int listen_fd = setup_listen_socket(PORT);
    if (listen_fd < 0) {
        fprintf(stderr, "Failed to set up listening socket.\n");
        return 1;
    }

    // 2. Initialize io_uring
    struct io_uring ring;
    if (io_uring_queue_init(QUEUE_DEPTH, &ring, 0) < 0) {
        perror("io_uring_queue_init");
        close(listen_fd);
        return 1;
    }

    printf("TCP Echo Server listening on port %d\n", PORT);

    // 3. Submit the first accept request
    submit_accept_request(&ring, listen_fd);

    // 4. Main event loop
    while (1) {
        // Submit queued operations and wait for a completion
        io_uring_submit_and_wait(&ring, 1);

        struct io_uring_cqe *cqe;
        unsigned head;
        unsigned count = 0;

        // Iterate over all completed events in the completion queue
        io_uring_for_each_cqe(&ring, head, cqe) {
            count++;
            request_t *req = (request_t *)cqe->user_data;

            // A negative result in the CQE indicates an error.
            if (cqe->res < 0) {
                // The value is the negated errno.
                fprintf(stderr, "Async operation failed: %s\n", strerror(-cqe->res));
                
                // If an error occurs, we must clean up associated resources.
                if (req) {
                    // For a read or write error, the connection is likely compromised,
                    // so we close the client socket. For an accept error, there is no
                    // client socket to close yet.
                    if (req->event_type != EVENT_TYPE_ACCEPT) {
                        close(req->client_fd);
                    }
                    // Free the request structure itself.
                    free(req);
                }
                // Continue to the next completion event.
                continue;
            }

            switch (req->event_type) {
                case EVENT_TYPE_ACCEPT: {
                    int client_fd = cqe->res;
                    printf("New connection accepted, fd: %d\n", client_fd);

                    // A new connection is accepted, submit a read request for it
                    submit_read_request(&ring, client_fd);

                    // Immediately submit another accept request to keep listening
                    submit_accept_request(&ring, listen_fd);
                    free(req); // Free the original accept request data
                    break;
                }
                case EVENT_TYPE_READ: {
                    int bytes_read = cqe->res;
                    if (bytes_read == 0) {
                        // Client closed the connection
                        printf("Client on fd %d disconnected.\n", req->client_fd);
                        close(req->client_fd);
                        free(req);
                    } else {
                        // Echo the data back by submitting a write request
                        submit_write_request(&ring, req, bytes_read);
                    }
                    break;
                }
                case EVENT_TYPE_WRITE: {
                    // Write is complete, submit a new read request to wait for more data
                    submit_read_request(&ring, req->client_fd);
                    free(req); // Free the write request data
                    break;
                }
            }
        }
        // Mark the processed CQEs as consumed
        io_uring_cq_advance(&ring, count);
    }

    // Cleanup
    io_uring_queue_exit(&ring);
    close(listen_fd);
    return 0;
}

// Helper function to create and configure the listening socket
static int setup_listen_socket(int port) {
    int fd;
    struct sockaddr_in serv_addr;

    fd = socket(AF_INET, SOCK_STREAM, 0);
    if (fd < 0) {
        perror("socket");
        return -1;
    }

    int opt = 1;
    setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, &opt, sizeof(opt));

    memset(&serv_addr, 0, sizeof(serv_addr));
    serv_addr.sin_family = AF_INET;
    serv_addr.sin_addr.s_addr = INADDR_ANY;
    serv_addr.sin_port = htons(port);

    if (bind(fd, (struct sockaddr *)&serv_addr, sizeof(serv_addr)) < 0) {
        perror("bind");
        close(fd);
        return -1;
    }

    if (listen(fd, 128) < 0) {
        perror("listen");
        close(fd);
        return -1;
    }

    return fd;
}

// Prepares and submits an accept request to the io_uring
static void submit_accept_request(struct io_uring *ring, int listen_fd) {
    struct io_uring_sqe *sqe = io_uring_get_sqe(ring);
    if (!sqe) {
        fprintf(stderr, "Could not get SQE for accept.\n");
        return;
    }

    request_t *req = (request_t *)malloc(sizeof(request_t));
    if (!req) {
        fprintf(stderr, "Failed to allocate memory for accept request.\n");
        // Can't recover, but let's not submit a broken request.
        // In a real app, you'd have a better memory management strategy.
        io_uring_sqe_set_data(sqe, NULL); // Nullify to avoid use-after-free
        return;
    }
    req->event_type = EVENT_TYPE_ACCEPT;

    io_uring_prep_accept(sqe, listen_fd, NULL, NULL, 0);
    io_uring_sqe_set_data(sqe, req);
}

// Prepares and submits a read request to the io_uring
static void submit_read_request(struct io_uring *ring, int client_fd) {
    struct io_uring_sqe *sqe = io_uring_get_sqe(ring);
    if (!sqe) {
        fprintf(stderr, "Could not get SQE for read.\n");
        return;
    }

    request_t *req = (request_t *)malloc(sizeof(request_t));
    if (!req) {
        fprintf(stderr, "Failed to allocate memory for read request.\n");
        io_uring_sqe_set_data(sqe, NULL);
        return;
    }
    req->event_type = EVENT_TYPE_READ;
    req->client_fd = client_fd;
    req->iov.iov_base = req->buffer;
    req->iov.iov_len = BUFFER_SIZE;

    io_uring_prep_readv(sqe, client_fd, &req->iov, 1, 0);
    io_uring_sqe_set_data(sqe, req);
}

// Prepares and submits a write request to the io_uring
static void submit_write_request(struct io_uring *ring, request_t *req, int bytes_read) {
    struct io_uring_sqe *sqe = io_uring_get_sqe(ring);
    if (!sqe) {
        fprintf(stderr, "Could not get SQE for write.\n");
        // The original read request `req` will be leaked here.
        // A real app needs robust resource cleanup.
        return;
    }

    req->event_type = EVENT_TYPE_WRITE;
    req->iov.iov_len = bytes_read; // Set the length to the number of bytes we actually read

    io_uring_prep_writev(sqe, req->client_fd, &req->iov, 1, 0);
    io_uring_sqe_set_data(sqe, req);
}