Implement support for SCRAM auth to sqlc.

Author: BearLemma

Cargo.lock

@@ -11,9 +11,11 @@ crate-type = ["cdylib"]
 anyhow = "1.0.66"
 bytes = "1.3.0"
 fallible-iterator = "0.2.0"
+fn-error-context = "0.2.0"
 postgres-protocol = "0.6.4"
 postgres-types = "0.2.4"
 tracing = "0.1.37"
 tracing-subscriber = "0.3.16"
 unwrap_or = "1.0.0"
 url = { version = "2.3.1", default-features = false }
+scram = "0.6.0"

sqlc/Cargo.toml

@@ -12,6 +12,7 @@ use postgres_protocol::message::backend::DataRowBody;
 use std::cell::RefCell;
 use std::collections::VecDeque;
 use std::ffi::{CStr, CString};
+use std::fmt::Debug;
 use std::ops::Range;
 use std::os::raw::{c_char, c_int, c_void};
 use std::rc::Rc;
@@ -22,9 +23,9 @@ thread_local! {
     static ERRMSG: RefCell<Option<CString>> = RefCell::new(None);
 }
 
-fn set_error_message<T: ToString>(e: T) {
+fn set_error_message<T: Debug>(e: T) {
     ERRMSG.with(|errmsg| {
-        errmsg.replace(Some(CString::new(e.to_string()).unwrap()));
+        errmsg.replace(Some(CString::new(format!("{e:?}")).unwrap()));
     });
 }
 

sqlc/src/lib.rs

@@ -1,6 +1,7 @@
 use anyhow::{Context, Result};
 use bytes::BytesMut;
 use fallible_iterator::FallibleIterator;
+use fn_error_context::context as fn_context;
 use postgres_protocol::message::backend::DataRowBody;
 use postgres_protocol::message::{backend, frontend};
 use postgres_types::Type;
@@ -29,13 +30,15 @@ pub struct Connection {
     stream: TcpStream,
     rx_buf: BytesMut,
     username: String,
+    password: Option<String>,
 }
 
 impl Connection {
     pub fn connect(addr: &str) -> Result<Self> {
         let url = Url::parse(addr)?;
         let host = url.host_str().unwrap();
         let port = url.port().unwrap();
+        let password = url.password().map(|p| p.to_owned());
         let stream = TcpStream::connect((host, port))
             .with_context(|| format!("Unable to connect to {addr}"))?;
         let rx_buf = BytesMut::with_capacity(1024);
@@ -44,6 +47,7 @@ impl Connection {
             stream,
             rx_buf,
             username,
+            password,
         })
     }
 
@@ -66,13 +70,21 @@ impl Connection {
     pub fn wait_until_ready(&mut self) -> Result<(Metadata, VecDeque<DataRowBody>)> {
         let mut metadata = Metadata::new();
         let mut rows = VecDeque::default();
+        loop {
+            let msg = self.receive_message()?;
+            if !self.process_msg(msg, &mut metadata, &mut rows)? {
+                return Ok((metadata, rows));
+            }
+        }
+    }
+
+    #[fn_context("failed to receive the next message from postgres server")]
+    fn receive_message(&mut self) -> Result<backend::Message> {
         loop {
             let msg = backend::Message::parse(&mut self.rx_buf)?;
             match msg {
                 Some(msg) => {
-                    if !self.process_msg(msg, &mut metadata, &mut rows) {
-                        return Ok((metadata, rows));
-                    }
+                    return Ok(msg);
                 }
                 None => {
                     // FIXME: Optimize with spare_capacity_mut() to make zero-copy.
@@ -89,7 +101,7 @@ impl Connection {
         msg: backend::Message,
         metadata: &mut Metadata,
         rows: &mut VecDeque<DataRowBody>,
-    ) -> bool {
+    ) -> Result<bool> {
         match msg {
             backend::Message::AuthenticationCleartextPassword => todo!(),
             backend::Message::AuthenticationGss => todo!(),
@@ -101,7 +113,10 @@ impl Connection {
             backend::Message::AuthenticationScmCredential => todo!(),
             backend::Message::AuthenticationSspi => todo!(),
             backend::Message::AuthenticationGssContinue(_) => todo!(),
-            backend::Message::AuthenticationSasl(_) => todo!(),
+            backend::Message::AuthenticationSasl(body) => {
+                trace!("TRACE postgres -> AuthenticationSasl");
+                self.run_sasl_auth(body)?;
+            }
             backend::Message::AuthenticationSaslContinue(_) => todo!(),
             backend::Message::AuthenticationSaslFinal(_) => todo!(),
             backend::Message::BackendKeyData(_) => {
@@ -121,8 +136,9 @@ impl Connection {
                 rows.push_back(row);
             }
             backend::Message::EmptyQueryResponse => todo!(),
-            backend::Message::ErrorResponse(_) => {
+            backend::Message::ErrorResponse(body) => {
                 trace!("TRACE postgres -> ErrorResponse");
+                anyhow::bail!(self.parse_err(body)?)
             }
             backend::Message::NoData => todo!(),
             backend::Message::NoticeResponse(_) => {
@@ -137,19 +153,80 @@ impl Connection {
             backend::Message::PortalSuspended => todo!(),
             backend::Message::ReadyForQuery(_) => {
                 trace!("TRACE postgres -> ReadyForQuery");
-                return false;
+                return Ok(false);
             }
             backend::Message::RowDescription(row_description) => {
                 trace!("TRACE postgres -> RowDescription");
                 let mut fields = row_description.fields();
-                while let Some(field) = fields.next().unwrap() {
+                while let Some(field) = fields.next()? {
                     metadata.col_names.push(field.name().into());
                     let ty = Type::from_oid(field.type_oid()).unwrap();
                     metadata.col_types.push(ty);
                 }
             }
             _ => todo!(),
         }
-        true
+        Ok(true)
+    }
+
+    #[fn_context("failed to authenticate to SQL server using SASL authentication protocol")]
+    fn run_sasl_auth(&mut self, body: backend::AuthenticationSaslBody) -> Result<()> {
+        let mechanisms: Vec<_> = body.mechanisms().collect()?;
+        anyhow::ensure!(
+            mechanisms.contains(&"SCRAM-SHA-256"),
+            "our client supports only 'SCRAM-SHA-256' SASL auth protocol, but the server supports only {mechanisms:?}"
+        );
+
+        let username = self.username.clone();
+        let password = self
+            .password
+            .clone()
+            .context("password must be provided when server enforces SASL auth")?;
+        let scram = scram::ScramClient::new(&username, &password, None);
+        let (scram, cli_message) = scram.client_first();
+
+        let mut buff = BytesMut::new();
+        frontend::sasl_initial_response("SCRAM-SHA-256", cli_message.as_bytes(), &mut buff)?;
+        self.stream.write_all(&buff)?;
+
+        trace!("TRACE postgres -> AuthenticationSasl -> client first message sent");
+
+        let body = match self.receive_message()? {
+            backend::Message::AuthenticationSaslContinue(body) => body,
+            backend::Message::ErrorResponse(body) => anyhow::bail!(self.parse_err(body)?),
+            _ => anyhow::bail!(
+                "received unexpected message from server. Expected 'AuthenticationSaslContinue'.",
+            ),
+        };
+
+        let scram = scram.handle_server_first(std::str::from_utf8(body.data())?)?;
+        let (scram, client_final) = scram.client_final();
+
+        buff.clear();
+        frontend::sasl_response(client_final.as_bytes(), &mut buff)?;
+        self.stream.write_all(&buff)?;
+
+        // Receive the last message from server.
+        let body = match self.receive_message()? {
+            backend::Message::AuthenticationSaslFinal(body) => body,
+            backend::Message::ErrorResponse(body) => anyhow::bail!(self.parse_err(body)?),
+            _ => anyhow::bail!(
+                "received unexpected message from server. Expected 'AuthenticationSaslFinal'.",
+            ),
+        };
+
+        // Checks the final response from the server
+        scram.handle_server_final(std::str::from_utf8(body.data())?)?;
+
+        trace!("TRACE postgres -> AuthenticationSasl -> authentication successful");
+        Ok(())
+    }
+
+    fn parse_err(&self, body: backend::ErrorResponseBody) -> Result<String> {
+        let err_fields: Vec<_> = body.fields().map(|f| Ok(f.value().to_string())).collect()?;
+        let err_msg =
+            format!("server responded with error response. Provided error fields: {err_fields:?}");
+        trace!("TRACE postgres -> Error ocurred: {err_msg}");
+        Ok(err_msg)
     }
 }

sqlc/src/postgres.rs

@@ -1,3 +1,5 @@
+DB_URI ?= 'postgres://127.0.0.1:5432'
+
 all: setup test
 .PHONY: all
 
@@ -6,5 +8,5 @@ setup:
 .PHONY: setup
 
 test:
-	LD_PRELOAD=../../../target/debug/libsqlc.so DB_URI=postgres://127.0.0.1:5432 bundle exec rspec sqlite_spec.rb
+	LD_PRELOAD=../../../target/debug/libsqlc.so DB_URI=$(DB_URI) bundle exec rspec sqlite_spec.rb
 .PHONY: test

testing/client/ruby/Makefile

@@ -11,3 +11,11 @@ bundle install
 ```console
 bundle exec rspec sqlite_spec.rb
 ```
+
+The default database URL can be configured using DB_URI env variable. It's especially
+important if your local postgres requires authentication. In that case, you
+can use
+
+```console
+DB_URI=postgres://asd:[email protected]:5432 bundle exec rspec sqlite_spec.rb
+````

testing/client/ruby/README.md

@@ -1,3 +1,5 @@
+DB_URI ?= 'postgres://127.0.0.1:5432'
+
 all: setup test
 .PHONY: all
 
@@ -6,5 +8,5 @@ setup:
 .PHONY: setup
 
 test:
-	bundle exec rspec postgresql_spec.rb
+	DB_URI=$(DB_URI) bundle exec rspec postgresql_spec.rb
 .PHONY: test

testing/server/ruby/Makefile

@@ -11,3 +11,11 @@ bundle install
 ```console
 bundle exec rspec postgresql_spec.rb
 ```
+
+The default database URL can be configured using DB_URI env variable. It's especially
+important if your local postgres requires authentication. In that case, you
+can use
+
+```console
+DB_URI=postgres://asd:[email protected]:5432 bundle exec rspec postgresql_spec.rb
+````

testing/server/ruby/README.md

@@ -1,17 +1,23 @@
 require "pg"
 
+db_uri = ENV["DB_URI"]
+
+if db_uri.nil?
+    raise "Please configure database via the `DB_URI` environment variable."
+end
+
 describe "PostgreSQL client" do
   it "connects" do
-    conn = PG.connect(host: "127.0.0.1", port: 5432)
+    conn = PG.connect(db_uri)
   end
 
   it "performs schema changes" do
-    conn = PG.connect(host: "127.0.0.1", port: 5432)
+    conn = PG.connect(db_uri)
     conn.exec("CREATE TABLE IF NOT EXISTS users (username TEXT, pass TEXT)")
   end
 
   it "queries tables" do
-    conn = PG.connect(host: "127.0.0.1", port: 5432)
+    conn = PG.connect(db_uri)
     conn.exec("CREATE TABLE IF NOT EXISTS users (username TEXT, pass TEXT)")
     conn.exec("DELETE FROM users")
     conn.exec("INSERT INTO users VALUES ('me', 'my_pass')")