ollie-core/pkg/elevate/broker_test.go

package elevate

import (
	"encoding/binary"
	"encoding/json"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"
)

func startTestBroker(t *testing.T) (*Broker, string) {
	t.Helper()
	dir := t.TempDir()
	sockPath := filepath.Join(dir, "elevate.sock")
	policyPath := filepath.Join(dir, "policy.yaml")

	b, err := NewBroker(BrokerConfig{
		SocketPath:   sockPath,
		PolicyPath:   policyPath,
		Logf:         t.Logf,
		SessionValid: func(id string) bool { return id == "test-session" },
	})
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(b.Close)
	return b, sockPath
}

func sendRequest(t *testing.T, sockPath, cmd, cwd string) net.Conn {
	return sendRequestWithSession(t, sockPath, cmd, cwd, "test-session")
}

func sendRequestWithSession(t *testing.T, sockPath, cmd, cwd, session string) net.Conn {
	t.Helper()
	conn, err := net.DialTimeout("unix", sockPath, 2*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	req, _ := json.Marshal(struct {
		Cmd     string            `json:"cmd"`
		Cwd     string            `json:"cwd"`
		Env     map[string]string `json:"env"`
		Session string            `json:"session,omitempty"`
	}{Cmd: cmd, Cwd: cwd, Env: map[string]string{"PATH": "/usr/bin"}, Session: session})
	req = append(req, '\n')
	conn.Write(req) //nolint:errcheck
	return conn
}

func readFrames(t *testing.T, conn net.Conn) (string, int) {
	t.Helper()
	var output string
	exitCode := -1
	header := make([]byte, 5)
	for {
		conn.SetReadDeadline(time.Now().Add(5 * time.Second)) //nolint:errcheck
		_, err := conn.Read(header)
		if err != nil {
			break
		}
		length := binary.BigEndian.Uint32(header[1:5])
		payload := make([]byte, length)
		if length > 0 {
			conn.Read(payload) //nolint:errcheck
		}
		switch header[0] {
		case 'd':
			output += string(payload)
		case 'x':
			fmt.Sscanf(string(payload), "%d", &exitCode)
			return output, exitCode
		}
	}
	return output, exitCode
}

func TestBrokerAutoApprove(t *testing.T) {
	b, sockPath := startTestBroker(t)

	// Add a policy rule
	b.GlobalPolicy().AddGlobal(Rule{Cmd: "echo *"})

	// Send a matching command
	conn := sendRequest(t, sockPath, "echo hello", os.TempDir())
	defer conn.Close()

	output, exitCode := readFrames(t, conn)
	if exitCode != 0 {
		t.Errorf("expected exit 0, got %d; output: %s", exitCode, output)
	}
	if output != "hello\n" {
		t.Errorf("expected 'hello\n', got %q", output)
	}
}

func TestBrokerDenyNoSession(t *testing.T) {
	_, sockPath := startTestBroker(t)

	// Send a request with no session ID
	conn := sendRequestWithSession(t, sockPath, "echo hello", os.TempDir(), "")
	defer conn.Close()

	output, exitCode := readFrames(t, conn)
	if exitCode != 1 {
		t.Errorf("expected exit 1, got %d", exitCode)
	}
	if !strings.Contains(output, "no session identity") {
		t.Errorf("expected 'no session identity' in output, got %q", output)
	}
}

func TestBrokerDenyInvalidSession(t *testing.T) {
	_, sockPath := startTestBroker(t)

	// Send a request with a session ID that doesn't exist
	conn := sendRequestWithSession(t, sockPath, "echo hello", os.TempDir(), "nonexistent-session")
	defer conn.Close()

	output, exitCode := readFrames(t, conn)
	if exitCode != 1 {
		t.Errorf("expected exit 1, got %d", exitCode)
	}
	if !strings.Contains(output, "session not found") {
		t.Errorf("expected 'session not found' in output, got %q", output)
	}
}

func TestBrokerDenyOnTimeout(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping timeout test in short mode")
	}

	// Use a very short TTL for testing
	dir := t.TempDir()
	sockPath := filepath.Join(dir, "elevate.sock")
	policyPath := filepath.Join(dir, "policy.yaml")

	b, err := NewBroker(BrokerConfig{
		SocketPath: sockPath,
		PolicyPath: policyPath,
		Logf:       t.Logf,
	})
	if err != nil {
		t.Fatal(err)
	}
	defer b.Close()

	// No policy match → goes to pending
	conn := sendRequest(t, sockPath, "rm -rf /", os.TempDir())
	defer conn.Close()

	// Manually deny after a short wait
	time.Sleep(100 * time.Millisecond)
	pending := b.Pending()
	if len(pending) != 1 {
		t.Fatalf("expected 1 pending, got %d", len(pending))
	}
	b.Resolve(pending[0].ID, ResolveDeny)

	output, exitCode := readFrames(t, conn)
	if exitCode != 1 {
		t.Errorf("expected exit 1 on deny, got %d", exitCode)
	}
	_ = output
}

func TestBrokerApproveAndExecute(t *testing.T) {
	b, sockPath := startTestBroker(t)

	conn := sendRequest(t, sockPath, "echo approved", os.TempDir())
	defer conn.Close()

	// Wait for it to appear in pending
	time.Sleep(100 * time.Millisecond)
	pending := b.Pending()
	if len(pending) != 1 {
		t.Fatalf("expected 1 pending, got %d", len(pending))
	}

	// Approve it
	b.Resolve(pending[0].ID, ResolveApprove)

	output, exitCode := readFrames(t, conn)
	if exitCode != 0 {
		t.Errorf("expected exit 0, got %d; output: %s", exitCode, output)
	}
	if output != "approved\n" {
		t.Errorf("expected 'approved\n', got %q", output)
	}
}

func TestBrokerPersist(t *testing.T) {
	b, sockPath := startTestBroker(t)

	// First request: no policy, goes to pending
	conn1 := sendRequest(t, sockPath, "echo persist-me", os.TempDir())
	defer conn1.Close()

	time.Sleep(100 * time.Millisecond)
	pending := b.Pending()
	if len(pending) != 1 {
		t.Fatalf("expected 1 pending, got %d", len(pending))
	}

	// Persist it
	b.Resolve(pending[0].ID, ResolvePersist)

	output, exitCode := readFrames(t, conn1)
	if exitCode != 0 {
		t.Errorf("first request: expected exit 0, got %d", exitCode)
	}
	if output != "persist-me\n" {
		t.Errorf("first request: expected 'persist-me\n', got %q", output)
	}

	// Second request with same command: should auto-approve (no pending)
	conn2 := sendRequest(t, sockPath, "echo persist-me", os.TempDir())
	defer conn2.Close()

	output2, exitCode2 := readFrames(t, conn2)
	if exitCode2 != 0 {
		t.Errorf("second request: expected exit 0, got %d", exitCode2)
	}
	if output2 != "persist-me\n" {
		t.Errorf("second request: expected 'persist-me\n', got %q", output2)
	}

	// Verify no pending requests
	if len(b.Pending()) != 0 {
		t.Error("expected no pending after persist")
	}
}

func TestBrokerRateLimitPerTurn(t *testing.T) {
	b, _ := startTestBroker(t)

	sessionID := "test-session"

	// Simulate 3 requests (all should be allowed)
	for i := 0; i < MaxPerTurn; i++ {
		if !b.checkTurnLimit(sessionID) {
			t.Fatalf("request %d should not be rate-limited", i+1)
		}
	}

	// 4th request should be rate-limited
	if b.checkTurnLimit(sessionID) {
		t.Error("4th request should be rate-limited")
	}

	// Reset turn counter
	b.ResetTurn(sessionID)

	// After reset, requests should be allowed again
	if !b.checkTurnLimit(sessionID) {
		t.Error("request after ResetTurn should be allowed")
	}
}

func TestBrokerRateLimitDenialMessage(t *testing.T) {
	b, _ := startTestBroker(t)

	// Test the unit logic directly:
	// identifySession returns "" for test connections (no PEERCRED mapping),
	// and checkTurnLimit allows unknown sessions. Rate limiting only applies
	// when session identity is established.
	sid := "rate-test"
	for i := 0; i < MaxPerTurn; i++ {
		if !b.checkTurnLimit(sid) {
			t.Fatalf("request %d should be allowed", i+1)
		}
	}
	if b.checkTurnLimit(sid) {
		t.Fatal("should be rate-limited after MaxPerTurn")
	}

	// Verify the counter survives and then resets
	b.ResetTurn(sid)
	if !b.checkTurnLimit(sid) {
		t.Fatal("should be allowed after reset")
	}
}

func TestBrokerRemoveSessionCleansUpTurnCount(t *testing.T) {
	b, _ := startTestBroker(t)

	sid := "cleanup-test"

	// Use up some turn budget
	b.checkTurnLimit(sid)
	b.checkTurnLimit(sid)

	// Remove session
	b.RemoveSession(sid)

	// Counter should be gone — new requests allowed
	for i := 0; i < MaxPerTurn; i++ {
		if !b.checkTurnLimit(sid) {
			t.Fatalf("request %d after RemoveSession should be allowed", i+1)
		}
	}
}