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