Skip to main content

kellnr_db/
con_string.rs

1use std::fmt::{Display, Formatter};
2use std::path::{Path, PathBuf};
3
4use kellnr_settings::{SessionAge, Settings};
5use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, utf8_percent_encode};
6
7#[derive(Debug, Clone, PartialEq, Eq, Hash)]
8pub enum ConString {
9    Postgres(PgConString),
10    Sqlite(SqliteConString),
11}
12
13impl Display for ConString {
14    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
15        let con_string = match self {
16            ConString::Postgres(p) => p.to_string(),
17            ConString::Sqlite(s) => s.to_string(),
18        };
19        write!(f, "{con_string}")
20    }
21}
22
23impl ConString {
24    pub fn admin_pwd(&self) -> String {
25        match self {
26            ConString::Postgres(p) => p.admin.pwd.clone(),
27            ConString::Sqlite(s) => s.admin_pwd.clone(),
28        }
29    }
30
31    pub fn admin_token(&self) -> Option<String> {
32        match self {
33            ConString::Postgres(p) => p.admin.token.clone(),
34            ConString::Sqlite(s) => s.admin_token.clone(),
35        }
36    }
37
38    /// Maximum lifetime of a session before it is considered expired.
39    pub fn session_age(&self) -> SessionAge {
40        match self {
41            ConString::Postgres(p) => p.session_age,
42            ConString::Sqlite(s) => s.session_age,
43        }
44    }
45}
46
47#[derive(Debug, Clone, PartialEq, Eq, Hash)]
48pub struct AdminUser {
49    pub pwd: String,
50    pub token: Option<String>,
51}
52
53impl AdminUser {
54    pub fn new(pwd: String, token: Option<String>) -> Self {
55        Self { pwd, token }
56    }
57}
58
59#[derive(Debug, Clone, PartialEq, Eq, Hash)]
60pub(crate) enum PgConStringInner {
61    /// A connection string constructed from individual fields.
62    Structured {
63        addr: String,
64        port: u16,
65        db: String,
66        user: String,
67        pwd: String,
68    },
69    /// A raw URL that gets passed directly to sqlx.
70    Raw(String),
71}
72
73#[derive(Debug, Clone, PartialEq, Eq, Hash)]
74pub struct PgConString {
75    inner: PgConStringInner,
76    admin: AdminUser,
77    pub session_age: SessionAge,
78}
79
80impl PgConString {
81    pub fn new(
82        addr: &str,
83        port: u16,
84        db: &str,
85        user: &str,
86        pwd: &str,
87        admin: AdminUser,
88        session_age: SessionAge,
89    ) -> Self {
90        Self {
91            inner: PgConStringInner::Structured {
92                addr: addr.to_owned(),
93                port,
94                db: db.to_owned(),
95                user: user.to_owned(),
96                pwd: pwd.to_owned(),
97            },
98            admin,
99            session_age,
100        }
101    }
102}
103
104impl From<&Settings> for PgConString {
105    fn from(s: &Settings) -> Self {
106        let inner = if let Some(string) = s.postgresql.connection_string.clone() {
107            PgConStringInner::Raw(string)
108        } else {
109            PgConStringInner::Structured {
110                addr: s.postgresql.address.clone(),
111                port: s.postgresql.port,
112                db: s.postgresql.db.clone(),
113                user: s.postgresql.user.clone(),
114                pwd: s.postgresql.pwd.clone(),
115            }
116        };
117
118        Self {
119            inner,
120            admin: AdminUser {
121                pwd: s.setup.admin_pwd.clone(),
122                token: s.setup.admin_token.clone(),
123            },
124            session_age: s.registry.session_age(),
125        }
126    }
127}
128
129/// Characters that must be percent-encoded in the user, password and database
130/// parts of the connection URL. Everything except the RFC 3986 unreserved characters.
131const URL_COMPONENT: &AsciiSet = &NON_ALPHANUMERIC
132    .remove(b'-')
133    .remove(b'.')
134    .remove(b'_')
135    .remove(b'~');
136
137impl Display for PgConString {
138    fn fmt(&self, f: &mut Formatter) -> Result<(), std::fmt::Error> {
139        match &self.inner {
140            PgConStringInner::Raw(s) => write!(f, "{s}"),
141            PgConStringInner::Structured {
142                user,
143                pwd,
144                addr,
145                port,
146                db,
147            } => {
148                write!(
149                    f,
150                    "postgres://{}:{}@{}:{}/{}",
151                    utf8_percent_encode(user, URL_COMPONENT),
152                    utf8_percent_encode(pwd, URL_COMPONENT),
153                    addr,
154                    port,
155                    utf8_percent_encode(db, URL_COMPONENT)
156                )
157            }
158        }
159    }
160}
161
162#[derive(Debug, Clone, PartialEq, Eq, Hash)]
163pub struct SqliteConString {
164    pub path: PathBuf,
165    pub admin_pwd: String,
166    pub admin_token: Option<String>,
167    pub session_age: SessionAge,
168}
169
170impl SqliteConString {
171    pub fn new(
172        path: &Path,
173        admin_pwd: &str,
174        admin_token: Option<String>,
175        session_age: SessionAge,
176    ) -> Self {
177        Self {
178            path: path.to_owned(),
179            admin_pwd: admin_pwd.to_owned(),
180            admin_token,
181            session_age,
182        }
183    }
184}
185
186impl From<&Settings> for SqliteConString {
187    fn from(settings: &Settings) -> Self {
188        Self {
189            path: settings.sqlite_path(),
190            admin_pwd: settings.setup.admin_pwd.clone(),
191            admin_token: settings.setup.admin_token.clone(),
192            session_age: settings.registry.session_age(),
193        }
194    }
195}
196
197impl Display for SqliteConString {
198    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
199        if self.path.to_str() == Some(":memory:") {
200            write!(f, "sqlite::memory:")
201        } else {
202            write!(f, "sqlite://{}?mode=rwc", self.path.display())
203        }
204    }
205}
206
207#[cfg(test)]
208mod tests {
209    use std::str::FromStr;
210
211    use percent_encoding::percent_decode_str;
212    use sea_orm::sqlx::postgres::PgConnectOptions;
213    use url::Url;
214
215    use super::*;
216
217    fn pg_con_string(user: &str, pwd: &str, db: &str) -> PgConString {
218        PgConString::new(
219            "localhost",
220            5432,
221            db,
222            user,
223            pwd,
224            AdminUser::new("admin".to_string(), None),
225            SessionAge::from_secs(60),
226        )
227    }
228
229    #[test]
230    fn pg_con_string_plain_credentials_are_unchanged() {
231        let con = pg_con_string("kellnr", "secret", "kellnr-db");
232
233        assert_eq!(
234            con.to_string(),
235            "postgres://kellnr:secret@localhost:5432/kellnr-db"
236        );
237    }
238
239    #[test]
240    fn pg_con_string_encodes_special_characters() {
241        let con = pg_con_string("us@er", "a/b@c:d%e#f?g", "my db");
242
243        assert_eq!(
244            con.to_string(),
245            "postgres://us%40er:a%2Fb%40c%3Ad%25e%23f%3Fg@localhost:5432/my%20db"
246        );
247    }
248
249    #[test]
250    fn pg_con_string_round_trips_through_sqlx() {
251        let user = "us@er:name";
252        let pwd = "cDlJH15F/Xj@:%2F#?&=+ รครถ";
253        let db = "kellnr/db";
254        let con = pg_con_string(user, pwd, db);
255
256        let opts = PgConnectOptions::from_str(&con.to_string()).unwrap();
257        // sqlx has no getter for the password, so decode it the same way sqlx does.
258        let url = Url::parse(&con.to_string()).unwrap();
259        let decoded_pwd = percent_decode_str(url.password().unwrap())
260            .decode_utf8()
261            .unwrap();
262
263        assert_eq!(opts.get_username(), user);
264        assert_eq!(decoded_pwd, pwd);
265        assert_eq!(opts.get_host(), "localhost");
266        assert_eq!(opts.get_port(), 5432);
267        assert_eq!(opts.get_database(), Some(db));
268    }
269}