common_meta/kv_backend/
util.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/// Removes sensitive information like passwords from connection strings.
16///
17/// This function sanitizes connection strings by removing credentials:
18/// - For URL format (mysql://user:password@host:port/db): Removes everything before '@'
19/// - For parameter format (host=localhost password=secret): Removes the password parameter
20/// - For URL format without credentials (mysql://host:port/db): Removes the protocol prefix
21///
22/// # Arguments
23///
24/// * `conn_str` - The connection string to sanitize
25///
26/// # Returns
27///
28/// A sanitized version of the connection string with sensitive information removed
29pub fn sanitize_connection_string(conn_str: &str) -> String {
30    // Case 1: URL format with credentials (mysql://user:password@host:port/db)
31    // Extract everything after the '@' symbol
32    if let Some(at_pos) = conn_str.find('@') {
33        return conn_str[at_pos + 1..].to_string();
34    }
35
36    // Case 2: Parameter format with password (host=localhost password=secret dbname=mydb)
37    // Filter out any parameter that starts with "password="
38    if conn_str.contains("password=") {
39        return conn_str
40            .split_whitespace()
41            .filter(|param| !param.starts_with("password="))
42            .collect::<Vec<_>>()
43            .join(" ");
44    }
45
46    // Case 3: URL format without credentials (mysql://host:port/db)
47    // Extract everything after the protocol prefix
48    if let Some(host_part) = conn_str.split("://").nth(1) {
49        return host_part.to_string();
50    }
51
52    // Case 4: Already sanitized or unknown format
53    // Return as is
54    conn_str.to_string()
55}
56
57#[cfg(test)]
58mod tests {
59    use super::*;
60
61    #[test]
62    fn test_sanitize_connection_string() {
63        // Test URL format with username/password
64        let conn_str = "mysql://user:password123@localhost:3306/db";
65        assert_eq!(sanitize_connection_string(conn_str), "localhost:3306/db");
66
67        // Test URL format without credentials
68        let conn_str = "mysql://localhost:3306/db";
69        assert_eq!(sanitize_connection_string(conn_str), "localhost:3306/db");
70
71        // Test parameter format with password
72        let conn_str = "host=localhost port=5432 user=postgres password=secret dbname=mydb";
73        assert_eq!(
74            sanitize_connection_string(conn_str),
75            "host=localhost port=5432 user=postgres dbname=mydb"
76        );
77
78        // Test parameter format without password
79        let conn_str = "host=localhost port=5432 user=postgres dbname=mydb";
80        assert_eq!(
81            sanitize_connection_string(conn_str),
82            "host=localhost port=5432 user=postgres dbname=mydb"
83        );
84    }
85}