149 lines
3.4 KiB
Go
149 lines
3.4 KiB
Go
package sandbox_test
|
|
|
|
import (
|
|
"encoding/json"
|
|
"os"
|
|
"path/filepath"
|
|
"testing"
|
|
|
|
"github.com/VictorVargas/rony-llm-agent/pkg/llm"
|
|
"github.com/VictorVargas/rony-llm-agent/pkg/tools"
|
|
"github.com/VictorVargas/rony-llm-agent/pkg/tools/sandbox"
|
|
)
|
|
|
|
func TestNewSandbox(t *testing.T) {
|
|
dir := t.TempDir()
|
|
|
|
sb, err := sandbox.NewSandbox(dir)
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
if sb == nil {
|
|
t.Fatal("expected non-nil sandbox")
|
|
}
|
|
}
|
|
|
|
func TestValidatePath_Allowed(t *testing.T) {
|
|
dir := t.TempDir()
|
|
sb, err := sandbox.NewSandbox(dir)
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
|
|
// Create a file inside the sandbox
|
|
testFile := filepath.Join(dir, "test.txt")
|
|
if err := os.WriteFile(testFile, []byte("hello"), 0644); err != nil {
|
|
t.Fatalf("creating test file: %v", err)
|
|
}
|
|
|
|
call := llm.ToolCall{
|
|
Name: "read_file",
|
|
Arguments: json.RawMessage(`{"path": "test.txt"}`),
|
|
}
|
|
|
|
err = sb.ValidateToolCall(tools.Tool{}, call)
|
|
if err != nil {
|
|
t.Fatalf("expected no error, got: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestValidatePath_EscapesSandbox(t *testing.T) {
|
|
dir := t.TempDir()
|
|
sb, err := sandbox.NewSandbox(dir)
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
|
|
call := llm.ToolCall{
|
|
Name: "read_file",
|
|
Arguments: json.RawMessage(`{"path": "/etc/passwd"}`),
|
|
}
|
|
|
|
err = sb.ValidateToolCall(tools.Tool{}, call)
|
|
if err == nil {
|
|
t.Fatal("expected sandbox violation error")
|
|
}
|
|
}
|
|
|
|
func TestValidatePath_SymlinkEscape(t *testing.T) {
|
|
dir := t.TempDir()
|
|
|
|
// Create a symlink that escapes the sandbox
|
|
escapeDir := t.TempDir()
|
|
symlinkPath := filepath.Join(dir, "link")
|
|
if err := os.Symlink(escapeDir, symlinkPath); err != nil {
|
|
t.Fatalf("creating symlink: %v", err)
|
|
}
|
|
|
|
sb, err := sandbox.NewSandbox(dir)
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
|
|
// The sandbox should allow symlinks inside the root
|
|
call := llm.ToolCall{
|
|
Name: "read_file",
|
|
Arguments: json.RawMessage(`{"path": "link"}`),
|
|
}
|
|
|
|
err = sb.ValidateToolCall(tools.Tool{}, call)
|
|
if err != nil {
|
|
t.Fatalf("expected symlink to be allowed (path is inside sandbox), got: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestValidatePath_NonExistentFile(t *testing.T) {
|
|
dir := t.TempDir()
|
|
sb, err := sandbox.NewSandbox(dir)
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
|
|
// Non-existent file inside sandbox should be allowed
|
|
call := llm.ToolCall{
|
|
Name: "read_file",
|
|
Arguments: json.RawMessage(`{"path": "nonexistent.txt"}`),
|
|
}
|
|
|
|
err = sb.ValidateToolCall(tools.Tool{}, call)
|
|
if err != nil {
|
|
t.Fatalf("expected no error for non-existent file, got: %v", err)
|
|
}
|
|
}
|
|
|
|
func TestValidatePath_InvalidJSON(t *testing.T) {
|
|
dir := t.TempDir()
|
|
sb, err := sandbox.NewSandbox(dir)
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
|
|
call := llm.ToolCall{
|
|
Name: "read_file",
|
|
Arguments: json.RawMessage(`not json`),
|
|
}
|
|
|
|
err = sb.ValidateToolCall(tools.Tool{}, call)
|
|
if err == nil {
|
|
t.Fatal("expected error for invalid JSON")
|
|
}
|
|
}
|
|
|
|
func TestValidatePath_NoPathsInArgs(t *testing.T) {
|
|
dir := t.TempDir()
|
|
sb, err := sandbox.NewSandbox(dir)
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
|
|
// Arguments with no paths should not cause errors
|
|
call := llm.ToolCall{
|
|
Name: "math_add",
|
|
Arguments: json.RawMessage(`{"a": 5, "b": 10}`),
|
|
}
|
|
|
|
err = sb.ValidateToolCall(tools.Tool{}, call)
|
|
if err != nil {
|
|
t.Fatalf("expected no error, got: %v", err)
|
|
}
|
|
}
|