// Unless explicitly stated otherwise all files in this repository are licensed
// under the Apache License Version 2.0.
// This product includes software developed at Datadog (https://www.datadoghq.com/).
// Copyright 2026-present Datadog, Inc.

package tests_test

import (
	"context"
	"strings"
	"testing"
	"time"

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

// --- Infinite-loop / runaway loop containment ---

// While-true with no body output terminates within the ctx deadline.
func TestWhilePentestInfiniteEmptyBody(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
	defer cancel()
	start := time.Now()
	_, _, _ = whileRunCtx(ctx, t, `while true; do :; done`)
	assert.Less(t, time.Since(start), 5*time.Second, "loop did not terminate after deadline")
}

// While-true producing output respects the runner's stdout cap and the ctx
// deadline. Stdout must not grow past a generous upper bound.
func TestWhilePentestRunawayOutput(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	stdout, _, _ := whileRunCtx(ctx, t, `while true; do echo x; done`)
	const generousUpperBound = 1 << 25 // 32 MiB
	assert.Less(t, len(stdout), generousUpperBound, "stdout grew past cap; possible unbounded buffer")
}

// Nested-while-true with the inner producing output: same containment as the
// flat case.
func TestWhilePentestNestedRunaway(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	stdout, _, _ := whileRunCtx(ctx, t, `while true; do while true; do echo y; done; done`)
	const generousUpperBound = 1 << 25
	assert.Less(t, len(stdout), generousUpperBound)
}

// Until-false with no body output: same containment as while-true.
func TestUntilPentestInfiniteEmptyBody(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
	defer cancel()
	start := time.Now()
	_, _, _ = whileRunCtx(ctx, t, `until false; do :; done`)
	assert.Less(t, time.Since(start), 5*time.Second)
}

// --- Stack / recursion safety ---

// 200 levels of nested while loops with an innermost break should unwind
// without stack overflow.
func TestWhilePentestDeeplyNested(t *testing.T) {
	const depth = 200
	var b strings.Builder
	for i := 0; i < depth; i++ {
		b.WriteString("while true; do ")
	}
	b.WriteString("echo ok; break ")
	b.WriteString(itoa(depth))
	b.WriteString("; ")
	for i := 0; i < depth; i++ {
		b.WriteString("done; ")
	}
	b.WriteString("echo done")
	stdout, _, code := whileRun(t, b.String())
	assert.Equal(t, 0, code)
	assert.Equal(t, "ok\ndone\n", stdout)
}

// 200 alternating while/until levels — stress the in-loop save/restore path
// across a heterogeneous nesting chain.
func TestWhilePentestAlternatingNested(t *testing.T) {
	const depth = 200
	var b strings.Builder
	for i := 0; i < depth; i++ {
		if i%2 == 0 {
			b.WriteString("while true; do ")
		} else {
			b.WriteString("until false; do ")
		}
	}
	b.WriteString("echo ok; break ")
	b.WriteString(itoa(depth))
	b.WriteString("; ")
	for i := 0; i < depth; i++ {
		b.WriteString("done; ")
	}
	b.WriteString("echo done")
	stdout, _, code := whileRun(t, b.String())
	assert.Equal(t, 0, code)
	assert.Equal(t, "ok\ndone\n", stdout)
}

// --- Excess break/continue arguments are clamped to the actual nesting depth ---

// "break 1000" at innermost loop must NOT crash, panic, or hang. Bash clamps
// to "break the innermost loop".
func TestWhilePentestBreakHugeNumber(t *testing.T) {
	stdout, _, code := whileRun(t, `while true; do echo x; break 1000; done; echo done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "x\ndone\n", stdout)
}

// "continue 1000" — same deal but with continue.
func TestWhilePentestContinueHugeNumber(t *testing.T) {
	stdout, _, code := whileRun(t, `i=; while [ "$i" != aa ]; do i="${i}a"; echo "$i"; continue 1000; done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "a\naa\n", stdout)
}

// --- High iteration count (50k iterations of a counter loop) ---

// Long counter-driven while loop completes without leaking memory or
// goroutines. We use string-concat as the counter so this exercises the body
// dispatch path (assignments + tests) at scale.
func TestWhilePentestManyIterations(t *testing.T) {
	// 5000 'a's — concatenation ops are O(n^2) on string-concat, so going
	// much higher pushes the test runtime past CI budgets.
	const target = 5000
	var bound strings.Builder
	for i := 0; i < target; i++ {
		bound.WriteString("a")
	}
	script := `i=; while [ "$i" != ` + bound.String() + ` ]; do i="${i}a"; done; echo done`
	stdout, _, code := whileRun(t, script)
	assert.Equal(t, 0, code)
	assert.Equal(t, "done\n", stdout)
}

// --- Cond with side-effecting pipeline ---

// Pipeline cond runs the producer and consumer per iteration; after a break
// the loop exits cleanly with no goroutine leak.
func TestWhilePentestPipelineCondLeakCheck(t *testing.T) {
	stdout, _, code := whileRun(t, `i=; while echo cond | grep -q cond && [ "$i" != aa ]; do i="${i}a"; echo "$i"; done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "a\naa\n", stdout)
}

// --- Body-side break/continue in pathological positions ---

// break inside an && short-circuit inside the body still exits the loop.
func TestWhilePentestBreakAfterAnd(t *testing.T) {
	stdout, _, code := whileRun(t, `while true; do true && break; done; echo done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "done\n", stdout)
}

// break inside an || short-circuit inside the body still exits the loop.
func TestWhilePentestBreakAfterOr(t *testing.T) {
	stdout, _, code := whileRun(t, `while true; do false || break; done; echo done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "done\n", stdout)
}

// continue followed by more body stmts — the trailing stmts must be skipped
// (matches loopStmtsBroken's early return on contnEnclosing).
func TestWhilePentestContinueSkipsTrailingBody(t *testing.T) {
	stdout, _, code := whileRun(t, `i=; while [ "$i" != aa ]; do i="${i}a"; continue; echo unreachable-1; echo unreachable-2; done`)
	assert.Equal(t, 0, code)
	assert.Equal(t, "", stdout)
}

// --- ctx-cancellation mid-cond ---

// A ctx cancelled while the cond is being evaluated should still produce a
// finite-time return.
func TestWhilePentestCancelDuringCond(t *testing.T) {
	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
	defer cancel()
	start := time.Now()
	_, _, _ = whileRunCtx(ctx, t, `while echo x; do :; done`)
	assert.Less(t, time.Since(start), 5*time.Second)
}

// --- Helpers ---

func itoa(n int) string {
	if n == 0 {
		return "0"
	}
	neg := n < 0
	if neg {
		n = -n
	}
	var buf [20]byte
	i := len(buf)
	for n > 0 {
		i--
		buf[i] = byte('0' + n%10)
		n /= 10
	}
	if neg {
		i--
		buf[i] = '-'
	}
	return string(buf[i:])
}
