Skip to main content

cli/data/
path.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Shared path helpers for export/import data files.
16
17use crate::data::export_v2::schema::{DDL_DIR, SCHEMA_DIR};
18
19pub(crate) fn ddl_path_for_schema(schema: &str) -> String {
20    format!(
21        "{}/{}/{}.sql",
22        SCHEMA_DIR,
23        DDL_DIR,
24        encode_path_segment(schema)
25    )
26}
27
28pub(crate) fn data_dir_for_schema_chunk(schema: &str, chunk_id: u32) -> String {
29    format!("data/{}/{}/", encode_path_segment(schema), chunk_id)
30}
31
32pub(crate) fn encode_path_segment(value: &str) -> String {
33    let mut encoded = String::with_capacity(value.len());
34    for byte in value.bytes() {
35        match byte {
36            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' => {
37                encoded.push(byte as char);
38            }
39            _ => {
40                encoded.push('%');
41                encoded.push(hex_char(byte >> 4));
42                encoded.push(hex_char(byte & 0x0F));
43            }
44        }
45    }
46    encoded
47}
48
49fn hex_char(nibble: u8) -> char {
50    match nibble {
51        0..=9 => (b'0' + nibble) as char,
52        10..=15 => (b'A' + (nibble - 10)) as char,
53        _ => unreachable!("nibble must be in 0..=15"),
54    }
55}
56
57#[cfg(test)]
58mod tests {
59    use super::*;
60
61    #[test]
62    fn test_encode_path_segment_preserves_safe_ascii() {
63        assert_eq!(encode_path_segment("test_db"), "test_db");
64    }
65
66    #[test]
67    fn test_encode_path_segment_escapes_path_traversal_chars() {
68        assert_eq!(encode_path_segment("../evil"), "%2E%2E%2Fevil");
69        assert_eq!(encode_path_segment(r"..\\evil"), "%2E%2E%5C%5Cevil");
70    }
71
72    #[test]
73    fn test_ddl_path_for_schema_encodes_schema_segment() {
74        assert_eq!(ddl_path_for_schema("public"), "schema/ddl/public.sql");
75        assert_eq!(
76            ddl_path_for_schema("../evil"),
77            "schema/ddl/%2E%2E%2Fevil.sql"
78        );
79    }
80
81    #[test]
82    fn test_data_dir_for_schema_chunk_encodes_schema_segment() {
83        assert_eq!(data_dir_for_schema_chunk("public", 1), "data/public/1/");
84        assert_eq!(
85            data_dir_for_schema_chunk("../evil", 7),
86            "data/%2E%2E%2Fevil/7/"
87        );
88    }
89}