iris-context/root/root.go

package root

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"syscall"

	"github.com/spf13/cobra"
	_ "github.com/versenilvis/iris/commands"
	"github.com/versenilvis/iris/internal/config"
	"github.com/versenilvis/iris/internal/logger"
	"golang.org/x/term"
)

var (
	rootCmd = &cobra.Command{
		Use:   "iris",
		Short: "IRIS is an awesome cli auto-completion tool",
		Long: `IRIS (a.k.a Intelligent Real-time Input Suggestion) is a shell auto-autocompletion tool.
It works exactly like coding editor suggestion menu drop down.`,
		Run: func(cmd *cobra.Command, args []string) {
			defer func() {
				if r := recover(); r != nil {
					WriteCrashLog(r)
					restoreTerminal()
					printCrashNotice()
					startRescueShell()
					os.Exit(2)
				}
			}()
			if pidStr := os.Getenv("IRIS_PID"); pidStr != "" {
				if pid, err := strconv.Atoi(pidStr); err == nil && pid > 0 {
					if logDir, err := config.CachePath(); err == nil {
						argsFile := filepath.Join(logDir, "reload-args")
						_ = os.WriteFile(argsFile, []byte(strings.Join(os.Args[1:], "\n")), 0600)
					}
					_ = syscall.Kill(pid, syscall.SIGUSR1)
					fmt.Println("\r\033[K\033[36m[IRIS] Sent reload signal to parent session.\033[0m")
					return
				}
			}
			runWrapper()
		},
	}
	shellFlag      string
	shellLoginFlag bool
	debugMode      bool
)

func init() {
	rootCmd.PersistentFlags().StringVarP(&shellFlag, "shell", "s", "", "shell to use (bash, zsh, fish)")
	rootCmd.PersistentFlags().BoolVar(&shellLoginFlag, "shell-login", false, "run the selected shell as a login shell")
	rootCmd.PersistentFlags().BoolVarP(&debugMode, "debug", "d", false, "enable debug logging to iris.log")

	rootCmd.PersistentPreRun = func(cmd *cobra.Command, args []string) {
		if shellFlag != "" {
			config.Get().Core.Shell = shellFlag
		}
		if cmd.Flags().Changed("shell-login") {
			config.Get().Core.ShellLogin = shellLoginFlag
		}
		logDir, err := config.CachePath()
		if err == nil {
			logger.Init(filepath.Join(logDir, "iris.log"), debugMode || config.Get().Core.Debug)
			logger.Infof("IRIS session started: os=%s, arch=%s, go=%s, pid=%d", runtime.GOOS, runtime.GOARCH, runtime.Version(), os.Getpid())
			cfg := config.Get()
			logger.Debugf("IRIS loaded config: shell=%q, shell-login=%v, mode=%q, ghost-text=%v, max-suggestions=%d", cfg.Core.Shell, cfg.Core.ShellLogin, cfg.Core.Mode, cfg.UI.GhostText, cfg.UI.MaxSuggestions)
		}
	}
}

// runWatchdog spawns the watchdog parent process
func runWatchdog() {
	exe, err := os.Executable()
	if err != nil {
		runOriginal()
		return
	}

	cmdStdin := os.Stdin
	if !term.IsTerminal(int(cmdStdin.Fd())) {
		if tty, ttyErr := os.OpenFile("/dev/tty", os.O_RDWR, 0); ttyErr == nil {
			cmdStdin = tty
		}
	}

	// save original terminal settings in parent process if Stdin is a terminal
	var watchdogOldState *term.State
	if term.IsTerminal(int(cmdStdin.Fd())) {
		var errState error
		watchdogOldState, errState = term.MakeRaw(int(cmdStdin.Fd()))
		if errState == nil {
			_ = term.Restore(int(cmdStdin.Fd()), watchdogOldState)
		}
	}

	r, w, err := os.Pipe()
	if err != nil {
		runOriginal()
		return
	}

	cmd := exec.CommandContext(context.Background(), exe, os.Args[1:]...)
	cmd.Env = append(os.Environ(), "IRIS_IS_CHILD=true")
	cmd.Stdin = cmdStdin
	cmd.Stdout = os.Stdout
	cmd.Stderr = w

	err = cmd.Start()
	if err != nil {
		runOriginal()
		return
	}

	_ = w.Close()

	// copy child stderr to both our buffer and the real stderr, filtering out panics
	var stderrBuf bytes.Buffer
	origStderr := os.Stderr
	tempBuf := make([]byte, 1024)
	suppress := false
	for {
		n, errRead := r.Read(tempBuf)
		if n > 0 {
			_, _ = stderrBuf.Write(tempBuf[:n])
			if stderrBuf.Len() > 64*1024 {
				// discard oldest bytes to avoid memory leak
				over := stderrBuf.Len() - 64*1024
				_ = stderrBuf.Next(over)
			}
			if !suppress {
				currentContent := stderrBuf.Bytes()
				searchStart := 0
				if len(currentContent) > n+12 {
					searchStart = len(currentContent) - (n + 12)
				}
				searchSlice := currentContent[searchStart:]
				idxPanic := bytes.Index(searchSlice, []byte("panic:"))
				idxFatal := bytes.Index(searchSlice, []byte("fatal error:"))
				triggerIdx := -1
				if idxPanic != -1 {
					triggerIdx = searchStart + idxPanic
				} else if idxFatal != -1 {
					triggerIdx = searchStart + idxFatal
				}

				if triggerIdx != -1 {
					suppress = true
					printedLen := len(currentContent) - n
					if triggerIdx > printedLen {
						_, _ = origStderr.Write(currentContent[printedLen:triggerIdx])
					}
				} else {
					_, _ = origStderr.Write(tempBuf[:n])
				}
			}
		}
		if errRead != nil {
			break
		}
	}

	// check if child exited abnormally or crashed
	errWait := cmd.Wait()
	if errWait != nil {
		content := stderrBuf.Bytes()
		if bytes.Contains(content, []byte("panic:")) || bytes.Contains(content, []byte("fatal error:")) {
			WriteCrashLog(string(content))
			// restore terminal state if watchdog saved it
			if watchdogOldState != nil {
				_ = term.Restore(int(cmdStdin.Fd()), watchdogOldState)
			}
			printCrashNotice()
			startRescueShell()
			os.Exit(2)
		}

		var exitErr *exec.ExitError
		if errors.As(errWait, &exitErr) {
			os.Exit(exitErr.ExitCode())
		}
		os.Exit(1)
	}
}

// runOriginal runs the normal command execution
func runOriginal() {
	if os.Getenv("IRIS_RELOADED") == "true" {
		fmt.Printf("\r\033[K\033[35m[IRIS] reloading...\033[0m\n")
		_ = os.Unsetenv("IRIS_RELOADED")
	}

	if err := rootCmd.Execute(); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}

func Execute() {
	_ = config.MigrateFromLegacyJSON()
	cfg, err := config.Load()
	if err != nil {
		fmt.Fprintf(os.Stderr, "[IRIS] config error: %v\n", err)
	}
	config.Init(cfg)

	if os.Getenv("IRIS_IS_CHILD") != "true" {
		runWatchdog()
		return
	}

	runOriginal()
}