yb_tokio_postgres/
connect_raw.rs

1use crate::codec::{BackendMessage, BackendMessages, FrontendMessage, PostgresCodec};
2use crate::config::{self, Config};
3use crate::connect_tls::connect_tls;
4use crate::maybe_tls_stream::MaybeTlsStream;
5use crate::tls::{TlsConnect, TlsStream};
6use crate::{Client, Connection, Error};
7use bytes::BytesMut;
8use fallible_iterator::FallibleIterator;
9use futures_channel::mpsc;
10use futures_util::{ready, Sink, SinkExt, Stream, TryStreamExt};
11use postgres_protocol::authentication;
12use postgres_protocol::authentication::sasl;
13use postgres_protocol::authentication::sasl::ScramSha256;
14use postgres_protocol::message::backend::{AuthenticationSaslBody, Message};
15use postgres_protocol::message::frontend;
16use std::borrow::Cow;
17use std::collections::{HashMap, VecDeque};
18use std::io;
19use std::pin::Pin;
20use std::task::{Context, Poll};
21use tokio::io::{AsyncRead, AsyncWrite};
22use tokio_util::codec::Framed;
23
24pub struct StartupStream<S, T> {
25    inner: Framed<MaybeTlsStream<S, T>, PostgresCodec>,
26    buf: BackendMessages,
27    delayed: VecDeque<BackendMessage>,
28}
29
30impl<S, T> Sink<FrontendMessage> for StartupStream<S, T>
31where
32    S: AsyncRead + AsyncWrite + Unpin,
33    T: AsyncRead + AsyncWrite + Unpin,
34{
35    type Error = io::Error;
36
37    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
38        Pin::new(&mut self.inner).poll_ready(cx)
39    }
40
41    fn start_send(mut self: Pin<&mut Self>, item: FrontendMessage) -> io::Result<()> {
42        Pin::new(&mut self.inner).start_send(item)
43    }
44
45    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
46        Pin::new(&mut self.inner).poll_flush(cx)
47    }
48
49    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
50        Pin::new(&mut self.inner).poll_close(cx)
51    }
52}
53
54impl<S, T> Stream for StartupStream<S, T>
55where
56    S: AsyncRead + AsyncWrite + Unpin,
57    T: AsyncRead + AsyncWrite + Unpin,
58{
59    type Item = io::Result<Message>;
60
61    fn poll_next(
62        mut self: Pin<&mut Self>,
63        cx: &mut Context<'_>,
64    ) -> Poll<Option<io::Result<Message>>> {
65        loop {
66            match self.buf.next() {
67                Ok(Some(message)) => return Poll::Ready(Some(Ok(message))),
68                Ok(None) => {}
69                Err(e) => return Poll::Ready(Some(Err(e))),
70            }
71
72            match ready!(Pin::new(&mut self.inner).poll_next(cx)) {
73                Some(Ok(BackendMessage::Normal { messages, .. })) => self.buf = messages,
74                Some(Ok(BackendMessage::Async(message))) => return Poll::Ready(Some(Ok(message))),
75                Some(Err(e)) => return Poll::Ready(Some(Err(e))),
76                None => return Poll::Ready(None),
77            }
78        }
79    }
80}
81
82pub async fn connect_raw<S, T>(
83    stream: S,
84    tls: T,
85    has_hostname: bool,
86    config: &Config,
87) -> Result<(Client, Connection<S, T::Stream>), Error>
88where
89    S: AsyncRead + AsyncWrite + Unpin,
90    T: TlsConnect<S>,
91{
92    let stream = connect_tls(stream, config.ssl_mode, tls, has_hostname).await?;
93
94    let mut stream = StartupStream {
95        inner: Framed::new(stream, PostgresCodec),
96        buf: BackendMessages::empty(),
97        delayed: VecDeque::new(),
98    };
99
100    let user = config
101        .user
102        .as_deref()
103        .map_or_else(|| Cow::Owned(whoami::username()), Cow::Borrowed);
104
105    startup(&mut stream, config, &user).await?;
106    authenticate(&mut stream, config, &user).await?;
107    let (process_id, secret_key, parameters) = read_info(&mut stream).await?;
108
109    let (sender, receiver) = mpsc::unbounded();
110    let client = Client::new(sender, config.ssl_mode, process_id, secret_key);
111    let connection = Connection::new(stream.inner, stream.delayed, parameters, receiver);
112
113    Ok((client, connection))
114}
115
116async fn startup<S, T>(
117    stream: &mut StartupStream<S, T>,
118    config: &Config,
119    user: &str,
120) -> Result<(), Error>
121where
122    S: AsyncRead + AsyncWrite + Unpin,
123    T: AsyncRead + AsyncWrite + Unpin,
124{
125    let mut params = vec![("client_encoding", "UTF8")];
126    params.push(("user", user));
127    if let Some(dbname) = &config.dbname {
128        params.push(("database", &**dbname));
129    }
130    if let Some(options) = &config.options {
131        params.push(("options", &**options));
132    }
133    if let Some(application_name) = &config.application_name {
134        params.push(("application_name", &**application_name));
135    }
136
137    let mut buf = BytesMut::new();
138    frontend::startup_message(params, &mut buf).map_err(Error::encode)?;
139
140    stream
141        .send(FrontendMessage::Raw(buf.freeze()))
142        .await
143        .map_err(Error::io)
144}
145
146async fn authenticate<S, T>(
147    stream: &mut StartupStream<S, T>,
148    config: &Config,
149    user: &str,
150) -> Result<(), Error>
151where
152    S: AsyncRead + AsyncWrite + Unpin,
153    T: TlsStream + Unpin,
154{
155    match stream.try_next().await.map_err(Error::io)? {
156        Some(Message::AuthenticationOk) => {
157            can_skip_channel_binding(config)?;
158            return Ok(());
159        }
160        Some(Message::AuthenticationCleartextPassword) => {
161            can_skip_channel_binding(config)?;
162
163            let pass = config
164                .password
165                .as_ref()
166                .ok_or_else(|| Error::config("password missing".into()))?;
167
168            authenticate_password(stream, pass).await?;
169        }
170        Some(Message::AuthenticationMd5Password(body)) => {
171            can_skip_channel_binding(config)?;
172
173            let pass = config
174                .password
175                .as_ref()
176                .ok_or_else(|| Error::config("password missing".into()))?;
177
178            let output = authentication::md5_hash(user.as_bytes(), pass, body.salt());
179            authenticate_password(stream, output.as_bytes()).await?;
180        }
181        Some(Message::AuthenticationSasl(body)) => {
182            authenticate_sasl(stream, body, config).await?;
183        }
184        Some(Message::AuthenticationKerberosV5)
185        | Some(Message::AuthenticationScmCredential)
186        | Some(Message::AuthenticationGss)
187        | Some(Message::AuthenticationSspi) => {
188            return Err(Error::authentication(
189                "unsupported authentication method".into(),
190            ))
191        }
192        Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
193        Some(_) => return Err(Error::unexpected_message()),
194        None => return Err(Error::closed()),
195    }
196
197    match stream.try_next().await.map_err(Error::io)? {
198        Some(Message::AuthenticationOk) => Ok(()),
199        Some(Message::ErrorResponse(body)) => Err(Error::db(body)),
200        Some(_) => Err(Error::unexpected_message()),
201        None => Err(Error::closed()),
202    }
203}
204
205fn can_skip_channel_binding(config: &Config) -> Result<(), Error> {
206    match config.channel_binding {
207        config::ChannelBinding::Disable | config::ChannelBinding::Prefer => Ok(()),
208        config::ChannelBinding::Require => Err(Error::authentication(
209            "server did not use channel binding".into(),
210        )),
211    }
212}
213
214async fn authenticate_password<S, T>(
215    stream: &mut StartupStream<S, T>,
216    password: &[u8],
217) -> Result<(), Error>
218where
219    S: AsyncRead + AsyncWrite + Unpin,
220    T: AsyncRead + AsyncWrite + Unpin,
221{
222    let mut buf = BytesMut::new();
223    frontend::password_message(password, &mut buf).map_err(Error::encode)?;
224
225    stream
226        .send(FrontendMessage::Raw(buf.freeze()))
227        .await
228        .map_err(Error::io)
229}
230
231async fn authenticate_sasl<S, T>(
232    stream: &mut StartupStream<S, T>,
233    body: AuthenticationSaslBody,
234    config: &Config,
235) -> Result<(), Error>
236where
237    S: AsyncRead + AsyncWrite + Unpin,
238    T: TlsStream + Unpin,
239{
240    let password = config
241        .password
242        .as_ref()
243        .ok_or_else(|| Error::config("password missing".into()))?;
244
245    let mut has_scram = false;
246    let mut has_scram_plus = false;
247    let mut mechanisms = body.mechanisms();
248    while let Some(mechanism) = mechanisms.next().map_err(Error::parse)? {
249        match mechanism {
250            sasl::SCRAM_SHA_256 => has_scram = true,
251            sasl::SCRAM_SHA_256_PLUS => has_scram_plus = true,
252            _ => {}
253        }
254    }
255
256    let channel_binding = stream
257        .inner
258        .get_ref()
259        .channel_binding()
260        .tls_server_end_point
261        .filter(|_| config.channel_binding != config::ChannelBinding::Disable)
262        .map(sasl::ChannelBinding::tls_server_end_point);
263
264    let (channel_binding, mechanism) = if has_scram_plus {
265        match channel_binding {
266            Some(channel_binding) => (channel_binding, sasl::SCRAM_SHA_256_PLUS),
267            None => (sasl::ChannelBinding::unsupported(), sasl::SCRAM_SHA_256),
268        }
269    } else if has_scram {
270        match channel_binding {
271            Some(_) => (sasl::ChannelBinding::unrequested(), sasl::SCRAM_SHA_256),
272            None => (sasl::ChannelBinding::unsupported(), sasl::SCRAM_SHA_256),
273        }
274    } else {
275        return Err(Error::authentication("unsupported SASL mechanism".into()));
276    };
277
278    if mechanism != sasl::SCRAM_SHA_256_PLUS {
279        can_skip_channel_binding(config)?;
280    }
281
282    let mut scram = ScramSha256::new(password, channel_binding);
283
284    let mut buf = BytesMut::new();
285    frontend::sasl_initial_response(mechanism, scram.message(), &mut buf).map_err(Error::encode)?;
286    stream
287        .send(FrontendMessage::Raw(buf.freeze()))
288        .await
289        .map_err(Error::io)?;
290
291    let body = match stream.try_next().await.map_err(Error::io)? {
292        Some(Message::AuthenticationSaslContinue(body)) => body,
293        Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
294        Some(_) => return Err(Error::unexpected_message()),
295        None => return Err(Error::closed()),
296    };
297
298    scram
299        .update(body.data())
300        .map_err(|e| Error::authentication(e.into()))?;
301
302    let mut buf = BytesMut::new();
303    frontend::sasl_response(scram.message(), &mut buf).map_err(Error::encode)?;
304    stream
305        .send(FrontendMessage::Raw(buf.freeze()))
306        .await
307        .map_err(Error::io)?;
308
309    let body = match stream.try_next().await.map_err(Error::io)? {
310        Some(Message::AuthenticationSaslFinal(body)) => body,
311        Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
312        Some(_) => return Err(Error::unexpected_message()),
313        None => return Err(Error::closed()),
314    };
315
316    scram
317        .finish(body.data())
318        .map_err(|e| Error::authentication(e.into()))?;
319
320    Ok(())
321}
322
323async fn read_info<S, T>(
324    stream: &mut StartupStream<S, T>,
325) -> Result<(i32, i32, HashMap<String, String>), Error>
326where
327    S: AsyncRead + AsyncWrite + Unpin,
328    T: AsyncRead + AsyncWrite + Unpin,
329{
330    let mut process_id = 0;
331    let mut secret_key = 0;
332    let mut parameters = HashMap::new();
333
334    loop {
335        match stream.try_next().await.map_err(Error::io)? {
336            Some(Message::BackendKeyData(body)) => {
337                process_id = body.process_id();
338                secret_key = body.secret_key();
339            }
340            Some(Message::ParameterStatus(body)) => {
341                parameters.insert(
342                    body.name().map_err(Error::parse)?.to_string(),
343                    body.value().map_err(Error::parse)?.to_string(),
344                );
345            }
346            Some(msg @ Message::NoticeResponse(_)) => {
347                stream.delayed.push_back(BackendMessage::Async(msg))
348            }
349            Some(Message::ReadyForQuery(_)) => return Ok((process_id, secret_key, parameters)),
350            Some(Message::ErrorResponse(body)) => return Err(Error::db(body)),
351            Some(_) => return Err(Error::unexpected_message()),
352            None => return Err(Error::closed()),
353        }
354    }
355}
yb_tokio_postgres/connect_raw.rs

yb_tokio_postgres/
connect_raw.rs
