path: root/src/client.c

#include "client.h"
#include "proto.h"
#include "unit.h"
#include "log.h"
#include "umbrella.h"
#include "console/rcon.h"

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <fcntl.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <sys/epoll.h>
#include <sys/stat.h>

/* ── Socket setup ────────────────────────────────────────────────────────── */

int client_listen(void) {
    /* Remove stale socket if present */
    unlink(UMBRELLA_SOCK_PATH);

    int fd = socket(AF_UNIX, SOCK_STREAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0);
    if (fd < 0) {
        log_error("socket: %s", strerror(errno));
        return -1;
    }

    struct sockaddr_un addr;
    memset(&addr, 0, sizeof(addr));
    addr.sun_family = AF_UNIX;
    strncpy(addr.sun_path, UMBRELLA_SOCK_PATH, sizeof(addr.sun_path) - 1);

    if (bind(fd, (struct sockaddr *)&addr, sizeof(addr)) != 0) {
        log_error("bind %s: %s", UMBRELLA_SOCK_PATH, strerror(errno));
        close(fd);
        return -1;
    }

    /* Owner read/write, group read/write, others nothing */
    chmod(UMBRELLA_SOCK_PATH, 0660);

    if (listen(fd, 16) != 0) {
        log_error("listen: %s", strerror(errno));
        close(fd);
        return -1;
    }

    log_info("Listening on %s", UMBRELLA_SOCK_PATH);
    return fd;
}

/* ── Client lifecycle ────────────────────────────────────────────────────── */

int client_accept(int listen_fd) {
    if (g.client_count >= MAX_CLIENTS) {
        log_warn("Max clients reached, rejecting connection");
        /* Accept and immediately close */
        int tmp = accept(listen_fd, NULL, NULL);
        if (tmp >= 0) close(tmp);
        return -1;
    }

    int fd = accept4(listen_fd, NULL, NULL, SOCK_NONBLOCK | SOCK_CLOEXEC);
    if (fd < 0) {
        if (errno != EAGAIN && errno != EWOULDBLOCK)
            log_error("accept: %s", strerror(errno));
        return -1;
    }

    /* Get peer credentials for access control */
    struct ucred cred;
    socklen_t cred_len = sizeof(cred);
    int uid = -1;
    if (getsockopt(fd, SOL_SOCKET, SO_PEERCRED, &cred, &cred_len) == 0)
        uid = (int)cred.uid;

    /* Find a free client slot */
    int slot = -1;
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (g.clients[i].fd == 0) { slot = i; break; }
    }
    if (slot < 0) {
        close(fd);
        return -1;
    }

    g.clients[slot].fd = fd;
    g.clients[slot].uid = uid;
    g.clients[slot].attached[0] = '\0';
    g.client_count++;

    /* Register with epoll */
    struct epoll_event ev;
    ev.events = EPOLLIN | EPOLLRDHUP;
    ev.data.fd = fd;
    epoll_ctl(g.epoll_fd, EPOLL_CTL_ADD, fd, &ev);

    log_debug("Client connected (fd=%d uid=%d)", fd, uid);
    return 0;
}

Client *client_find(int fd) {
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (g.clients[i].fd == fd)
            return &g.clients[i];
    }
    return NULL;
}

void client_remove(int fd) {
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (g.clients[i].fd == fd) {
            log_debug("Client disconnected (fd=%d)", fd);
            epoll_ctl(g.epoll_fd, EPOLL_CTL_DEL, fd, NULL);
            close(fd);
            memset(&g.clients[i], 0, sizeof(Client));
            g.client_count--;
            return;
        }
    }
}

void client_broadcast_output(const char *unit_name,
                             const char *data, int history) {
    for (int i = 0; i < MAX_CLIENTS; i++) {
        if (g.clients[i].fd == 0) continue;
        if (strcmp(g.clients[i].attached, unit_name) == 0)
            proto_send_output(g.clients[i].fd, unit_name, data, history);
    }
}

/* ── Command handlers ────────────────────────────────────────────────────── */

static const char *probe_rcon_state(Unit *u) {
    if (u->console.type != CONSOLE_RCON)
        return unit_state_str(u->state);

    /* Check systemd service state first -- cheap, no network */
    if (u->service[0]) {
        char cmd[256];
        snprintf(cmd, sizeof(cmd),
                 "systemctl is-active --quiet %s 2>/dev/null", u->service);
        if (system(cmd) != 0)
            return "stopped";
    }

    /* Service is active -- probe RCON to confirm reachability */
    char response[64] = {0};
    int r = rcon_exec(u->console.host, u->console.port,
                      u->console.password, "echo umbrella_probe",
                      response, sizeof(response));
    return (r == 0) ? "running" : "unreachable";
}

static void cmd_list(Client *c) {
    /* Build a JSON array of unit summaries, using active probing for state */
    char buf[PROTO_MAX_MSG];
    int pos = 0;

    pos += snprintf(buf + pos, sizeof(buf) - pos,
                    "{\"type\":\"list\",\"units\":[");

    for (int i = 0; i < g.unit_count; i++) {
        Unit *u = &g.units[i];
        if (i > 0)
            pos += snprintf(buf + pos, sizeof(buf) - pos, ",");
        pos += snprintf(buf + pos, sizeof(buf) - pos,
                        "{\"name\":\"%s\",\"display\":\"%s\","
                        "\"state\":\"%s\"}",
                        u->name, u->display, probe_rcon_state(u));
    }

    snprintf(buf + pos, sizeof(buf) - pos, "]}");
    proto_send(c->fd, buf);
}

static void cmd_status(Client *c, const char *unit_name) {
    Unit *u = unit_find(unit_name);
    if (!u) { proto_send_error(c->fd, "unit not found"); return; }

    const char *state = probe_rcon_state(u);

    char buf[512];
    snprintf(buf, sizeof(buf),
             "{\"type\":\"status\",\"unit\":\"%s\",\"display\":\"%s\","
             "\"state\":\"%s\",\"pid\":%d}",
             u->name, u->display, state, (int)u->pid);
    proto_send(c->fd, buf);
}
static void cmd_attach(Client *c, const char *unit_name) {
    Unit *u = unit_find(unit_name);
    if (!u) { proto_send_error(c->fd, "unit not found"); return; }

    strncpy(c->attached, unit_name, MAX_NAME - 1);
    log_info("Client fd=%d attached to %s", c->fd, unit_name);

    /* Send ring buffer history first */
    RingBuffer *rb = u->output;
    if (rb && rb->count > 0) {
        int start = (rb->head - rb->count + RING_BUF_LINES) % RING_BUF_LINES;
        for (int i = 0; i < rb->count; i++) {
            int idx = (start + i) % RING_BUF_LINES;
            proto_send_output(c->fd, unit_name, rb->lines[idx], 1);
        }
    }

    proto_send_ok(c->fd);
}

static void cmd_detach(Client *c) {
    if (c->attached[0])
        log_info("Client fd=%d detached from %s", c->fd, c->attached);
    c->attached[0] = '\0';
    proto_send_ok(c->fd);
}

static void cmd_input(Client *c, const char *unit_name, const char *data) {
    Unit *u = unit_find(unit_name);
    if (!u) { proto_send_error(c->fd, "unit not found"); return; }

    if (u->console.type == CONSOLE_RCON) {
        /* Send via RCON — open a connection, send, close */
        char response[4096] = {0};
        int r = rcon_exec(u->console.host, u->console.port,
                          u->console.password, data,
                          response, sizeof(response));
        if (r != 0) {
            proto_send_error(c->fd, "RCON command failed");
            return;
        }
        /* Echo the response back as output */
        if (response[0])
            proto_send_output(c->fd, unit_name, response, 0);
        proto_send_ok(c->fd);

    } else if (u->console.type == CONSOLE_STDIN) {
        if (u->stdin_fd < 0) {
            proto_send_error(c->fd, "process not running");
            return;
        }
        ssize_t n = write(u->stdin_fd, data, strlen(data));
        if (n < 0) {
            proto_send_error(c->fd, "write to stdin failed");
            return;
        }
        proto_send_ok(c->fd);
    }
}

static void cmd_action(Client *c, const char *unit_name,
                       const char *action_name) {
    Unit *u = unit_find(unit_name);
    if (!u) { proto_send_error(c->fd, "unit not found"); return; }

    /* Find the action */
    const char *script = NULL;
    for (int i = 0; i < u->action_count; i++) {
        if (strcmp(u->actions[i].name, action_name) == 0) {
            script = u->actions[i].script;
            break;
        }
    }
    if (!script) { proto_send_error(c->fd, "action not found"); return; }

    log_info("Running action '%s' for unit %s: %s",
             action_name, unit_name, script);

    /* Fork a child to run the script */
    pid_t pid = fork();
    if (pid < 0) {
        proto_send_error(c->fd, "fork failed");
        return;
    }
    if (pid == 0) {
        /* Child: exec the script */
        execl("/bin/bash", "bash", script, NULL);
        _exit(127);
    }

    /* Parent: don't wait — SIGCHLD handler will reap it */
    proto_send_ok(c->fd);
}

/* ── Main dispatch ───────────────────────────────────────────────────────── */

int client_handle(int fd) {
    Client *c = client_find(fd);
    if (!c) return -1;

    char buf[PROTO_MAX_MSG];
    int n = proto_recv(fd, buf, sizeof(buf));

    if (n <= 0) return -1;  /* 0 = disconnect, -1 = error */

    /* Extract command and fields */
    char cmd[32] = {0};
    char unit[MAX_NAME] = {0};
    char action[MAX_NAME] = {0};
    char data[RING_BUF_LINE_MAX] = {0};

    json_get_str(buf, "cmd",    cmd,    sizeof(cmd));
    json_get_str(buf, "unit",   unit,   sizeof(unit));
    json_get_str(buf, "action", action, sizeof(action));
    json_get_str(buf, "data",   data,   sizeof(data));

    log_debug("Client fd=%d cmd=%s unit=%s", fd, cmd, unit);

    if      (strcmp(cmd, "list")   == 0) cmd_list(c);
    else if (strcmp(cmd, "status") == 0) cmd_status(c, unit);
    else if (strcmp(cmd, "attach") == 0) cmd_attach(c, unit);
    else if (strcmp(cmd, "detach") == 0) cmd_detach(c);
    else if (strcmp(cmd, "input")  == 0) cmd_input(c, unit, data);
    else if (strcmp(cmd, "action") == 0) cmd_action(c, unit, action);
    else proto_send_error(fd, "unknown command");

    return 0;
}