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(),