From 5ff763ecdac1b570a05539b9901e81cb54952cfa Mon Sep 17 00:00:00 2001
From: Matt Low <matt@mlow.ca>
Date: Sun, 26 Nov 2023 06:35:22 +0000
Subject: [PATCH] Only allow read_dir (and other file access) within current
 working dir

Hopefully, anyway :)
---
 pkg/cli/functions.go | 65 +++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 61 insertions(+), 4 deletions(-)

diff --git a/pkg/cli/functions.go b/pkg/cli/functions.go
index b49afb1..1f46093 100644
--- a/pkg/cli/functions.go
+++ b/pkg/cli/functions.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"strings"
 
 	openai "github.com/sashabaranov/go-openai"
 )
@@ -133,11 +134,67 @@ func ExecuteToolCalls(toolCalls []openai.ToolCall) ([]Message, error) {
 	return toolResults, nil
 }
 
+// isPathContained attempts to verify whether `path` is the same as or
+// contained within `directory`. It is overly cautious, returning false even if
+// `path` IS contained within `directory`, but the two paths use different
+// casing, and we happen to be on a case-insensitive filesystem.
+// This is ultimately to attempt to stop an LLM from going outside of where I
+// tell it to. Additional layers of security should be considered.. run in a
+// VM/container.
+func isPathContained(directory string, path string) (bool, error) {
+	// Clean and resolve symlinks for both paths
+	absPath, err := filepath.Abs(path)
+	if err != nil {
+		return false, err
+	}
+	realPath, err := filepath.EvalSymlinks(absPath)
+	if err != nil {
+		return false, err
+	}
+
+	absDirectory, err := filepath.Abs(directory)
+	if err != nil {
+		return false, err
+	}
+	realDirectory, err := filepath.EvalSymlinks(absDirectory)
+	if err != nil {
+		return false, err
+	}
+
+	// Case insensitive checks
+	if !strings.EqualFold(realPath, realDirectory) &&
+		!strings.HasPrefix(strings.ToLower(realPath), strings.ToLower(realDirectory)+string(os.PathSeparator)) {
+		return false, nil
+	}
+
+	return true, nil
+}
+
+func isPathWithinCWD(path string) (bool, *FunctionResult) {
+	cwd, err := os.Getwd()
+	if err != nil {
+		return false, &FunctionResult{Message: "Failed to determine current working directory"}
+	}
+	if ok, err := isPathContained(cwd, path); !ok {
+		if err != nil {
+			return false, &FunctionResult{Message: fmt.Sprintf("Could not determine whether path '%s' is within the current working directory: %s", path, err.Error())}
+		}
+		return false, &FunctionResult{Message: fmt.Sprintf("Path '%s' is not within the current working directory", path)}
+	}
+	return true, nil
+}
+
 func ReadDir(path string) string {
-	// TODO: ensure it is not possible to escape to directories above CWD
-	// TODO: implement whitelist - list of directories which model is allowed to work in
-	targetPath := filepath.Join(".", path)
-	files, err := os.ReadDir(targetPath)
+	// TODO(?): implement whitelist - list of directories which model is allowed to work in
+	if path == "" {
+		path = "."
+	}
+	ok, res := isPathWithinCWD(path)
+	if !ok {
+		return resultToJson(*res)
+	}
+
+	files, err := os.ReadDir(path)
 	if err != nil {
 		return resultToJson(FunctionResult{
 			Message: err.Error(),