headscale/integration/dockertestutil/execute.go

package dockertestutil

import (
	"bytes"
	"errors"
	"fmt"
	"os"
	"strconv"
	"sync"
	"time"

	"github.com/juanfont/headscale/hscontrol/util"
	"github.com/ory/dockertest/v3"
)

// ScaleTimeout scales an integration timeout for the running environment: the
// nested-docker nix-VM checks are slower than a host docker daemon and set
// HEADSCALE_INTEGRATION_TIMEOUT_SCALE to widen every budget; otherwise CI uses
// 2x and local dev the base. One knob covers exec, convergence, peer sync, and
// pool waits. Lives here (the lowest integration package) so every layer can
// reach it.
func ScaleTimeout(base time.Duration) time.Duration {
	scale := os.Getenv("HEADSCALE_INTEGRATION_TIMEOUT_SCALE")
	if scale != "" {
		f, err := strconv.ParseFloat(scale, 64)
		if err == nil && f > 0 {
			return time.Duration(float64(base) * f)
		}
	}

	if util.IsCI() {
		return base * 2
	}

	return base
}

// defaultExecuteTimeout returns the timeout for docker exec commands.
func defaultExecuteTimeout() time.Duration {
	return ScaleTimeout(10 * time.Second)
}

var (
	ErrDockertestCommandFailed  = errors.New("dockertest command failed")
	ErrDockertestCommandTimeout = errors.New("dockertest command timed out")
)

type ExecuteCommandConfig struct {
	timeout time.Duration
}

type ExecuteCommandOption func(*ExecuteCommandConfig) error

func ExecuteCommandTimeout(timeout time.Duration) ExecuteCommandOption {
	return ExecuteCommandOption(func(conf *ExecuteCommandConfig) error {
		conf.timeout = timeout
		return nil
	})
}

// buffer is a goroutine safe [bytes.Buffer].
type buffer struct {
	store bytes.Buffer
	mutex sync.Mutex
}

// Write appends the contents of p to the buffer, growing the buffer as needed. It returns
// the number of bytes written.
func (b *buffer) Write(p []byte) (int, error) {
	b.mutex.Lock()
	defer b.mutex.Unlock()

	return b.store.Write(p)
}

// String returns the contents of the unread portion of the buffer
// as a string.
func (b *buffer) String() string {
	b.mutex.Lock()
	defer b.mutex.Unlock()

	return b.store.String()
}

func ExecuteCommand(
	resource *dockertest.Resource,
	cmd []string,
	env []string,
	options ...ExecuteCommandOption,
) (string, string, error) {
	stdout := buffer{}
	stderr := buffer{}

	execConfig := ExecuteCommandConfig{
		timeout: defaultExecuteTimeout(),
	}

	for _, opt := range options {
		err := opt(&execConfig)
		if err != nil {
			return "", "", fmt.Errorf("execute-command/options: %w", err)
		}
	}

	type result struct {
		exitCode int
		err      error
	}

	resultChan := make(chan result, 1)

	// Run your long running function in it's own goroutine and pass back it's
	// response into our channel.
	go func() {
		exitCode, err := resource.Exec(
			cmd,
			dockertest.ExecOptions{
				Env:    append(env, "HEADSCALE_LOG_LEVEL=info"),
				StdOut: &stdout,
				StdErr: &stderr,
			},
		)

		resultChan <- result{exitCode, err}
	}()

	// Listen on our channel AND a timeout channel - which ever happens first.
	select {
	case res := <-resultChan:
		if res.err != nil {
			return stdout.String(), stderr.String(), fmt.Errorf("command failed, stderr: %s: %w", stderr.String(), res.err)
		}

		if res.exitCode != 0 {
			// Uncomment for debugging
			// log.Println("Command: ", cmd)
			// log.Println("stdout: ", stdout.String())
			// log.Println("stderr: ", stderr.String())
			return stdout.String(), stderr.String(), fmt.Errorf("command failed, stderr: %s: %w", stderr.String(), ErrDockertestCommandFailed)
		}

		return stdout.String(), stderr.String(), nil
	case <-time.After(execConfig.timeout):
		return stdout.String(), stderr.String(), fmt.Errorf("command failed, stderr: %s: %w", stderr.String(), ErrDockertestCommandTimeout)
	}
}