Luna/shell/main.cpp

#include "sh.h"
#include <errno.h>
#include <luna/String.h>
#include <luna/Vector.h>
#include <os/ArgumentParser.h>
#include <os/File.h>
#include <os/FileSystem.h>
#include <os/Process.h>
#include <pwd.h>
#include <signal.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/utsname.h>
#include <termios.h>
#include <unistd.h>

#include <luna/HashMap.h>

extern shell_builtin_t builtin_cd;
extern shell_builtin_t builtin_exit;
extern shell_builtin_t builtin_set;
extern shell_builtin_t builtin_unset;

static HashMap<StringView, shell_builtin_t> s_builtins;

using os::File;

static Result<Vector<String>> split_command_into_args(StringView cmd)
{
    return cmd.split(" \n"_sv);
}

static Result<void> execute_command(StringView command)
{
    if (strcspn(command.chars(), " #") == 0) return {};

    auto args = TRY(split_command_into_args(command));
    if (args.size() < 1) exit(0);

    return os::Process::exec(args[0].view(), args.slice());
}

static Result<void> init_builtins()
{
    s_builtins = {};

    TRY(s_builtins.try_set("cd"_sv, builtin_cd));
    TRY(s_builtins.try_set("exit"_sv, builtin_exit));
    TRY(s_builtins.try_set("set"_sv, builtin_set));
    TRY(s_builtins.try_set("unset"_sv, builtin_unset));

    return {};
}

static Result<void> builtin_wrapper(const Vector<String>& args, shell_builtin_t builtin)
{
    Vector<const char*> argv;
    for (const auto& arg : args) { TRY(argv.try_append(arg.chars())); }
    TRY(argv.try_append(nullptr));

    auto rc = builtin((int)args.size(), const_cast<char**>(argv.data()));
    if (rc.has_error())
    {
        errno = rc.error();
        perror(argv[0]);
    }

    return {};
}

// Do nothing, but we need a handler so read() returns EINTR.
static void sigint_handler(int)
{
}

struct utsname g_sysinfo;

const char* hostname = "";
const char* username = "";
char prompt_end = '$';

Result<int> luna_main(int argc, char** argv)
{
    StringView path;
    StringView command;
    bool interactive { false };

    SharedPtr<File> input_file;

    os::ArgumentParser parser;
    parser.add_description("The Luna system's command shell."_sv);
    parser.add_system_program_info("sh"_sv);
    parser.add_positional_argument(path, "path"_sv, "-"_sv);
    parser.add_value_argument(command, 'c', "command"_sv, "execute a single command and then exit"_sv);
    parser.add_switch_argument(interactive, 'i', "interactive"_sv, "run an interactive shell"_sv);
    parser.parse(argc, argv);

    // TODO: This does not properly handle builtins.
    if (!command.is_empty()) TRY(execute_command(command));

    if (path == "-") { input_file = File::standard_input(); }
    else
    {
        input_file = TRY(File::open(path, File::ReadOnly));
        input_file->set_close_on_exec();
    }

    if (isatty(input_file->fd())) interactive = true;

    if (interactive)
    {
        // Set up everything to form a prompt.
        uname(&g_sysinfo);
        hostname = g_sysinfo.nodename;

        if (getuid() == 0) prompt_end = '#';

        struct passwd* pw = getpwuid(getuid());
        if (pw) { username = pw->pw_name; }
        else { username = getenv("USER"); }
        endpwent();

        signal(SIGTTOU, SIG_IGN);
        signal(SIGINT, sigint_handler);

        tcsetpgrp(input_file->fd(), getpgid(0));
    }

    TRY(init_builtins());

    while (1)
    {
        if (interactive)
        {
            auto cwd = TRY(os::FileSystem::working_directory());
            os::print("\x1b[%dm%s\x1b[m@\x1b[36m%s\x1b[m:\x1b[1;34m%s\x1b[m%c ", getuid() == 0 ? 31 : 35, username,
                      hostname, cwd.chars(), prompt_end);
        }

        auto maybe_cmd = input_file->read_line();
        if (maybe_cmd.has_error())
        {
            if (maybe_cmd.error() == EINTR)
            {
                os::print("\n");
                continue;
            }
            return maybe_cmd.release_error();
        }

        auto cmd = maybe_cmd.release_value();
        if (cmd.is_empty())
        {
            if (interactive) puts("exit");
            break;
        }

        if (strspn(cmd.chars(), " \n") == cmd.length()) continue;

        if (strcspn(cmd.chars(), " #") == 0) continue;

        auto args = TRY(split_command_into_args(cmd.view()));
        if (args.size() < 1) continue;

        if (!strchr(args[0].chars(), '/'))
        {
            auto maybe_builtin = s_builtins.try_get(args[0].view());

            if (maybe_builtin.has_value())
            {
                auto builtin = maybe_builtin.value();

                TRY(builtin_wrapper(args, builtin));

                continue;
            }
        }

        pid_t child = TRY(os::Process::fork());

        if (child == 0)
        {
            if (interactive)
            {
                setpgid(0, 0);
                tcsetpgrp(input_file->fd(), getpid());
            }
            TRY(os::Process::exec(args[0].view(), args.slice()));
        }

        int status;
        TRY(os::Process::wait(child, &status));

        if (interactive) tcsetpgrp(input_file->fd(), getpgid(0));

        if (WIFSIGNALED(status))
        {
            int sig = WTERMSIG(status);
            if (sig != SIGINT && sig != SIGQUIT) os::println("[sh] Process %d exited: %s", child, strsignal(sig));
            else
                os::print("\n");
        }
    }

    return 0;
}