arsh/main.c

#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <stdint.h>
#include <stdio.h>
#include <sys/wait.h>
#include <signal.h>
#include <setjmp.h>

constexpr size_t max_lenght = 1024;
const char* token_separators = " \t\n\r";

static sigjmp_buf restart_jmp_buf;
static constexpr int restart_value = 69;
static volatile sig_atomic_t jump_active = 0;

void sigint_handler(int signo)
{
  if(!jump_active)
  {
    return;
  }
  siglongjmp(restart_jmp_buf, restart_value);
}

/*Checked calloc: if allocation fails -> exists*/
void* ccalloc(size_t size, size_t n)
{
  void* ret = calloc(size, n);
  if(!ret)
  {
    perror("calloc failed");
    exit(1);
  }
  return ret;
}

static constexpr size_t max_line = 1024;
static char string_buffer[max_line];
char* readline(char* print, FILE* stream)
{
  if(print)
  {
    printf("%s", print);
  }
  if(fgets(string_buffer, max_line, stream) == NULL)
  {
    if (feof(stdin))
    {
        return NULL;   // Ctrl+D
    }
    if(ferror(stream))
    {
      perror("fgets error");
      clearerr(stdin);
      return NULL;
    }
    return NULL;
  }
  
  size_t length = strlen(string_buffer);
  //Remove trailing \n
  if (length > 0 && string_buffer[length - 1] == '\n')
  {
    string_buffer[length - 1] = '\0';
  }
  return string_buffer;
}

/**
 * Struct to represent a command and its arguments.
 */
struct command
{
  /**
   * IO redirections; redirect[i] should be used as fd i in the child.
   * A value of -1 indicates no redirect.
   */
  int redirect[2];
  /** The arguments; must be NULL-terminated. */
  char* argv[];
};
/** The name of the executable. */
char* command_name(struct command* cmd)
{
  return cmd->argv[0];
}

/**
 * Struct to represent a pipeline of commands. The intention is that cmd[i]'s
 * output goes to cmd[i+1]'s input.
 */
struct pipeline
{
  /** The total number of commands. */
  size_t n_cmds;
  struct command* cmds[];
};

void* parser_memory[max_lenght] = {0};
size_t parser_memory_index = 0;
void parser_allocated(void* ptr)
{
  parser_memory[parser_memory_index] = ptr;
  if(parser_memory_index == max_lenght)
  {
    fprintf(stderr, "Not enough memory for parser: Memory will leak!");
    exit(1);
  }
  parser_memory_index++;
}

void parser_free()
{
  for(size_t i = 0; i < parser_memory_index; i++)
  {
    free(parser_memory[i]);
  }
  memset(parser_memory, 0, max_lenght);
  parser_memory_index = 0;
}

/*
 * For internal use. Returns the next non-empty token according in *line,
 * splitting on characters in TOKEN_SEP. Returns NULL if no token remains.
 * Updates line to point to remainder after removal of the token.  Modifies the
 * string *line.
 */
char* next_non_empty(char **line) 
{
  char *tok;

  /* Consume empty tokens. */
  while ((tok = strsep(line, token_separators)) && !*tok);

  return tok;
}

/**
 * Parses str into a freshly allocated command struct and returns a pointer to it.
 * The redirects in the returned command will be set to -1, ie no redirect.
 */
struct command* parse_command(char* str) 
{
  /* Copy the input line in case the caller wants it later. */
  char* copy = strndup(str, max_lenght);
  parser_allocated(copy);
  char* token;
  int i = 0;

  /*
   * Being lazy and allocating way too much memory for the args array.
   * Using calloc to ensure it's zero-initialised, which is important because
   * execvp expects a NULL-terminated array of arguments.
   */
  struct command* ret = ccalloc(sizeof(struct command) + strlen(copy) * sizeof(char*), 1);
  parser_allocated(ret);

  while ((token = next_non_empty(&copy))) 
  {
    ret->argv[i++] = token;
  }
  ret->redirect[0] = ret->redirect[1] = -1;
  return ret;
}

/**
 * Parses str into a freshly allocated pipeline_struct and returns a pointer to
 * it.  All commands in cmds will also be freshy allocated, and have their
 * redirects set to -1, ie no redirect.
 */
struct pipeline* parse_pipeline(char *str)
{
  char* copy = strndup(str, max_lenght);
  parser_allocated(copy);
  char* cmd_str;
  int n_cmds = 0;
  int i = 0;
  struct pipeline* ret;

  /*
   * Count pipe characters that appear in pipeline to know how much space to
   * allocate for the cmds array.
   */
  for (char* cur = str; *cur; cur++)
  {
    if (*cur == '|') ++n_cmds;
  }

  ++n_cmds;  /* There is one more command than there are pipe characters. */

  ret = ccalloc(sizeof(struct pipeline) + n_cmds * sizeof(struct command*), 1);
  parser_allocated(ret);
  ret->n_cmds = n_cmds;

  while((cmd_str = strsep(&str, "|")))
  {
    ret->cmds[i++] = parse_command(cmd_str);
  }

  return ret;
}

void close_ALL_the_pipes(int n_pipes, int (*pipes)[2])
{
  for (int i = 0; i < n_pipes; ++i)
  {
    close(pipes[i][0]);
    close(pipes[i][1]);
  }
}

int exec_with_redir(struct command* command, int n_pipes, int (*pipes)[2])
{
  int fd = -1;
  if ((fd = command->redirect[0]) != -1)
  {
    dup2(fd, STDIN_FILENO);
  }
  if ((fd = command->redirect[1]) != -1)
  {
    dup2(fd, STDOUT_FILENO);
  }
  close_ALL_the_pipes(n_pipes, pipes);
  return execvp(command_name(command), command->argv);
}

static struct sigaction s_old;

pid_t run_with_redir(struct command* command, int n_pipes, int (*pipes)[2])
{
  pid_t child_pid = fork();
  
  if(child_pid < 0)
  {
    perror("Fork failed");
    exit(1);
  }
  if (child_pid) 
  {  /* We are the parent. */
      return child_pid;
  } 
  else 
  {  // We are the child. */
    sigaction(SIGINT, &s_old, NULL);
    exec_with_redir(command, n_pipes, pipes);
    perror(command_name(command));
    exit(1);
  }
}

int cd(char* path)
{
  return chdir(path);
}

int main(int argc, const char* argv[])
{
  FILE* stream = stdin;
  char* prompt = "arsh> ";
  if(argc == 2)
  {
    FILE* file = fopen(argv[1], "r");
    if(file)
    {
      stream = file;
      prompt = NULL;
    }
  }

  struct sigaction s;
  s.sa_handler = sigint_handler;
  sigemptyset(&s.sa_mask);
  s.sa_flags = SA_RESTART;
  sigaction(SIGINT, &s, &s_old);

  int stat_loc;

  if(sigsetjmp(restart_jmp_buf, 1) == restart_value)
  {
    printf("\n");
  }
  jump_active = 1;
  
  while(true)
  {
    char* input = readline(prompt, stream);
    if(input == NULL) //CTRL + D
    {
      if(prompt)
      {
	printf("\nexit\n");
      }
      exit(0);
    }
    
    // After that line you need cleanup
    struct pipeline* pipeline = parse_pipeline(input);
    int n_pipes = pipeline->n_cmds - 1;
    
    /* pipes[i] redirects from pipeline->cmds[i] to pipeline->cmds[i+1]. */
    int (*pipes)[2] = ccalloc(sizeof(int[2]), n_pipes);
    parser_allocated(pipes);
    pid_t* pids = ccalloc(sizeof(pid_t), pipeline->n_cmds);
    parser_allocated(pids);
    for (int i = 1; i < pipeline->n_cmds; ++i)
    {
      pipe(pipes[i-1]);  
      pipeline->cmds[i]->redirect[STDIN_FILENO] = pipes[i-1][0];
      pipeline->cmds[i-1]->redirect[STDOUT_FILENO] = pipes[i-1][1];
    }
    for (int i = 0; i < pipeline->n_cmds; ++i)
    {
      struct command* cmd = pipeline->cmds[i];
      char* cmd_name = command_name(cmd);
      if(!cmd_name)
      {
	continue;
      }
      if(strlen(cmd_name) >= 1 && cmd_name[0] == '#')
      {
	goto cleanup;
      }
      
      if(strcmp(cmd_name, "cd") == 0)
      {
	char* path = cmd->argv[1] ? cmd->argv[1] : getenv("HOME");
	if (cd(path) < 0)
	{
	  char err_buf[max_lenght] = {0};
	  sprintf(err_buf, "cd: %s", path);
	  perror(err_buf);
	}

	goto cleanup; //skip fork()
      }
      if(strcmp(cmd_name, "exit") == 0)
      {
	int32_t status_code = cmd->argv[1] != NULL ? atoi(cmd->argv[1]) : 0;
	exit(status_code);
      }
      pids[i] = run_with_redir(cmd, n_pipes, pipes);
    }
    close_ALL_the_pipes(n_pipes, pipes);

    for(size_t i = 0; i < pipeline->n_cmds; i++)
    {
      waitpid(pids[i], &stat_loc, WUNTRACED);
    }

cleanup:

    parser_free();
  }
}