challenge-tester.go

package main

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"regexp"
	"slices"
	"strings"
	"time"

	"github.com/go-cmd/cmd"
	"github.com/spf13/cobra"
)

var errSomeChallengesFailed = errors.New("some challenges failed")

type TestRunner struct {
	ChallengesDir string
	Challenges    []string
	TokenPremium  string
	TokenFree     string
	Token         string
	Patterns      []string
	DryRun        bool
	Verbose       bool
	TestTimeout   time.Duration
	RecentSince   time.Duration
}

func main() {
	var runner TestRunner

	rootCmd := &cobra.Command{
		Use:   "challenge-tester",
		Short: "...",
		RunE: func(cmd *cobra.Command, args []string) error {
			if runner.TokenPremium == "" && runner.TokenFree == "" {
				return fmt.Errorf("either --as-premium-user or --as-free-user must be provided")
			}
			if runner.TokenPremium != "" {
				runner.Token = runner.TokenPremium
			} else {
				runner.Token = runner.TokenFree
			}

			if err := runner.Run(); err != nil {
				if errors.Is(err, errSomeChallengesFailed) {
					cmd.SilenceUsage = true
				}
				return err
			}

			return nil
		},
	}

	rootCmd.Flags().StringVarP(
		&runner.ChallengesDir,
		"challenges-dir",
		"d",
		"",
		"Directory containing challenge folders",
	)
	rootCmd.MarkFlagRequired("challenges-dir")

	rootCmd.Flags().StringSliceVarP(
		&runner.Challenges,
		"challenge",
		"c",
		[]string{},
		"Run tests for specific challenges (can be specified multiple times)",
	)

	rootCmd.Flags().StringVar(
		&runner.TokenPremium,
		"as-premium-user",
		"",
		"Run tests as a premium user with the given token",
	)

	rootCmd.Flags().StringVar(
		&runner.TokenFree,
		"as-free-user",
		"",
		"Run tests as a free user with the given token",
	)

	rootCmd.Flags().StringSliceVarP(
		&runner.Patterns,
		"pattern",
		"p",
		[]string{},
		"Only run tests for challenges that match the given pattern",
	)

	rootCmd.Flags().BoolVar(
		&runner.DryRun,
		"dry-run",
		false,
		"Only print the test plan without running tests",
	)

	rootCmd.Flags().DurationVarP(
		&runner.RecentSince,
		"recent-since",
		"r",
		time.Duration(60*time.Minute),
		"Only run tests for challenges that haven't been completed recently",
	)

	rootCmd.Flags().DurationVarP(
		&runner.TestTimeout,
		"test-timeout",
		"t",
		time.Duration(5*time.Minute),
		"Timeout for each test case (solution) in a challenge",
	)

	rootCmd.Flags().BoolVarP(
		&runner.Verbose,
		"verbose",
		"v",
		false,
		"Print all output from the challenge",
	)

	if err := rootCmd.Execute(); err != nil {
		log.Fatal(err)
	}
}

type Challenge struct {
	Name      string
	Path      string
	Premium   bool
	Solutions []string

	LastSuccessAt time.Time
}

type TestStats struct {
	TotalChallenges int
	SuccessCount    int
	SkippedCount    int
	FailureCount    int
}

func (tr *TestRunner) Run() error {
	if err := tr.authenticate(); err != nil {
		return err
	}

	var patterns []*regexp.Regexp
	for _, pattern := range tr.Patterns {
		p, err := regexp.Compile(pattern)
		if err != nil {
			return fmt.Errorf("invalid pattern %q: %v", pattern, err)
		}
		patterns = append(patterns, p)
	}

	challenges, err := tr.enumerateChallenges(patterns)
	if err != nil {
		return err
	}

	var stats TestStats
	if tr.DryRun {
		stats, err = tr.printTestPlan(challenges)
	} else {
		stats, err = tr.runTests(challenges)
	}
	if err != nil {
		return err
	}

	tr.printStats(stats)

	if stats.FailureCount > 0 {
		return errSomeChallengesFailed
	}

	return nil
}

func (tr *TestRunner) authenticate() error {
	sess, token, found := strings.Cut(tr.Token, ":")
	if !found {
		return fmt.Errorf("invalid token format")
	}

	authCmd := cmd.NewCmd("labctl", "auth", "login", "-s", sess, "-t", token)
	status := <-authCmd.Start()
	if status.Exit != 0 {
		return fmt.Errorf("labctl auth login failed: %s", status.Stderr)
	}

	return nil
}

func (tr *TestRunner) enumerateChallenges(patterns []*regexp.Regexp) ([]Challenge, error) {
	var challenges []Challenge
	solutionRegex := regexp.MustCompile(`^\.solution(-\d+)?\.sh$`)

	err := filepath.Walk(tr.ChallengesDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		if info.IsDir() {
			return nil
		}

		if solutionRegex.MatchString(filepath.Base(path)) {
			challengeName := filepath.Base(filepath.Dir(path))

			// If we're in multiple challenge mode, only process the specified challenges
			if len(tr.Challenges) > 0 && !slices.Contains(tr.Challenges, challengeName) {
				return nil
			}

			if len(patterns) > 0 && !slices.ContainsFunc(patterns, func(pattern *regexp.Regexp) bool {
				return pattern.MatchString(challengeName)
			}) {
				return nil
			}

			var challenge *Challenge
			for i := range challenges {
				if challenges[i].Name == challengeName {
					challenge = &challenges[i]
					break
				}
			}
			if challenge == nil {
				challenges = append(challenges, Challenge{
					Name: challengeName,
					Path: filepath.Join(tr.ChallengesDir, challengeName),
				})
				challenge = &challenges[len(challenges)-1]

				challenge.Premium, err = isPremiumChallenge(challenge.Path)
				if err != nil {
					return fmt.Errorf("failed to determine if challenge is premium: %v", err)
				}

				if lastSuccessAt, err := readSuccessMarker(challenge.Path); err == nil {
					challenge.LastSuccessAt = lastSuccessAt
				} else {
					return fmt.Errorf("failed to read success marker: %v", err)
				}
			}
			challenge.Solutions = append(challenge.Solutions, path)
		}

		return nil
	})

	return challenges, err
}

func (tr *TestRunner) printTestPlan(challenges []Challenge) (TestStats, error) {
	stats := TestStats{TotalChallenges: len(challenges)}

	fmt.Println("TEST PLAN:")
	for _, challenge := range challenges {
		fmt.Printf("CHALLENGE %s\n", challenge.Name)
		if challenge.Premium && tr.Token == tr.TokenFree {
			fmt.Printf("    --- SKIP (premium challenge)\n")
		} else if !challenge.LastSuccessAt.IsZero() && time.Since(challenge.LastSuccessAt) < tr.RecentSince {
			fmt.Printf("    --- SKIP (recently completed)\n")
		} else {
			for _, solution := range challenge.Solutions {
				testCase := challenge.Name + "/" + filepath.Base(solution)
				fmt.Printf("    --- SOLUTION %s\n", testCase)
			}
		}
	}

	return stats, nil
}

func (tr *TestRunner) runTests(challenges []Challenge) (TestStats, error) {
	stats := TestStats{TotalChallenges: len(challenges)}

	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	for _, challenge := range challenges {
		fmt.Printf("CHALLENGE %s\n", challenge.Name)

		if challenge.Premium && tr.Token == tr.TokenFree {
			status := <-cmd.NewCmd("labctl", "challenge", "start", "--safety-disclaimer-consent", challenge.Name).Start()
			output := strings.Join(status.Stdout, "\n") + strings.Join(status.Stderr, "\n")
			if containsIgnoreCase(output, "unable to start a premium challenge") {
				stats.SkippedCount++
				fmt.Printf("    --- SKIP (premium challenge)\n\n")
			} else {
				stats.FailureCount++
				fmt.Printf("    --- FAIL (unexpected output)\n")
				fmt.Println(output)
				fmt.Println()
			}

			continue
		}

		if !challenge.LastSuccessAt.IsZero() && time.Since(challenge.LastSuccessAt) < tr.RecentSince {
			fmt.Printf("    --- SKIP (recently completed)\n\n")
			stats.SkippedCount++
			continue
		}

		challengeSuccess, err := tr.runChallengeSolutions(challenge, ticker)
		if err != nil {
			return stats, err
		}

		if challengeSuccess {
			stats.SuccessCount++
			if err := writeSuccessMarker(challenge.Path); err != nil {
				return stats, fmt.Errorf("failed to write success marker: %v", err)
			}
		} else {
			stats.FailureCount++
			if err := deleteSuccessMarker(challenge.Path); err != nil {
				return stats, fmt.Errorf("failed to delete success marker: %v", err)
			}
		}

		fmt.Println()
	}

	return stats, nil
}

func (tr *TestRunner) runChallengeSolutions(challenge Challenge, ticker *time.Ticker) (bool, error) {
	challengeSuccess := true

	for _, solution := range challenge.Solutions {
		testCase := challenge.Name + "/" + filepath.Base(solution)
		fmt.Printf("    --- SOLUTION %s\n", testCase)

		solutionSuccess, err := tr.runSingleSolution(challenge, solution, ticker)
		if err != nil {
			return false, err
		}

		if !solutionSuccess {
			challengeSuccess = false
			<-cmd.NewCmd("labctl", "challenge", "stop", challenge.Name).Start()
		}
	}

	return challengeSuccess, nil
}

func (tr *TestRunner) runSingleSolution(challenge Challenge, solutionPath string, ticker *time.Ticker) (bool, error) {
	testCase := challenge.Name + "/" + filepath.Base(solutionPath)

	solution, err := os.ReadFile(solutionPath)
	if err != nil {
		return false, fmt.Errorf("failed to read solution file: %v", err)
	}

	chunks := parseSolutionChunks(solution)
	if len(chunks) == 0 {
		return false, fmt.Errorf("no solution content found")
	}

	// Build command for first chunk
	challengeArgs := []string{"challenge", "start", "--safety-disclaimer-consent"}
	if chunks[0].Machine != "" {
		challengeArgs = append(challengeArgs, "--machine", chunks[0].Machine)
	}
	if chunks[0].User != "" {
		challengeArgs = append(challengeArgs, "--user", chunks[0].User)
	}
	challengeArgs = append(challengeArgs, challenge.Name)

	challengeCmd := cmd.NewCmdOptions(
		cmd.Options{Buffered: false, Streaming: true},
		"labctl", challengeArgs...,
	)

	firstChunkContent := chunks[0].Content
	firstChunkContent = "set -exuo pipefail\n" + firstChunkContent + "\n"

	challengeStdin := newReadCloser(bytes.NewReader([]byte(firstChunkContent)))
	statusCh := challengeCmd.StartWithStdin(challengeStdin)

	// Start goroutine for subsequent chunks if any
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	if len(chunks) > 1 {
		go tr.executeSubsequentChunks(ctx, challenge.Name, chunks[1:])
	}

	var combinedOutput string
	outputCh := make(chan string, 1000)
	timeout := time.After(tr.TestTimeout)
	timedOut := false
	gone := 0

	for {
		select {
		case line := <-challengeCmd.Stdout:
			outputCh <- line

		case line := <-challengeCmd.Stderr:
			outputCh <- line

		case line := <-outputCh:
			if tr.Verbose {
				fmt.Println(line)
			}

			combinedOutput += line + "\n"
			if strings.Contains(combinedOutput, "Playground stopped") {
				challengeStdin.Close()
			}
			if strings.Contains(combinedOutput, "Couldn't start solving") {
				challengeStdin.Close()
			}

		case <-timeout:
			timedOut = true
			challengeStdin.Close()
			challengeCmd.Stop()

		case <-ticker.C:
			status := <-cmd.NewCmd("labctl", "challenge", "list", "-q").Start()
			if !strings.Contains(strings.Join(status.Stdout, "\n")+strings.Join(status.Stderr, "\n"), challenge.Name) {
				gone++ // need to see it's gone twice to reduce the chance of competing with a normal exit
				if gone > 1 {
					challengeStdin.Close()
					challengeCmd.Stop()
				}
			}

		case status := <-statusCh:
			combinedOutput += drainOutputCh(outputCh)
			combinedOutput += drainOutputCh(challengeCmd.Stdout)
			combinedOutput += drainOutputCh(challengeCmd.Stderr)
			combinedOutput = strings.TrimRight(combinedOutput, "\n")

			if status.Exit != 0 {
				extra := ""
				if timedOut {
					extra = " (timeout)"
				}
				if gone > 1 {
					extra = " (gone)"
				}
				fmt.Printf("        FAIL %s: exit code %d%s\n", testCase, status.Exit, extra)
				fmt.Println(combinedOutput)
				return false, nil
			} else if !containsIgnoreCase(combinedOutput, "challenge completed") {
				fmt.Printf("        FAIL %s: 'Challenge completed' not found in output\n", testCase)
				fmt.Println(combinedOutput)
				return false, nil
			} else {
				fmt.Printf("        SUCCESS %s\n", testCase)
				return true, nil
			}
		}
	}
}

func (tr *TestRunner) executeSubsequentChunks(ctx context.Context, challengeName string, chunks []SolutionChunk) {
	// Wait for playground to become available
	var playgroundID string
	var err error

	// Poll for playground with exponential backoff
	for attempt := 0; attempt < 10; attempt++ {
		select {
		case <-ctx.Done():
			return
		default:
		}

		playgroundID, err = getPlaygroundID(challengeName)
		if err == nil {
			break
		}

		// Exponential backoff: 1s, 2s, 4s, 8s, then 10s max
		sleepDuration := max(time.Duration(1<<uint(attempt)), 10) * time.Second

		select {
		case <-ctx.Done():
			return
		case <-time.After(sleepDuration):
		}
	}

	if err != nil {
		// Could not find playground, but don't fail the test - the main chunk might still succeed
		return
	}

	// Execute chunks sequentially
	for _, chunk := range chunks {
		select {
		case <-ctx.Done():
			return
		default:
		}

		sshArgs := []string{"ssh", playgroundID}
		if chunk.Machine != "" {
			sshArgs = append(sshArgs, "--machine", chunk.Machine)
		}
		if chunk.User != "" {
			sshArgs = append(sshArgs, "--user", chunk.User)
		}

		sshCmd := cmd.NewCmdOptions(
			cmd.Options{Buffered: false, Streaming: tr.Verbose},
			"labctl", sshArgs...,
		)
		chunkContent := "set -exuo pipefail\n" + chunk.Content + "\n"
		sshStdin := bytes.NewReader([]byte(chunkContent))
		sshStatusCh := sshCmd.StartWithStdin(sshStdin)

		if tr.Verbose {
			// Stream output in verbose mode
			chunkLabel := chunk.Machine
			if chunk.User != "" {
				chunkLabel = chunk.User + "@" + chunk.Machine
			}

			go func() {
				for line := range sshCmd.Stdout {
					fmt.Printf("[chunk %s] %s\n", chunkLabel, line)
				}
			}()
			go func() {
				for line := range sshCmd.Stderr {
					fmt.Printf("[chunk %s] %s\n", chunkLabel, line)
				}
			}()
		}

		select {
		case <-ctx.Done():
			sshCmd.Stop()
			return
		case sshStatus := <-sshStatusCh:
			if sshStatus.Exit != 0 {
				// Chunk failed, but don't affect the main solution result
				// The failure will be visible in the logs
				return
			}
		}
	}
}

func (tr *TestRunner) printStats(stats TestStats) {
	fmt.Printf("\nTEST RESULTS:\n")
	fmt.Printf("Challenges: %d\n", stats.TotalChallenges)
	fmt.Printf("   Success: %d\n", stats.SuccessCount)
	fmt.Printf("   Skipped: %d\n", stats.SkippedCount)
	fmt.Printf("   Failure: %d\n", stats.FailureCount)
}

func containsIgnoreCase(s, substr string) bool {
	s, substr = strings.ToLower(s), strings.ToLower(substr)
	return strings.Contains(s, substr)
}

type readCloser struct {
	ctx    context.Context
	cancel context.CancelFunc

	inner io.Reader
}

func newReadCloser(inner io.Reader) readCloser {
	ctx, cancel := context.WithCancel(context.Background())
	return readCloser{ctx: ctx, cancel: cancel, inner: inner}
}

func (r readCloser) Read(p []byte) (int, error) {
	n, err := r.inner.Read(p)
	if err != nil && err == io.EOF {
		<-r.ctx.Done()
		return n, err
	}

	return n, err
}

func (r *readCloser) Close() error {
	r.cancel()
	return nil
}

func writeSuccessMarker(challengePath string) error {
	now := time.Now().Format("2006-01-02T15:04:05Z")
	return os.WriteFile(filepath.Join(challengePath, ".success"), []byte(now), 0o644)
}

func readSuccessMarker(challengePath string) (time.Time, error) {
	data, err := os.ReadFile(filepath.Join(challengePath, ".success"))
	if err != nil && os.IsNotExist(err) {
		return time.Time{}, nil
	}
	if err != nil {
		return time.Time{}, err
	}

	return time.Parse("2006-01-02T15:04:05Z", string(data))
}

func deleteSuccessMarker(challengePath string) error {
	if err := os.Remove(filepath.Join(challengePath, ".success")); err != nil && !os.IsNotExist(err) {
		return err
	}
	return nil
}

var premiumChallengePattern = regexp.MustCompile(`(?m)^\s*#\s*premium:\s*true\s*$`)

func isPremiumChallenge(challengePath string) (bool, error) {
	content, err := os.ReadFile(filepath.Join(challengePath, "index.md"))
	if err != nil {
		if os.IsNotExist(err) { // contributed challenges don't have index.md, and are all free (hopefully)
			return false, nil
		}
		return false, err
	}

	return premiumChallengePattern.Match(content), nil
}

type SolutionChunk struct {
	Content string
	Machine string
	User    string
}

func parseSolutionChunks(content []byte) []SolutionChunk {
	lines := strings.Split(string(content), "\n")
	var chunks []SolutionChunk
	var currentChunk SolutionChunk
	var currentContent []string

	sessionRegex := regexp.MustCompile(`^\s*#\s*session:\s*(?:([^@\s]+)@)?([^@\s]+)\s*$`)

	for _, line := range lines {
		if matches := sessionRegex.FindStringSubmatch(line); matches != nil {
			// Save previous chunk if it has content
			if len(currentContent) > 0 {
				currentChunk.Content = strings.Join(currentContent, "\n")
				chunks = append(chunks, currentChunk)
				currentContent = nil
			}

			// Start new chunk
			currentChunk = SolutionChunk{
				User:    matches[1], // may be empty
				Machine: matches[2],
			}
		} else {
			currentContent = append(currentContent, line)
		}
	}

	// Add the last chunk
	if len(currentContent) > 0 {
		currentChunk.Content = strings.Join(currentContent, "\n")
		chunks = append(chunks, currentChunk)
	}

	return chunks
}

func getPlaygroundID(challengeName string) (string, error) {
	status := <-cmd.NewCmd("labctl", "playground", "list").Start()
	if status.Exit != 0 {
		return "", fmt.Errorf("labctl playground list failed: %s", strings.Join(status.Stderr, "\n"))
	}

	output := strings.Join(status.Stdout, "\n")
	lines := strings.Split(output, "\n")

	for _, line := range lines {
		if !strings.Contains(line, " running ") {
			continue
		}

		if !strings.HasSuffix(line, "/"+challengeName) {
			continue
		}

		return strings.Fields(line)[0], nil
	}

	return "", fmt.Errorf("playground for challenge %s not found", challengeName)
}

func drainOutputCh(ch <-chan string) string {
	var output string
	for len(ch) > 0 {
		output += <-ch + "\n"
	}
	return output
}