bantam-cpp/tools.cpp

/* tools.cpp — bash/read/write tools and oversized result offloading. */
#include "bantam.hpp"
#include <algorithm>
#include <array>
#include <cerrno>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <dirent.h>
#include <fcntl.h>
#include <poll.h>
#include <signal.h>
#include <sys/stat.h>
#include <sys/wait.h>
#include <unistd.h>
#include <vector>

/* ----------------------------- shell ---------------------------------- */
/* Resolve the Bourne-again shell from the system PATH. Returns a static
 * buffer holding the absolute path to `bash`, or "/bin/sh" when bash cannot
 * be located anywhere on PATH. Never returns a hardcoded bash path. */
static const char *resolve_shell() {
  static char buf[4096];
  static const char *cached = nullptr;
  if (cached) return cached;
  const char *path = getenv("PATH");
  if (path && *path) {
    std::string p(path);
    size_t start = 0;
    while (start <= p.size()) {
      size_t end = p.find(':', start);
      std::string dir = (end == std::string::npos) ? p.substr(start)
                                                   : p.substr(start, end - start);
      if (!dir.empty() && dir.back() != '/') dir += '/';
      std::string cand = dir + "bash";
      struct stat st;
      if (stat(cand.c_str(), &st) == 0 && S_ISREG(st.st_mode) &&
          (st.st_mode & (S_IXUSR | S_IXGRP | S_IXOTH)) != 0) {
        snprintf(buf, sizeof buf, "%s", cand.c_str());
        cached = buf;
        return cached;
      }
      if (end == std::string::npos) break;
      start = end + 1;
    }
  }
  cached = "/bin/sh";
  return cached;
}

std::string shell_with_workdir(std::string_view cmd, std::string_view workdir,
                               int timeout, volatile int *interrupted) {
  std::string fcmd = filter_text(cmd);
  if (timeout <= 0) timeout = 120;
  const char *sh = resolve_shell();

  int pipefd[2];
  if (pipe(pipefd) != 0) return "[tool error: pipe failed]";
  pid_t pid = fork();
  if (pid < 0) {
    close(pipefd[0]);
    close(pipefd[1]);
    return "[tool error: fork failed]";
  }
  if (pid == 0) {
    /* child: create a dedicated process group so timeout/SIGINT kills the
     * entire process tree spawned by the shell. */
    setpgid(0, 0);
    dup2(pipefd[1], STDOUT_FILENO);
    dup2(pipefd[1], STDERR_FILENO);
    close(pipefd[0]);
    close(pipefd[1]);
    if (!workdir.empty()) chdir(std::string(workdir).c_str());
    const char *shname = (strrchr(sh, '/') ? strrchr(sh, '/') + 1 : sh);
    execl(sh, shname, "-c", fcmd.c_str(), (char *)nullptr);
    _exit(127);
  }
  close(pipefd[1]);

  int flags = fcntl(pipefd[0], F_GETFL, 0);
  fcntl(pipefd[0], F_SETFL, flags | O_NONBLOCK);

  std::string out;
  struct pollfd pfd; pfd.fd = pipefd[0]; pfd.events = POLLIN;
  time_t deadline = time(nullptr) + timeout;
  bool killed = false;
  while (true) {
    if (*interrupted) break;
    int remain = (int)(deadline - time(nullptr));
    if (remain <= 0) { killed = true; break; }
    int wait_ms = remain * 1000;
    if (wait_ms > 200) wait_ms = 200;
    int pr = poll(&pfd, 1, wait_ms);
    if (pr < 0) { if (errno == EINTR) continue; break; }
    if (pr == 0) continue;
    char buf[4096];
    ssize_t n = read(pipefd[0], buf, sizeof buf);
    if (n > 0) out.append(buf, (size_t)n);
    else if (n == 0) break; /* EOF: child exited */
    else if (errno == EAGAIN || errno == EWOULDBLOCK) continue;
    else break;
  }
  if (*interrupted || killed) {
    kill(-pid, SIGKILL);
    kill(pid, SIGKILL);
    /* drain remaining output briefly */
    struct pollfd dpfd; dpfd.fd = pipefd[0]; dpfd.events = POLLIN;
    for (int i = 0; i < 20; i++) {
      int pr = poll(&dpfd, 1, 50);
      if (pr <= 0) break;
      char buf[4096];
      ssize_t n = read(pipefd[0], buf, sizeof buf);
      if (n > 0) out.append(buf, (size_t)n);
      else break;
    }
  }
  int status = 0;
  while (waitpid(pid, &status, 0) < 0 && errno == EINTR) {}
  close(pipefd[0]);

  std::string filtered = filter_text(out);
  /* trim trailing whitespace */
  while (!filtered.empty() &&
         (filtered.back() == '\n' || filtered.back() == '\r' ||
          filtered.back() == ' '  || filtered.back() == '\t'))
    filtered.pop_back();

  if (*interrupted) {
    return filtered + "\n\nexit: -1";
  }
  if (killed) {
    return sfmt("%s\n\n[shell timeout after %ds]\nexit: -1", filtered.c_str(), timeout);
  }
  int code = 0;
  if (WIFEXITED(status)) code = WEXITSTATUS(status);
  else if (WIFSIGNALED(status)) code = -1;
  return sfmt("%s\n\nexit: %d", filtered.c_str(), code);
}

/* ----------------------------- read ----------------------------------- */
std::string read_file_or_dir(std::string_view path, int offset, int limit, int *style) {
  std::string p(path);
  /* trim */
  size_t s = 0;
  while (s < p.size() && (p[s] == ' ' || p[s] == '\t')) s++;
  p = p.substr(s);
  if (p.empty()) { *style = 31; return "[tool error: read requires 'filePath' parameter]"; }
  struct stat st;
  if (stat(p.c_str(), &st) != 0) {
    *style = 31;
    return sfmt("[tool error: read %s: %s]", p.c_str(), strerror(errno));
  }
  if (S_ISDIR(st.st_mode)) {
    struct DirEnt {
      std::string name;
      int is_dir; /* -1 = unknown, 0 = non-dir, 1 = dir */
    };
    std::vector<DirEnt> entries;
    {
      DIR *dir = opendir(p.c_str());
      if (!dir) { *style = 31; return "[tool error: read (directory listing failed)]"; }
      struct dirent *de;
      while ((de = readdir(dir))) {
        if (de->d_name[0] == '.' &&
            (de->d_name[1] == '\0' || (de->d_name[1] == '.' && de->d_name[2] == '\0')))
          continue;
        int isdir = -1;
#if defined(DT_DIR) && defined(DT_UNKNOWN) && defined(DT_LNK)
        if (de->d_type != DT_UNKNOWN && de->d_type != DT_LNK) {
          isdir = (de->d_type == DT_DIR) ? 1 : 0;
        }
#endif
        entries.push_back(DirEnt{de->d_name, isdir});
      }
      closedir(dir);
    }
    std::sort(entries.begin(), entries.end(), [](const DirEnt &a, const DirEnt &b) {
      return a.name < b.name;
    });
    std::string b;
    for (auto &ent : entries) {
      bool isdir = (ent.is_dir == 1);
      if (ent.is_dir < 0) {
        std::string full = p + "/" + ent.name;
        struct stat ds;
        isdir = (stat(full.c_str(), &ds) == 0 && S_ISDIR(ds.st_mode));
      }
      b += ent.name;
      if (isdir) b += '/';
      b += '\n';
    }
    if (!b.empty()) b.pop_back(); /* trailing newline */
    *style = 2;
    return b;
  }
  /* regular file */
  FILE *f = std::fopen(p.c_str(), "rb");
  if (!f) { *style = 31; return sfmt("[tool error: read %s: %s]", p.c_str(), strerror(errno)); }
  if (offset < 1) offset = 1;
  if (limit <= 0) limit = 2000;
  std::string b;
  char *line_buf = nullptr;
  size_t cap = 0;
  ssize_t n;
  int line_num = 0, read_count = 0;
  while ((n = getline(&line_buf, &cap, f)) >= 0) {
    line_num++;
    if (line_num < offset) continue;
    read_count++;
    size_t len = (size_t)n;
    while (len > 0 && (line_buf[len - 1] == '\n' || line_buf[len - 1] == '\r')) len--;
    std::string_view lview = utf8_truncate(std::string_view(line_buf, len), 2000);
    str_appendf(b, "%d: ", line_num);
    b.append(lview);
    b += '\n';
    if (read_count >= limit) break;
  }
  free(line_buf);
  std::fclose(f);
  if (!b.empty()) b.pop_back(); /* trailing newline */
  *style = 2;
  return b;
}

/* ----------------------------- write ---------------------------------- */
/* Convert a 1-indexed line number to the byte offset of the start of that
 * line within `data`. Line 1 begins at byte 0; a line number past the end of
 * the file maps to data.size() (i.e. append at end of file). */
static size_t line_to_byte_offset(std::string_view data, long line_no) {
  if (line_no <= 1) return 0;
  size_t pos = 0;
  long cur = 1;
  while (pos < data.size()) {
    if (data[pos] == '\n') {
      cur++;
      if (cur == line_no) return pos + 1;
    }
    pos++;
  }
  return data.size();
}

/* Byte length of `nlines` complete lines beginning at byte offset `start`.
 * Counts the newlines terminating each line and stops at end of data. */
static size_t lines_to_byte_len(std::string_view data, size_t start, long nlines) {
  if (nlines <= 0 || start > data.size()) return 0;
  size_t pos = start;
  long consumed = 0;
  while (pos < data.size() && consumed < nlines) {
    if (data[pos] == '\n') consumed++;
    pos++;
  }
  return pos - start;
}

std::string write_file(std::string_view path, long line_offset, long byte_offset,
                       long del_lines, long del_bytes,
                       std::string_view content, int *err) {
  *err = 0;
  std::string p(path);
  size_t s = 0;
  while (s < p.size() && (p[s] == ' ' || p[s] == '\t')) s++;
  p = p.substr(s);
  if (p.empty()) { *err = 1; return "[tool error: path is required]"; }

  std::string data;
  struct stat st;
  if (stat(p.c_str(), &st) == 0) {
    if (S_ISDIR(st.st_mode)) {
      *err = 1;
      return sfmt("[tool error: write %s: Is a directory]", p.c_str());
    }
    if (!read_file_str(p, data)) {
      *err = 1;
      return sfmt("[tool error: write %s: %s]", p.c_str(), strerror(errno));
    }
  } else {
    auto slash = p.find_last_of('/');
    if (slash != std::string::npos) mkdir_p(std::string_view(p).substr(0, slash));
    else mkdir_p(".");
  }

  /* Resolve insertion byte offset. Line-based takes precedence over byte-based
   * ("first and foremost line-based writing"); if neither is given the whole
   * file is overwritten by the caller. */
  size_t ins_off = 0;
  if (line_offset >= 0) {
    ins_off = line_to_byte_offset(data, line_offset);
  } else if (byte_offset >= 0) {
    ins_off = (size_t)byte_offset;
    if (ins_off > data.size()) data.resize(ins_off, '\0');
  } else {
    ins_off = 0;
  }

  /* Resolve deletion span. Line-based takes precedence over byte-based. */
  size_t del_len = 0;
  if (del_lines >= 0) {
    del_len = lines_to_byte_len(data, ins_off, del_lines);
  } else if (del_bytes >= 0) {
    del_len = (size_t)del_bytes;
    if (ins_off + del_len > data.size())
      del_len = (ins_off < data.size()) ? (data.size() - ins_off) : 0;
  }

  size_t end_del = ins_off + del_len;
  size_t prefix_len = ins_off;
  size_t suffix_len = (end_del < data.size()) ? (data.size() - end_del) : 0;
  size_t clen = content.size();
  size_t newlen = prefix_len + clen + suffix_len;
  std::string nd;
  nd.reserve(newlen);
  nd.append(data, 0, prefix_len);
  nd.append(content);
  if (suffix_len) nd.append(data, end_del, suffix_len);
  /* Keep a single trailing newline for line-based / whole-file text writes.
   * Byte-based edits are exact byte operations and must not be mutated with an
   * injected newline, so the guard is skipped whenever a byte offset or byte
   * deletion is actually in effect. */
  bool byte_based = (line_offset < 0 && byte_offset >= 0) ||
                   (del_lines < 0 && del_bytes >= 0);
  if (!byte_based && !nd.empty() && nd.back() != '\n') nd.push_back('\n');
  if (!write_file_str(p, nd)) {
    *err = 1;
    return sfmt("[tool error: write %s: %s]", p.c_str(), strerror(errno));
  }
  return sfmt("Successfully wrote %zu bytes to %s", clen, p.c_str());
}

/* ----------------------------- offload -------------------------------- */
std::string offload_tool_result(std::string_view res, int max_res,
                                 std::vector<std::string> &tmps) {
  if (max_res <= 0) max_res = 65536;
  if (res.size() <= (size_t)max_res) return std::string(res);
  std::string dir = tool_res_dir();
  std::string tpl = dir + "/toolres-XXXXXX";
  int fd = mkstemp(tpl.data());
  if (fd < 0) return std::string(res);
  size_t written = 0;
  bool write_ok = true;
  while (written < res.size()) {
    ssize_t w = write(fd, res.data() + written, res.size() - written);
    if (w < 0) {
      if (errno == EINTR) continue;
      write_ok = false;
      break;
    }
    written += (size_t)w;
  }
  close(fd);
  if (!write_ok) { unlink(tpl.c_str()); return std::string(res); }
  tmps.push_back(tpl);
  return sfmt("[tool result too large for context: %zu bytes (limit %d). The full output was saved to %s. Read it with the read tool; to read it partially, use offset and limit. The file is %zu bytes long.]",
              res.size(), max_res, tpl.c_str(), res.size());
}
void cleanup_temps(std::vector<std::string> &tmps) {
  for (auto &t : tmps) unlink(t.c_str());
  tmps.clear();
}