job_control_pipe_shell.c

#include <assert.h>
#include <signal.h>
#include <stdbool.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>

#define VECTOR(T)                                                              \
  typedef struct {                                                             \
    int size;                                                                  \
    int capacity;                                                              \
    T *items;                                                                  \
  } T##_vector;                                                                \
  T##_vector *alloc_##T##_vector(int n) {                                      \
    T##_vector *iv = (T##_vector *)malloc(sizeof(T##_vector));                 \
    if (!iv) {                                                                 \
      perror("malloc");                                                        \
      exit(EXIT_FAILURE);                                                      \
    }                                                                          \
    iv->size = 0;                                                              \
    iv->capacity = n;                                                          \
    iv->items = (T *)malloc(n * sizeof(T));                                    \
    if (!iv->items) {                                                          \
      perror("malloc");                                                        \
      exit(EXIT_FAILURE);                                                      \
    }                                                                          \
    return iv;                                                                 \
  }                                                                            \
  static void grow_##T##_vector(T##_vector *iv) {                              \
    int new_cap = iv->capacity * 2;                                            \
    T *new_items = (T *)realloc(iv->items, new_cap * sizeof(T));               \
    if (!new_items) {                                                          \
      perror("realloc");                                                       \
      exit(EXIT_FAILURE);                                                      \
    }                                                                          \
    iv->capacity = new_cap;                                                    \
    iv->items = new_items;                                                     \
  }                                                                            \
  T *get_##T##_vector(T##_vector *iv, int i) {                                 \
    assert(i >= 0 && i < iv->capacity);                                        \
    if (i >= iv->size)                                                         \
      iv->size = i + 1;                                                        \
    return &iv->items[i];                                                      \
  }                                                                            \
  void append_##T##_vector(T##_vector *iv, T item) {                           \
    if (iv->size == iv->capacity)                                              \
      grow_##T##_vector(iv);                                                   \
    iv->items[iv->size++] = item;                                              \
  }                                                                            \
  void remove_##T##_vector(T##_vector *iv, int i) {                            \
    assert(i >= 0 && i < iv->size);                                            \
    memmove(&iv->items[i], &iv->items[i + 1], (--iv->size - i) * sizeof(T));   \
  }                                                                            \
  void clear_##T##_vector(T##_vector *iv) { iv->size = 0; }                    \
  void free_##T##_vector(T##_vector *iv) {                                     \
    free(iv->items);                                                           \
    free(iv);                                                                  \
  }

typedef struct {
  pid_t pid;
  char **argv;
  int in_fd;
  int out_fd;
} job;
VECTOR(job)

typedef struct {
  pid_t pgid;
  char *cmd;
  char **argv;
  job_vector *jv;
  bool bg;
} job_group;
VECTOR(job_group)

typedef struct {
  int idx;
} free_idx;
VECTOR(free_idx)

void free_job_group(job_group *jgptr) {
  free(jgptr->argv);
  free(jgptr->cmd);
  free_job_vector(jgptr->jv);
}

#define PROMPT "# "
job_group_vector *gv;
free_idx_vector *fv;
bool sigchld = false;

int wait_job_group(job_group *job_group) {
  int status;
  if (tcsetpgrp(STDIN_FILENO, job_group->pgid) == -1) {
    perror("tcsetpgrp");
    exit(EXIT_FAILURE);
  }
  for (int i = 0; i < job_group->jv->size; i++) {
    job *job = get_job_vector(job_group->jv, i);
    if (waitpid(job->pid, &status, job_group->bg ? WNOHANG : WUNTRACED) > 0) {
      if (WIFSTOPPED(status))
        job_group->bg = true;
      else if (WIFEXITED(status) || WIFSIGNALED(status))
        remove_job_vector(job_group->jv, i--);
    }
  }
  if (tcsetpgrp(STDIN_FILENO, getpgrp()) == -1) {
    perror("tcsetpgrp");
    exit(EXIT_FAILURE);
  }
  return job_group->jv->size == 0;
}

void getline_job_group(job_group *job_group, char *lineptr, size_t n) {
  char *token;

  job_group->cmd = lineptr;
  int tokens = 0;

  while ((token = strsep(&lineptr, " \t\n")))
    if (*token)
      tokens++;

  job_group->argv = (char **)malloc((tokens + 1) * sizeof(char *));
  if (!job_group->argv) {
    perror("malloc");
    exit(EXIT_FAILURE);
  }

  job tmp_job;
  int pfds[2];
  tmp_job.pid = 0;
  tmp_job.argv = job_group->argv;
  tmp_job.in_fd = STDIN_FILENO;
  tmp_job.out_fd = STDOUT_FILENO;
  lineptr = job_group->cmd;

  for (int i = 0; i < tokens; i++) {
    while (*lineptr == '\0')
      lineptr++;

    if (strcmp("|", lineptr) == 0) {
      if (pipe(pfds) == -1) {
        perror("pipe");
        exit(EXIT_FAILURE);
      }

      job_group->argv[i] = NULL;
      tmp_job.out_fd = pfds[1];
      append_job_vector(job_group->jv, tmp_job);
      // setup the next job
      tmp_job.argv = job_group->argv + i + 1;
      tmp_job.in_fd = pfds[0];
      tmp_job.out_fd = STDOUT_FILENO;
    } else if (strcmp("&", lineptr) == 0) {
      job_group->argv[i] = NULL;
      job_group->bg = true;
    } else
      job_group->argv[i] = lineptr;

    lineptr += strlen(lineptr);
  }
  job_group->argv[tokens] = NULL;
  tmp_job.out_fd = STDOUT_FILENO;
  append_job_vector(job_group->jv, tmp_job);
}

void execute_job(job *jb, pid_t pgid) {
  pid_t pid = fork();
  if (pid == -1) {
    perror("fork");
    exit(EXIT_FAILURE);
  }

  if (pid == 0) {
    signal(SIGINT, SIG_DFL);
    signal(SIGQUIT, SIG_DFL);
    signal(SIGTSTP, SIG_DFL);

    if (jb->in_fd != STDIN_FILENO) {
      dup2(jb->in_fd, STDIN_FILENO);
      close(jb->in_fd);
    }

    if (jb->out_fd != STDOUT_FILENO) {
      dup2(jb->out_fd, STDOUT_FILENO);
      close(jb->out_fd);
    }

    if (setpgid(0, pgid) == -1) {
      perror("setpgid");
      _exit(EXIT_FAILURE);
    }
    execvp(jb->argv[0], jb->argv);
    perror("execvp");
    _exit(EXIT_FAILURE);
  }

  if (setpgid(pid, pgid ? pgid : pid) == -1) {
    perror("setpgid");
    exit(EXIT_FAILURE);
  }
  if (jb->in_fd != STDIN_FILENO)
    close(jb->in_fd);

  if (jb->out_fd != STDOUT_FILENO)
    close(jb->out_fd);

  jb->pid = pid;
}

void execute_job_group(job_group *job_group) {
  if (job_group->jv->size < 1)
    return;

  job *first_job = get_job_vector(job_group->jv, 0);
  execute_job(first_job, 0); // set the pgid to the pid of the child
  job_group->pgid = first_job->pid;

  for (int i = 1; i < job_group->jv->size; i++)
    execute_job(get_job_vector(job_group->jv, i), job_group->pgid);
}

void sigchld_function() {
  int status, pid;
found:
  while ((pid = waitpid(-1, &status, WNOHANG)) > 0) {
    if (WIFEXITED(status) || WIFSIGNALED(status)) {
      if (WTERMSIG(status) == SIGSEGV)
        fprintf(stderr, "Segmentation fault\n");
      for (int i = 0; i < gv->size; i++) {
        job_group *jgptr = get_job_group_vector(gv, i);
        if (jgptr->pgid) {
          for (int j = 0; j < jgptr->jv->size; j++) {
            job *job = get_job_vector(jgptr->jv, j);
            // if the job is found, remove the whole group
            if (job->pid == pid) {
              kill(-jgptr->pgid, SIGCONT); // not needed, choosing to resume the
                                           // group if one child exits
              free_job_group(jgptr);
              jgptr->pgid = 0;
              append_free_idx_vector(fv, (free_idx){i});
              goto found;
            }
          }
        }
      }
    }
  }
}

void sigchld_handler(int sig) { sigchld = true; }

int main(int argc, char *argv[]) {
  FILE *fp = stdin;
  if (argc > 1) {
    fp = fopen(argv[1], "r");
    if (!fp) {
      perror("fopen");
      exit(EXIT_FAILURE);
    }
  }

  gv = alloc_job_group_vector(1);
  fv = alloc_free_idx_vector(1);

  if (signal(SIGINT, SIG_IGN) == SIG_ERR ||
      signal(SIGQUIT, SIG_IGN) == SIG_ERR ||
      signal(SIGTSTP, SIG_IGN) == SIG_ERR ||
      signal(SIGTTIN, SIG_IGN) == SIG_ERR ||
      signal(SIGTTOU, SIG_IGN) == SIG_ERR ||
      signal(SIGCHLD, sigchld_handler) == SIG_ERR) {
    perror("signal");
    exit(EXIT_FAILURE);
  }

  for (;;) {
    job_group job_group, *jgptr;
    job_group.bg = false;
    job_group.jv = alloc_job_vector(1);
    job_group.pgid = 0;

    if (fp == stdin)
      if (write(STDOUT_FILENO, PROMPT, sizeof(PROMPT)) == -1)
        perror("write");

    char *lineptr = NULL;
    size_t n = 0;
    if (getline(&lineptr, &n, fp) == -1) {
      free(lineptr);
      free_job_vector(job_group.jv);
      break;
    }
    if (sigchld) {
      sigchld_function();
      sigchld = false;
    }

    getline_job_group(&job_group, lineptr, n);
    bool bg;

    if (job_group.argv[0] == NULL) {
    } else if (strcmp(job_group.argv[0], "cd") == 0) {
      if (job_group.argv[1] == NULL) {
        if (chdir(getenv("HOME")) == -1)
          perror("chdir");
      } else if (chdir(job_group.argv[1]) == -1)
        perror("chdir");
    } else if (strcmp(job_group.argv[0], "exit") == 0) {
      free_job_group(&job_group);
      break;
    } else if (strcmp(job_group.argv[0], "jobs") == 0) {
      for (int i = 0; i < gv->size; i++) {
        jgptr = get_job_group_vector(gv, i);
        if (jgptr->pgid) {
          printf("%d:", i + 1);
          for (int j = 0; j < jgptr->jv->size; j++) {
            job *job = get_job_vector(jgptr->jv, j);
            for (int k = 0; job->argv[k] != NULL; k++)
              printf(" %s", job->argv[k]);
            if (job->out_fd != STDOUT_FILENO)
              printf(" |");
          }
          printf("\n");
        }
      }
    } else if ((bg = strcmp(job_group.argv[0], "bg") == 0) ||
               strcmp(job_group.argv[0], "fg") == 0) {
      char *endptr;
      int job_id;
      if (job_group.argv[1] != NULL) {
        job_id = strtol(job_group.argv[1], &endptr, 10);
        if (*endptr != '\0')
          fprintf(stderr, "no number in argument provided\n");
      } else {
        for (int i = gv->size - 1; i >= 0; i--) {
          jgptr = get_job_group_vector(gv, i);
          if (jgptr->pgid) {
            job_id = i + 1;
            break;
          }
        }
      }

      if (job_id > 0 && job_id <= gv->size) {
        jgptr = get_job_group_vector(gv, job_id - 1);
        if (jgptr->pgid) {
          jgptr->bg = bg;
          if (kill(-jgptr->pgid, SIGCONT) == -1)
            perror("kill");
          else if (wait_job_group(jgptr)) {
            free_job_group(jgptr);
            jgptr->pgid = 0;
            append_free_idx_vector(fv, (free_idx){job_id - 1});
          }
        }
      } else
        fprintf(stderr, "no job %d\n", job_id);
    } else
      goto not_builtin;

    free_job_group(&job_group);
    continue;

  not_builtin:
    execute_job_group(&job_group);
    if (wait_job_group(&job_group))
      free_job_group(&job_group);
    else {
      if (fv->size > 0) {
        free_idx *fptr = get_free_idx_vector(fv, 0);
        int job_group_idx = fptr->idx, fptr_idx = 0;
        for (int i = 1; i < fv->size; i++) {
          fptr = get_free_idx_vector(fv, i);
          if (fptr->idx < job_group_idx) {
            job_group_idx = fptr->idx;
            fptr_idx = i;
          }
        }
        remove_free_idx_vector(fv, fptr_idx);
        jgptr = get_job_group_vector(gv, job_group_idx);
        assert(jgptr->pgid == 0);
        *jgptr = job_group;
      } else
        append_job_group_vector(gv, job_group);
    }
  }

  for (int i = 0; i < gv->size; i++) {
    job_group *jgptr = get_job_group_vector(gv, i);
    if (jgptr->pgid)
      free_job_group(jgptr);
  }
  free_job_group_vector(gv);

  return 0;
}