refactor: agent rag use `documents` field other than `embeddings` dir (#658)

1 files changed, 52 insertions, 1 deletions

diff --git a/src/utils/mod.rs b/src/utils/mod.rs
index f8bfcc6..7d1d07a 100644
--- a/src/utils/mod.rs
+++ b/src/utils/mod.rs
@@ -17,7 +17,10 @@ pub use self::spinner::{create_spinner, Spinner};
 use fancy_regex::Regex;
 use is_terminal::IsTerminal;
 use lazy_static::lazy_static;
-use std::env;
+use std::{
+    env,
+    path::{self, Path, PathBuf},
+};
 
 lazy_static! {
     pub static ref CODE_BLOCK_RE: Regex = Regex::new(r"(?ms)```\w*(.*)```").unwrap();
@@ -151,6 +154,32 @@ pub fn dimmed_text(input: &str) -> String {
     nu_ansi_term::Style::new().dimmed().paint(input).to_string()
 }
 
+pub fn safe_join_path<T1: AsRef<Path>, T2: AsRef<Path>>(
+    base_path: T1,
+    sub_path: T2,
+) -> Option<PathBuf> {
+    let base_path = base_path.as_ref();
+    let sub_path = sub_path.as_ref();
+    if sub_path.is_absolute() {
+        return None;
+    }
+
+    let mut joined_path = PathBuf::from(base_path);
+
+    for component in sub_path.components() {
+        if path::Component::ParentDir == component {
+            return None;
+        }
+        joined_path.push(component);
+    }
+
+    if joined_path.starts_with(base_path) {
+        Some(joined_path)
+    } else {
+        None
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -161,4 +190,26 @@ mod tests {
         assert!(fuzzy_match("openai:gpt-4-turbo", "oai4"));
         assert!(!fuzzy_match("openai:gpt-4-turbo", "4gpt"));
     }
+
+    #[test]
+    #[cfg(not(target_os = "windows"))]
+    fn test_safe_join_path() {
+        assert_eq!(
+            safe_join_path("/home/user/dir1", "files/file1"),
+            Some(PathBuf::from("/home/user/dir1/files/file1"))
+        );
+        assert!(safe_join_path("/home/user/dir1", "/files/file1").is_none());
+        assert!(safe_join_path("/home/user/dir1", "../file1").is_none());
+    }
+
+    #[test]
+    #[cfg(target_os = "windows")]
+    fn test_safe_join_path() {
+        assert_eq!(
+            safe_join_path("C:\\Users\\user\\dir1", "files/file1"),
+            Some(PathBuf::from("C:\\Users\\user\\dir1\\files\\file1"))
+        );
+        assert!(safe_join_path("C:\\Users\\user\\dir1", "/files/file1").is_none());
+        assert!(safe_join_path("C:\\Users\\user\\dir1", "../file1").is_none());
+    }
 }