TLA Line data Source code
1 : use anyhow::{bail, Context};
2 : use itertools::Itertools;
3 : use std::borrow::Cow;
4 : use std::fmt;
5 : use url::Host;
6 :
7 : /// Parses a string of format either `host:port` or `host` into a corresponding pair.
8 : /// The `host` part should be a correct `url::Host`, while `port` (if present) should be
9 : /// a valid decimal u16 of digits only.
10 CBC 1542621 : pub fn parse_host_port<S: AsRef<str>>(host_port: S) -> Result<(Host, Option<u16>), anyhow::Error> {
11 1542621 : let (host, port) = match host_port.as_ref().rsplit_once(':') {
12 1542606 : Some((host, port)) => (
13 1542606 : host,
14 1542606 : // +80 is a valid u16, but not a valid port
15 7713022 : if port.chars().all(|c| c.is_ascii_digit()) {
16 1542605 : Some(port.parse::<u16>().context("Unable to parse port")?)
17 : } else {
18 1 : bail!("Port contains a non-ascii-digit")
19 : },
20 : ),
21 15 : None => (host_port.as_ref(), None), // No colons, no port specified
22 : };
23 1542620 : let host = Host::parse(host).context("Unable to parse host")?;
24 1542619 : Ok((host, port))
25 1542621 : }
26 :
27 : #[cfg(test)]
28 : mod tests_parse_host_port {
29 : use crate::parse_host_port;
30 : use url::Host;
31 :
32 1 : #[test]
33 1 : fn test_normal() {
34 1 : let (host, port) = parse_host_port("hello:123").unwrap();
35 1 : assert_eq!(host, Host::Domain("hello".to_owned()));
36 1 : assert_eq!(port, Some(123));
37 1 : }
38 :
39 1 : #[test]
40 1 : fn test_no_port() {
41 1 : let (host, port) = parse_host_port("hello").unwrap();
42 1 : assert_eq!(host, Host::Domain("hello".to_owned()));
43 1 : assert_eq!(port, None);
44 1 : }
45 :
46 1 : #[test]
47 1 : fn test_ipv6() {
48 1 : let (host, port) = parse_host_port("[::1]:123").unwrap();
49 1 : assert_eq!(host, Host::<String>::Ipv6(std::net::Ipv6Addr::LOCALHOST));
50 1 : assert_eq!(port, Some(123));
51 1 : }
52 :
53 1 : #[test]
54 1 : fn test_invalid_host() {
55 1 : assert!(parse_host_port("hello world").is_err());
56 1 : }
57 :
58 1 : #[test]
59 1 : fn test_invalid_port() {
60 1 : assert!(parse_host_port("hello:+80").is_err());
61 1 : }
62 : }
63 :
64 776907 : #[derive(Clone)]
65 : pub struct PgConnectionConfig {
66 : host: Host,
67 : port: u16,
68 : password: Option<String>,
69 : options: Vec<String>,
70 : }
71 :
72 : /// A simplified PostgreSQL connection configuration. Supports only a subset of possible
73 : /// settings for simplicity. A password getter or `to_connection_string` methods are not
74 : /// added by design to avoid accidentally leaking password through logging, command line
75 : /// arguments to a child process, or likewise.
76 : impl PgConnectionConfig {
77 1543652 : pub fn new_host_port(host: Host, port: u16) -> Self {
78 1543652 : PgConnectionConfig {
79 1543652 : host,
80 1543652 : port,
81 1543652 : password: None,
82 1543652 : options: vec![],
83 1543652 : }
84 1543652 : }
85 :
86 3259 : pub fn host(&self) -> &Host {
87 3259 : &self.host
88 3259 : }
89 :
90 2058 : pub fn port(&self) -> u16 {
91 2058 : self.port
92 2058 : }
93 :
94 UBC 0 : pub fn set_host(mut self, h: Host) -> Self {
95 0 : self.host = h;
96 0 : self
97 0 : }
98 :
99 0 : pub fn set_port(mut self, p: u16) -> Self {
100 0 : self.port = p;
101 0 : self
102 0 : }
103 :
104 CBC 1533969 : pub fn set_password(mut self, s: Option<String>) -> Self {
105 1533969 : self.password = s;
106 1533969 : self
107 1533969 : }
108 :
109 1534783 : pub fn extend_options<I: IntoIterator<Item = S>, S: Into<String>>(mut self, i: I) -> Self {
110 4602722 : self.options.extend(i.into_iter().map(|s| s.into()));
111 1534783 : self
112 1534783 : }
113 :
114 : /// Return a `<host>:<port>` string.
115 1413 : pub fn raw_address(&self) -> String {
116 1413 : format!("{}:{}", self.host(), self.port())
117 1413 : }
118 :
119 : /// Build a client library-specific connection configuration.
120 : /// Used for testing and when we need to add some obscure configuration
121 : /// elements at the last moment.
122 1195 : pub fn to_tokio_postgres_config(&self) -> tokio_postgres::Config {
123 1195 : // Use `tokio_postgres::Config` instead of `postgres::Config` because
124 1195 : // the former supports more options to fiddle with later.
125 1195 : let mut config = tokio_postgres::Config::new();
126 1195 : config.host(&self.host().to_string()).port(self.port);
127 1195 : if let Some(password) = &self.password {
128 52 : config.password(password);
129 1143 : }
130 1195 : if !self.options.is_empty() {
131 1190 : // These options are command-line options and should be escaped before being passed
132 1190 : // as an 'options' connection string parameter, see
133 1190 : // https://www.postgresql.org/docs/15/libpq-connect.html#LIBPQ-CONNECT-OPTIONS
134 1190 : //
135 1190 : // They will be space-separated, so each space inside an option should be escaped,
136 1190 : // and all backslashes should be escaped before that. Although we don't expect options
137 1190 : // with spaces at the moment, they're supported by PostgreSQL. Hence we support them
138 1190 : // in this typesafe interface.
139 1190 : //
140 1190 : // We use `Cow` to avoid allocations in the best case (no escaping). A fully imperative
141 1190 : // solution would require 1-2 allocations in the worst case as well, but it's harder to
142 1190 : // implement and this function is hardly a bottleneck. The function is only called around
143 1190 : // establishing a new connection.
144 1190 : #[allow(unstable_name_collisions)]
145 1190 : config.options(
146 1190 : &self
147 1190 : .options
148 1190 : .iter()
149 3582 : .map(|s| {
150 3582 : if s.contains(['\\', ' ']) {
151 2 : Cow::Owned(s.replace('\\', "\\\\").replace(' ', "\\ "))
152 : } else {
153 3580 : Cow::Borrowed(s.as_str())
154 : }
155 3582 : })
156 1190 : .intersperse(Cow::Borrowed(" ")) // TODO: use impl from std once it's stabilized
157 1190 : .collect::<String>(),
158 1190 : );
159 1190 : }
160 1195 : config
161 1195 : }
162 :
163 : /// Connect using postgres protocol with TLS disabled.
164 5 : pub fn connect_no_tls(&self) -> Result<postgres::Client, postgres::Error> {
165 5 : postgres::Config::from(self.to_tokio_postgres_config()).connect(postgres::NoTls)
166 5 : }
167 : }
168 :
169 : impl fmt::Debug for PgConnectionConfig {
170 3803 : fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
171 3803 : // We want `password: Some(REDACTED-STRING)`, not `password: Some("REDACTED-STRING")`
172 3803 : // so even if the password is `REDACTED-STRING` (quite unlikely) there is no confusion.
173 3803 : // Hence `format_args!()`, it returns a "safe" string which is not escaped by `Debug`.
174 3803 : f.debug_struct("PgConnectionConfig")
175 3803 : .field("host", &self.host)
176 3803 : .field("port", &self.port)
177 3803 : .field(
178 3803 : "password",
179 3803 : &self
180 3803 : .password
181 3803 : .as_ref()
182 3803 : .map(|_| format_args!("REDACTED-STRING")),
183 3803 : )
184 3803 : .finish()
185 3803 : }
186 : }
187 :
188 : #[cfg(test)]
189 : mod tests_pg_connection_config {
190 : use crate::PgConnectionConfig;
191 : use once_cell::sync::Lazy;
192 : use url::Host;
193 :
194 1 : static STUB_HOST: Lazy<Host> = Lazy::new(|| Host::Domain("stub.host.example".to_owned()));
195 :
196 1 : #[test]
197 1 : fn test_no_password() {
198 1 : let cfg = PgConnectionConfig::new_host_port(STUB_HOST.clone(), 123);
199 1 : assert_eq!(cfg.host(), &*STUB_HOST);
200 1 : assert_eq!(cfg.port(), 123);
201 1 : assert_eq!(cfg.raw_address(), "stub.host.example:123");
202 1 : assert_eq!(
203 1 : format!("{:?}", cfg),
204 1 : "PgConnectionConfig { host: Domain(\"stub.host.example\"), port: 123, password: None }"
205 1 : );
206 1 : }
207 :
208 1 : #[test]
209 1 : fn test_ipv6() {
210 1 : // May be a special case because hostname contains a colon.
211 1 : let cfg = PgConnectionConfig::new_host_port(Host::parse("[::1]").unwrap(), 123);
212 1 : assert_eq!(
213 1 : cfg.host(),
214 1 : &Host::<String>::Ipv6(std::net::Ipv6Addr::LOCALHOST)
215 1 : );
216 1 : assert_eq!(cfg.port(), 123);
217 1 : assert_eq!(cfg.raw_address(), "[::1]:123");
218 1 : assert_eq!(
219 1 : format!("{:?}", cfg),
220 1 : "PgConnectionConfig { host: Ipv6(::1), port: 123, password: None }"
221 1 : );
222 1 : }
223 :
224 1 : #[test]
225 1 : fn test_with_password() {
226 1 : let cfg = PgConnectionConfig::new_host_port(STUB_HOST.clone(), 123)
227 1 : .set_password(Some("password".to_owned()));
228 1 : assert_eq!(cfg.host(), &*STUB_HOST);
229 1 : assert_eq!(cfg.port(), 123);
230 1 : assert_eq!(cfg.raw_address(), "stub.host.example:123");
231 1 : assert_eq!(
232 1 : format!("{:?}", cfg),
233 1 : "PgConnectionConfig { host: Domain(\"stub.host.example\"), port: 123, password: Some(REDACTED-STRING) }"
234 1 : );
235 1 : }
236 :
237 1 : #[test]
238 1 : fn test_with_options() {
239 1 : let cfg = PgConnectionConfig::new_host_port(STUB_HOST.clone(), 123).extend_options([
240 1 : "hello",
241 1 : "world",
242 1 : "with space",
243 1 : "and \\ backslashes",
244 1 : ]);
245 1 : assert_eq!(cfg.host(), &*STUB_HOST);
246 1 : assert_eq!(cfg.port(), 123);
247 1 : assert_eq!(cfg.raw_address(), "stub.host.example:123");
248 1 : assert_eq!(
249 1 : cfg.to_tokio_postgres_config().get_options(),
250 1 : Some("hello world with\\ space and\\ \\\\\\ backslashes")
251 1 : );
252 1 : }
253 : }
|
