burke/zeus

package zeusclient

import (
    "io"
    "net"
    "os"
    "os/signal"
    "strconv"
    "strings"
    "sync"
    "syscall"

    "github.com/burke/ttyutils"
    "github.com/burke/zeus/go/messages"
    slog "github.com/burke/zeus/go/shinylog"
    "github.com/burke/zeus/go/unixsocket"
    "github.com/burke/zeus/go/zerror"
    "github.com/kr/pty"
)

const (
    sigInt  = 3
    sigQuit = 28
    sigTstp = 26
)

// man signal | grep 'terminate process' | awk '{print $2}' | xargs -I '{}' echo -n "syscall.{}, "
var terminatingSignals = []os.Signal{syscall.SIGHUP, syscall.SIGINT, syscall.SIGKILL, syscall.SIGPIPE, syscall.SIGALRM, syscall.SIGTERM, syscall.SIGXCPU, syscall.SIGXFSZ, syscall.SIGVTALRM, syscall.SIGPROF, syscall.SIGUSR1, syscall.SIGUSR2}

// Run runs the main Zeus command.
func Run(args []string, input io.Reader, output *os.File, stderr *os.File) int {
    if os.Getenv("RAILS_ENV") != "" {
        println("Warning: Specifying a Rails environment via RAILS_ENV has no effect for commands run with zeus.")
        println("As a safety precaution to protect you from nuking your development database,")
        println("Zeus will now cowardly refuse to proceed. Please unset RAILS_ENV and try again.")
        return 1
    }

    // setup stdout
    localStdout, remoteStdout, outputIsTerminal, err := socketsForOutput(output)
    if err != nil {
        slog.ErrorString(err.Error() + "\r")
        return 1
    }
    defer localStdout.Close()
    defer remoteStdout.Close()

    // setup stderr
    localStderr, remoteStderr, _, err := socketsForOutput(stderr)
    if err != nil {
        slog.ErrorString(err.Error() + "\r")
        return 1
    }
    defer localStderr.Close()
    defer remoteStderr.Close()

    addr, err := net.ResolveUnixAddr("unixgram", unixsocket.ZeusSockName())
    if err != nil {
        slog.ErrorString(err.Error() + "\r")
        return 1
    }

    conn, err := net.DialUnix("unix", nil, addr)
    if err != nil {
        zerror.ErrorCantConnectToMaster()
        return 1
    }
    usock := unixsocket.New(conn)

    msg := messages.CreateCommandAndArgumentsMessage(args, os.Getpid())
    usock.WriteMessage(msg)
    err = sendCommandLineArguments(usock, args)
    if err != nil {
        slog.ErrorString(err.Error() + "\r")
        return 1
    }

    usock.WriteFD(int(remoteStdout.Fd()))
    usock.WriteFD(int(remoteStderr.Fd()))

    msg, err = usock.ReadMessage()
    if err != nil {
        slog.ErrorString(err.Error() + "\r")
        return 1
    }

    parts := strings.Split(msg, "\000")
    commandPid, err := strconv.Atoi(parts[0])
    defer func() {
        if commandPid > 0 {
            // Just in case.
            syscall.Kill(commandPid, 9)
        }
    }()

    if err != nil {
        slog.ErrorString(err.Error() + "\r")
        return 1
    }

    if outputIsTerminal {
        c := make(chan os.Signal, 1)
        handledSignals := append(append(terminatingSignals, syscall.SIGWINCH), syscall.SIGCONT)
        signal.Notify(c, handledSignals...)
        go func() {
            for sig := range c {
                if sig == syscall.SIGCONT {
                    syscall.Kill(commandPid, syscall.SIGCONT)
                } else if sig == syscall.SIGWINCH {
                    ttyutils.MirrorWinsize(output, localStdout)
                    syscall.Kill(commandPid, syscall.SIGWINCH)
                } else { // member of terminatingSignals
                    print("\r")
                    syscall.Kill(commandPid, sig.(syscall.Signal))
                    os.Exit(1)
                }
            }
        }()
    }

    exitStatus := -1
    if len(parts) > 2 {
        exitStatus, err = strconv.Atoi(parts[0])
        if err != nil {
            slog.ErrorString(err.Error() + "\r")
            return 1
        }
    }

    var endOfIO sync.WaitGroup

    oldTerminalStateStdout, err := forwardOutput(localStdout, output, endOfIO)
    if oldTerminalStateStdout != nil {
        defer ttyutils.RestoreTerminalState(output.Fd(), oldTerminalStateStdout)
    }
    if err != nil {
        slog.ErrorString(err.Error() + "\r")
        return 1
    }

    oldTerminalStateStderr, err := forwardOutput(localStderr, stderr, endOfIO)
    if oldTerminalStateStderr != nil {
        defer ttyutils.RestoreTerminalState(stderr.Fd(), oldTerminalStateStderr)
    }
    if err != nil {
        slog.ErrorString(err.Error() + "\r")
        return 1
    }

    go func() {
        endOfIO.Add(1)
        buf := make([]byte, 8192)
        for {
            n, err := input.Read(buf)
            if err != nil {
                endOfIO.Done()
                break
            }
            if outputIsTerminal {
                for i := 0; i < n; i++ {
                    switch buf[i] {
                    case sigInt:
                        syscall.Kill(commandPid, syscall.SIGINT)
                    case sigQuit:
                        syscall.Kill(commandPid, syscall.SIGQUIT)
                    case sigTstp:
                        syscall.Kill(commandPid, syscall.SIGTSTP)
                        syscall.Kill(os.Getpid(), syscall.SIGTSTP)
                    }
                }
            }
            localStdout.Write(buf[:n])
        }
    }()

    endOfIO.Wait()

    if exitStatus == -1 {
        msg, err = usock.ReadMessage()
        if err != nil {
            slog.ErrorString(err.Error() + "\r")
            return 1
        }
        parts := strings.Split(msg, "\000")
        exitStatus, err = strconv.Atoi(parts[0])
        if err != nil {
            slog.ErrorString(err.Error() + "\r")
            return 1
        }
    }

    return exitStatus
}

func sendCommandLineArguments(usock *unixsocket.Usock, args []string) error {
    master, slave, err := unixsocket.Socketpair(syscall.SOCK_STREAM)
    if err != nil {
        return err
    }
    usock.WriteFD(int(slave.Fd()))
    if err != nil {
        return err
    }
    slave.Close()

    go func() {
        defer master.Close()
        argAsBytes := []byte{}
        for _, arg := range args[1:] {
            argAsBytes = append(argAsBytes, []byte(arg)...)
            argAsBytes = append(argAsBytes, byte(0))
        }
        _, err = master.Write(argAsBytes)
        if err != nil {
            slog.ErrorString("Could not send arguments across: " +
                err.Error() + "\r")
            return
        }
    }()

    return nil
}

func socketsForOutput(out *os.File) (local, remote *os.File, outIsTerminal bool, err error) {
    outIsTerminal = ttyutils.IsTerminal(out.Fd())

    if outIsTerminal {
        local, remote, err = pty.Open()
    } else {
        local, remote, err = unixsocket.Socketpair(syscall.SOCK_STREAM)
    }

    return
}

func forwardOutput(from, to *os.File, signalEnd sync.WaitGroup) (oldTerminalState *ttyutils.Termios, err error) {
    ttyutils.MirrorWinsize(to, from)

    if ttyutils.IsTerminal(to.Fd()) {
        oldTerminalState, err = ttyutils.MakeTerminalRaw(to.Fd())
        if err != nil {
            return
        }
    }

    go func() {
        signalEnd.Add(1)
        for {
            buf := make([]byte, 1024)
            n, err := from.Read(buf)

            if err == nil || (err == io.EOF && n > 0) {
                to.Write(buf[:n])
            } else {
                signalEnd.Done()
                break
            }
        }
    }()

    return
}