r/src/bash_executor.c

// retoor <retoor@molodetz.nl>

#define _GNU_SOURCE
#include "bash_executor.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <sys/types.h>
#include <poll.h>
#include <signal.h>
#include <errno.h>
#include <time.h>

#define DEFAULT_TIMEOUT 300

char *r_bash_execute(const char *command, bool interactive, int timeout_seconds) {
    if (!command) {
        return strdup("Error: null command");
    }

    if (timeout_seconds <= 0) {
        timeout_seconds = DEFAULT_TIMEOUT;
    }

    size_t len = strlen(command);
    char *cmd_with_nl = malloc(len + 2);
    if (!cmd_with_nl) {
        return strdup("Error: memory allocation failed");
    }

    strcpy(cmd_with_nl, command);
    if (len > 0 && cmd_with_nl[len - 1] != '\n') {
        cmd_with_nl[len] = '\n';
        cmd_with_nl[len + 1] = '\0';
    } else if (len == 0) {
        cmd_with_nl[0] = '\n';
        cmd_with_nl[1] = '\0';
    }

    char tmp_script[] = "/tmp/r_bash_XXXXXX.sh";
    int script_fd = mkstemps(tmp_script, 3);
    if (script_fd == -1) {
        free(cmd_with_nl);
        return strdup("Error: failed to create temp script");
    }

    if (write(script_fd, cmd_with_nl, strlen(cmd_with_nl)) == -1) {
        close(script_fd);
        free(cmd_with_nl);
        unlink(tmp_script);
        return strdup("Error: failed to write to temp script");
    }
    
    close(script_fd);
    free(cmd_with_nl);

    char *output = NULL;
    size_t total_size = 0;

    if (interactive) {
        // For interactive mode, we still use system() but it doesn't easily support capturing output while timing out
        // Given the requirement, we'll try to use a simple timeout for interactive too if we can, 
        // but typically interactive means user is at it.
        // However, user said "prevent hanging processes".
        
        char *run_cmd = NULL;
        if (asprintf(&run_cmd, "timeout %ds bash %s", timeout_seconds, tmp_script) == -1) {
            unlink(tmp_script);
            return strdup("Error: asprintf failed");
        }
        int status = system(run_cmd);
        free(run_cmd);

        if (WIFEXITED(status) && WEXITSTATUS(status) == 124) {
            output = strdup("Error: Command timed out in interactive mode.");
        } else {
            if (asprintf(&output, "Command exited with status %d", status) == -1) {
                output = strdup("Command completed.");
            }
        }
    } else {
        int pipe_fds[2];
        if (pipe(pipe_fds) == -1) {
            unlink(tmp_script);
            return strdup("Error: pipe failed");
        }

        pid_t pid = fork();
        if (pid == -1) {
            close(pipe_fds[0]);
            close(pipe_fds[1]);
            unlink(tmp_script);
            return strdup("Error: fork failed");
        }

        if (pid == 0) {
            // Child
            close(pipe_fds[0]);
            dup2(pipe_fds[1], STDOUT_FILENO);
            dup2(pipe_fds[1], STDERR_FILENO);
            close(pipe_fds[1]);

            char *args[] = {"bash", tmp_script, NULL};
            execvp("bash", args);
            exit(1);
        }

        // Parent
        close(pipe_fds[1]);
        int out_fd = pipe_fds[0];
        
        struct poll_pfd {
            int fd;
            short events;
            short revents;
        } pfd;
        pfd.fd = out_fd;
        pfd.events = POLLIN;

        time_t start_time = time(NULL);
        bool timed_out = false;

        while (true) {
            time_t now = time(NULL);
            int remaining = timeout_seconds - (int)(now - start_time);
            if (remaining <= 0) {
                timed_out = true;
                break;
            }

            int ret = poll((struct pollfd *)&pfd, 1, remaining * 1000);
            if (ret == -1) {
                if (errno == EINTR) continue;
                break;
            }
            if (ret == 0) {
                timed_out = true;
                break;
            }

            if (pfd.revents & POLLIN) {
                char buffer[4096];
                ssize_t bytes = read(out_fd, buffer, sizeof(buffer) - 1);
                if (bytes <= 0) break;
                buffer[bytes] = '\0';

                // Print to stderr for user
                fprintf(stderr, "\033[2m%s\033[0m", buffer);
                fflush(stderr);

                char *new_output = realloc(output, total_size + (size_t)bytes + 1);
                if (!new_output) {
                    break;
                }
                output = new_output;
                memcpy(output + total_size, buffer, (size_t)bytes);
                total_size += (size_t)bytes;
                output[total_size] = '\0';
            } else if (pfd.revents & (POLLHUP | POLLERR)) {
                break;
            }
        }

        if (timed_out) {
            kill(-pid, SIGKILL); // Kill process group if possible
            kill(pid, SIGKILL);
            
            const char *timeout_msg = "\n[Error: Command timed out after %d seconds]\n";
            char *msg = NULL;
            if (asprintf(&msg, timeout_msg, timeout_seconds) != -1) {
                size_t msg_len = strlen(msg);
                char *new_output = realloc(output, total_size + msg_len + 1);
                if (new_output) {
                    output = new_output;
                    strcpy(output + total_size, msg);
                    total_size += msg_len;
                }
                free(msg);
            }
            
            fprintf(stderr, "\033[1;31m%s\033[0m", "\n[Timeout reached, process terminated]\n");
        }

        close(out_fd);
        waitpid(pid, NULL, WNOHANG);
    }

    if (!output) {
        output = strdup("");
    }

    unlink(tmp_script);
    return output;
}