// Copyright 2019-2024 The Inspektor Gadget authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package command provides a generic way for running testing commands.
package command

import (
	"bytes"
	"errors"
	"io"
	"os/exec"
	"syscall"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
)

type Command struct {
	// Name of the command to be run, used to give information.
	Name string

	// ValidateOutput is a function used to verify the output. It must make the test fail in
	// case of error.
	ValidateOutput func(t *testing.T, output string)

	// ValidateStdErrOutput is a function used to verify the output. It must make the test fail in
	// case of error.
	ValidateStdErrOutput func(t *testing.T, output string)

	// StdOutWriter is an optional writer to which the command's standard output will be written.
	// It's used in situations where we want to process the output on-line without saving it.
	StdOutWriter io.Writer

	// StdErrLineObserver, if set, is invoked with each complete line of the command's
	// standard error as it is produced. It allows callers to react to output while the
	// command is still running (e.g. to detect a readiness marker).
	StdErrLineObserver func(line string)

	// StdErrStoreFilter, if set and returning true for a given standard error line, causes
	// that line to be omitted from the stderr buffer used for validation and logging. It's
	// used to drop verbose/debug noise while still allowing StdErrLineObserver to see it.
	StdErrStoreFilter func(line string) bool

	// StartAndStop indicates this command should first be started then stopped.
	// It corresponds to gadget like execsnoop which wait user to type Ctrl^C.
	StartAndStop bool

	// started indicates this command was started.
	// It is only used by command which have StartAndStop set.
	started bool

	// Cmd object is used when we want to start the command, then do
	// other stuff and wait for its completion or just run the command.
	Cmd *exec.Cmd

	// stdout contains command standard output when started using Startcommand().
	stdout bytes.Buffer

	// stderr contains command standard output when started using Startcommand().
	stderr bytes.Buffer

	// stderrProc is the live stderr line processor, set when StdErrLineObserver or
	// StdErrStoreFilter is used. It must be flushed once the process has exited so a final
	// line without a trailing newline is not lost.
	stderrProc *stderrLineProcessor

	// trackExit, when set before Start(), makes a background goroutine own Cmd.Wait() so the
	// command's exit can be observed (via ExitCh) while it is still "running" from the test's
	// point of view. It is used by the readiness gate to fail fast when a started command exits
	// before becoming ready, instead of waiting for the readiness timeout. Untracked commands
	// keep the original Stop()/kill() behaviour.
	trackExit bool
	exitCh    chan struct{}
	exitErr   error
}

// TrackExit enables background exit tracking (see the trackExit field). It must be called before
// Start().
func (c *Command) TrackExit() {
	c.trackExit = true
}

// ExitCh returns a channel that is closed once a tracked command's process has exited (and been
// reaped), or nil if exit tracking is not enabled. ExitErr may be read once it is closed.
func (c *Command) ExitCh() <-chan struct{} {
	return c.exitCh
}

// ExitErr returns the error from Cmd.Wait() for a tracked command. It must only be read after
// ExitCh() has been closed.
func (c *Command) ExitErr() error {
	return c.exitErr
}

func (c *Command) IsStartAndStop() bool {
	return c.StartAndStop
}

func (c *Command) Running() bool {
	return c.started
}

// initExecCmd configures c.Cmd to store the stdout and stderr in c.stdout and c.stderr so that we
// can use them on c.verifyOutput().
func (c *Command) initExecCmd() {
	c.Cmd.Stdout = &c.stdout
	c.Cmd.Stderr = &c.stderr

	if c.StdOutWriter != nil {
		c.Cmd.Stdout = c.StdOutWriter
	}

	// When a caller wants to observe stderr lines live or filter noise out of the stored
	// stderr, route stderr through a line processor that feeds the observer and only stores
	// the lines that pass the filter.
	if c.StdErrLineObserver != nil || c.StdErrStoreFilter != nil {
		c.stderrProc = &stderrLineProcessor{
			observer: c.StdErrLineObserver,
			filter:   c.StdErrStoreFilter,
			store:    &c.stderr,
		}
		c.Cmd.Stderr = c.stderrProc
	}

	// To be able to kill the process of /bin/sh and its child (the process of
	// c.Cmd), we need to send the termination signal to their process group ID
	// (PGID). However, child processes get the same PGID as their parents by
	// default, so in order to avoid killing also the integration tests process,
	// we set the fields Setpgid and Pgid of syscall.SysProcAttr before
	// executing /bin/sh. Doing so, the PGID of /bin/sh (and its children)
	// will be set to its process ID, see:
	// https://cs.opensource.google/go/go/+/refs/tags/go1.17.8:src/syscall/exec_linux.go;l=32-34.
	c.Cmd.SysProcAttr = &syscall.SysProcAttr{Setpgid: true, Pgid: 0}
}

// verifyOutput verifies the output of the command by using the
// ValidateOutput callback function provided by the user.
func (c *Command) verifyOutput(t *testing.T) {
	if c.ValidateOutput != nil {
		c.ValidateOutput(t, c.stdout.String())
	}
}

// verifyStderrOutput verifies the output of the command by using the
// ValidateStdErrOutput callback function provided by the user.
func (c *Command) verifyStderrOutput(t *testing.T) {
	if c.ValidateStdErrOutput != nil {
		c.ValidateStdErrOutput(t, c.stderr.String())
	}
}

// flushStderr flushes any final stderr line that was not terminated by a newline. It must be
// called only after the process has exited (so all writes to the processor are done), which is
// guaranteed after Cmd.Run()/Cmd.Wait() returns.
func (c *Command) flushStderr() {
	if c.stderrProc != nil {
		c.stderrProc.flush()
	}
}

// Run runs the Command on the given as parameter test.
func (c *Command) Run(t *testing.T) {
	c.initExecCmd()

	t.Logf("[%s] Run command(%s):\n", time.Now().UTC(), c.Name)
	err := c.Cmd.Run()
	c.flushStderr()
	t.Logf("[%s] Command returned(%s):\n%s\n%s\n",
		time.Now().UTC(), c.Name, c.stderr.String(), c.stdout.String())
	require.NoError(t, err, "failed to run command(%s)", c.Name)

	c.verifyStderrOutput(t)
	c.verifyOutput(t)
}

// Start starts the Command on the given as parameter test, you need to
// wait it using Stop().
func (c *Command) Start(t *testing.T) {
	if c.started {
		t.Logf("Warn(%s): trying to start command but it was already started\n", c.Name)
		return
	}

	c.initExecCmd()

	t.Logf("[%s] Start command(%s)", time.Now().UTC(), c.Name)
	err := c.Cmd.Start()
	require.NoError(t, err, "failed to start command(%s)", c.Name)

	c.started = true

	// When exit tracking is enabled, a background goroutine owns Cmd.Wait() so that the
	// command's exit (e.g. a gadget that fails to load) can be observed via ExitCh() while the
	// test is still waiting for it to become ready. kill() then reaps through this goroutine
	// instead of calling Cmd.Wait() a second time.
	if c.trackExit {
		c.exitCh = make(chan struct{})
		go func() {
			c.exitErr = c.Cmd.Wait()
			close(c.exitCh)
		}()
	}
}

// Stop stops a Command previously started with Start().
// To do so, it Kill() the process corresponding to this Cmd and then wait for
// its termination.
// Cmd output is then checked with regard to ExpectedString, ExpectedRegexp or ExpectedEntries.
func (c *Command) Stop(t *testing.T) {
	if !c.started {
		t.Logf("Warn(%s): trying to stop command but it was not started\n", c.Name)
		return
	}

	t.Logf("[%s] Stop command(%s)\n", time.Now().UTC(), c.Name)
	err := c.kill()
	c.flushStderr()
	t.Logf("[%s] Command returned(%s):\n%s\n%s\n",
		time.Now().UTC(), c.Name, c.stderr.String(), c.stdout.String())
	require.NoError(t, err, "failed to kill command(%s)", c.Name)

	c.verifyOutput(t)

	c.started = false
}

// kill kills a command by sending SIGKILL because we want to stop the process
// immediately and avoid that the signal is trapped.
func (c *Command) kill() error {
	const sig syscall.Signal = syscall.SIGKILL

	if c.Cmd == nil {
		return nil
	}

	// Tracked commands have a background goroutine that owns Cmd.Wait(). Reap through it
	// instead of calling Wait() again (which would panic).
	if c.exitCh != nil {
		select {
		case <-c.exitCh:
			// The process already exited on its own. Report its exit status just like the
			// untracked path would: a self-exit before readiness is already surfaced by
			// WaitForReady, but a gadget that becomes ready and then crashes during the
			// workload window must not be silently treated as a clean stop.
			return killErrorAllowingSignal(c.exitErr, sig)
		default:
		}

		// The process was still running when we checked; kill its group. It may still exit
		// and be reaped by the background goroutine between the check and here, in which case
		// the group is already gone (ESRCH); that is not an error for us.
		if err := syscall.Kill(-c.Cmd.Process.Pid, sig); err != nil && !errors.Is(err, syscall.ESRCH) {
			return err
		}
		<-c.exitCh
		return killErrorAllowingSignal(c.exitErr, sig)
	}

	// No need to kill, command has not been executed yet or it already exited
	if c.Cmd.ProcessState != nil && c.Cmd.ProcessState.Exited() {
		return nil
	}

	// Given that we set Setpgid, here we just need to send the PID of c.Cmd
	// (which is the same PGID) as a negative number to syscall.Kill(). As a
	// result, the signal will be received by all the processes with such PGID,
	// in our case, the process of /bin/sh and c.Cmd.
	err := syscall.Kill(-c.Cmd.Process.Pid, sig)
	if err != nil {
		return err
	}

	// In some cases, we do not have to wait here because the cmd was executed
	// with run(), which already waits. On the contrary, in the case it was
	// executed with start() thus ig.started is true, we need to wait indeed.
	if c.started {
		return killErrorAllowingSignal(c.Cmd.Wait(), sig)
	}

	return err
}

// killErrorAllowingSignal returns nil if err is nil or is the process being terminated by sig
// (which is what we expect when we SIGKILL a still-running command); otherwise it returns err.
func killErrorAllowingSignal(err error, sig syscall.Signal) error {
	if err == nil {
		return nil
	}

	// Verify if the error is about the signal we just sent. In that case,
	// do not return error, it is what we were expecting.
	var exiterr *exec.ExitError
	if ok := errors.As(err, &exiterr); !ok {
		return err
	}

	waitStatus, ok := exiterr.Sys().(syscall.WaitStatus)
	if !ok {
		return err
	}

	if waitStatus.Signal() != sig {
		return err
	}

	return nil
}

// stderrLineProcessor is an io.Writer that splits the incoming bytes into lines and, for each
// complete line, invokes an optional observer and optionally stores the line in a backing buffer.
// It's used to observe stderr live (e.g. to detect a readiness marker) while keeping the stored
// stderr free of verbose/debug noise.
type stderrLineProcessor struct {
	buf      []byte
	observer func(line string)
	filter   func(line string) bool
	store    *bytes.Buffer
}

func (p *stderrLineProcessor) Write(b []byte) (int, error) {
	p.buf = append(p.buf, b...)
	for {
		i := bytes.IndexByte(p.buf, '\n')
		if i < 0 {
			break
		}
		line := string(p.buf[:i])
		p.buf = append(p.buf[:0], p.buf[i+1:]...)
		p.handle(line, true)
	}
	return len(b), nil
}

func (p *stderrLineProcessor) handle(line string, hasNewline bool) {
	if p.observer != nil {
		p.observer(line)
	}
	if p.store != nil && (p.filter == nil || !p.filter(line)) {
		p.store.WriteString(line)
		if hasNewline {
			p.store.WriteByte('\n')
		}
	}
}

// flush emits any buffered bytes that were not terminated by a newline as a final line. It must
// be called only after the last Write (i.e. after the process has exited) to avoid data races.
func (p *stderrLineProcessor) flush() {
	if len(p.buf) == 0 {
		return
	}
	line := string(p.buf)
	p.buf = p.buf[:0]
	p.handle(line, false)
}
