Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@ async-trait = "0.1"
connection-string = "0.2"
num-traits = "0.2"
uuid = "1.0"
zeroize = "1.8.2"

[target.'cfg(windows)'.dependencies]
winauth = { version = "0.0.4", optional = true }
Expand Down
86 changes: 86 additions & 0 deletions docker/test-server.sh
Original file line number Diff line number Diff line change
@@ -0,0 +1,86 @@
#!/usr/bin/env bash
#
# Start a SQL Server for the test suite, with podman or docker.
#
# ./docker/test-server.sh up # build, start, wait until it accepts connections
# ./docker/test-server.sh down # stop and remove
# ./docker/test-server.sh logs # follow the server log
#
# Then:
#
# export TIBERIUS_TEST_CONNECTION_STRING='server=tcp:localhost,1433;user=SA;password=<YourStrong@Passw0rd>;IntegratedSecurity=true;TrustServerCertificate=true'
# cargo test
#
# IMAGE selects the flavour; the default works on both x86_64 and arm64.
# The full SQL Server images are x86_64 only, so on an arm64 machine
# (Apple silicon) they either refuse to run or run under emulation.

set -euo pipefail

ENGINE="${ENGINE:-$(command -v podman >/dev/null 2>&1 && echo podman || echo docker)}"
NAME="${NAME:-tiberius-test-mssql}"
PORT="${PORT:-1433}"
IMAGE="${IMAGE:-azure-sql-edge}"
PASSWORD='<YourStrong@Passw0rd>'
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

case "${1:-up}" in
up)
echo "engine: $ENGINE image: $IMAGE port: $PORT"
"$ENGINE" build -q -f "$HERE/docker-$IMAGE.dockerfile" -t "$NAME:local" "$HERE"
"$ENGINE" rm -f "$NAME" >/dev/null 2>&1 || true
"$ENGINE" run -d --name "$NAME" \
-e ACCEPT_EULA=Y \
-e "MSSQL_SA_PASSWORD=$PASSWORD" \
-e "SA_PASSWORD=$PASSWORD" \
-p "$PORT:1433" \
"$NAME:local" >/dev/null

# The port opens well before the server will answer, so poll the log
# rather than the socket.
#
# The log is captured into a variable and matched there, rather than
# piped into `grep -q`. Under `set -o pipefail`, `grep -q` exits on the
# first match, the writer upstream dies of SIGPIPE, and the pipeline
# reports failure even though the match succeeded — so the wait never
# ends.
echo -n "waiting for SQL Server"
for _ in $(seq 1 120); do
logs="$("$ENGINE" logs "$NAME" 2>&1 || true)"

case "$logs" in
*"SQL Server is now ready for client connections"*)
echo " — ready"
exit 0
;;
esac

running="$("$ENGINE" ps --format '{{.Names}}' || true)"
case "$running" in
*"$NAME"*) ;;
*)
echo " — container exited:"
"$ENGINE" logs --tail 30 "$NAME" || true
exit 1
;;
esac

echo -n .
sleep 2
done
echo " — gave up; last lines:"
"$ENGINE" logs --tail 30 "$NAME"
exit 1
;;
down)
"$ENGINE" rm -f "$NAME" >/dev/null 2>&1 || true
echo "removed $NAME"
;;
logs)
"$ENGINE" logs -f "$NAME"
;;
*)
echo "usage: $0 {up|down|logs}" >&2
exit 2
;;
esac
62 changes: 60 additions & 2 deletions src/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -251,10 +251,13 @@ impl<S: AsyncRead + AsyncWrite + Unpin + Send> Client<S> {
Ok(result)
}

/// Execute a `BULK INSERT` statement, efficiantly storing a large number of
/// Execute a `BULK INSERT` statement, efficiently storing a large number of
/// rows to a specified table. Note: make sure the input row follows the same
/// schema as the table, otherwise calling `send()` will return an error.
///
/// This is equivalent to `bulk_insert_columns(table, &["*"])`, inserting into
/// all of a table's columns.
///
/// # Example
///
/// ```
Expand Down Expand Up @@ -299,12 +302,67 @@ impl<S: AsyncRead + AsyncWrite + Unpin + Send> Client<S> {
pub async fn bulk_insert<'a>(
&'a mut self,
table: &'a str,
) -> crate::Result<BulkLoadRequest<'a, S>> {
self.bulk_insert_columns(table, &["*"]).await
}

/// Execute a `BULK INSERT` statement, efficiently storing a large number of
/// rows to a specified table. Note: make sure the input row follows the same
/// schema as the column list, otherwise calling `send()` will return an error.
///
/// # Example
///
/// ```
/// # use tiberius::{Config, IntoRow};
/// # use tokio_util::compat::TokioAsyncWriteCompatExt;
/// # use std::env;
/// # #[tokio::main]
/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
/// # let c_str = env::var("TIBERIUS_TEST_CONNECTION_STRING").unwrap_or(
/// # "server=tcp:localhost,1433;integratedSecurity=true;TrustServerCertificate=true".to_owned(),
/// # );
/// # let config = Config::from_ado_string(&c_str)?;
/// # let tcp = tokio::net::TcpStream::connect(config.get_addr()).await?;
/// # tcp.set_nodelay(true)?;
/// # let mut client = tiberius::Client::connect(config, tcp.compat_write()).await?;
/// let create_table = r#"
/// CREATE TABLE ##bulk_test_columns (
/// id INT IDENTITY PRIMARY KEY,
/// foo INT NOT NULL,
/// bar FLOAT NOT NULL
/// )
/// "#;
///
/// client.simple_query(create_table).await?;
///
/// // Start the bulk insert with the client.
/// let mut req = client.bulk_insert_columns("##bulk_test_columns", &["foo", "bar"]).await?;
///
/// for (i, j) in [(0i32, 0f64), (1i32, 1f64), (2i32, 2f64)] {
/// let row = (i, j).into_row();
///
/// // The request will handle flushing to the wire in an optimal way,
/// // balancing between memory usage and IO performance.
/// req.send(row).await?;
/// }
///
/// // The request must be finalized.
/// let res = req.finalize().await?;
/// assert_eq!(3, res.total());
/// # Ok(())
/// # }
/// ```
pub async fn bulk_insert_columns<'a>(
&'a mut self,
table: &'a str,
columns: &'a [&'a str],
) -> crate::Result<BulkLoadRequest<'a, S>> {
// Start the bulk request
self.connection.flush_stream().await?;

// retrieve column metadata from server
let query = format!("SELECT TOP 0 * FROM {}", table);
let columns = columns.join(", ");
let query = format!("SELECT TOP 0 {columns} FROM {table}");

let req = BatchRequest::new(query, self.connection.context().transaction_descriptor());

Expand Down
35 changes: 27 additions & 8 deletions src/client/auth.rs
Original file line number Diff line number Diff line change
@@ -1,18 +1,15 @@
use std::fmt::Debug;
use zeroize::Zeroizing;

#[derive(Clone, PartialEq, Eq)]
pub struct SqlServerAuth {
user: String,
password: String,
password: Zeroizing<String>,
}

impl SqlServerAuth {
pub(crate) fn user(&self) -> &str {
&self.user
}

pub(crate) fn password(&self) -> &str {
&self.password
pub(crate) fn into_credentials(self) -> (String, Zeroizing<String>) {
(self.user, self.password)
}
}

Expand Down Expand Up @@ -79,7 +76,7 @@ impl AuthMethod {
pub fn sql_server(user: impl ToString, password: impl ToString) -> Self {
Self::SqlServer(SqlServerAuth {
user: user.to_string(),
password: password.to_string(),
password: Zeroizing::new(password.to_string()),
})
}

Expand All @@ -104,3 +101,25 @@ impl AuthMethod {
Self::AADToken(token.to_string())
}
}

#[cfg(test)]
mod tests {
use super::AuthMethod;
use zeroize::Zeroize;

#[test]
fn sql_server_password_can_be_consumed_and_zeroized() {
let AuthMethod::SqlServer(auth) = AuthMethod::sql_server("sa", "secret") else {
unreachable!();
};

let (user, mut password) = auth.into_credentials();

assert_eq!("sa", user);
assert_eq!("secret", password.as_str());

password.zeroize();

assert!(password.is_empty());
}
}
18 changes: 18 additions & 0 deletions src/client/config.rs
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ pub struct Config {
pub(crate) trust: TrustConfig,
pub(crate) auth: AuthMethod,
pub(crate) readonly: bool,
pub(crate) packet_size: Option<u32>,
}

#[derive(Clone, Debug)]
Expand Down Expand Up @@ -65,6 +66,7 @@ impl Default for Config {
trust: TrustConfig::Default,
auth: AuthMethod::None,
readonly: false,
packet_size: None,
}
}
}
Expand Down Expand Up @@ -115,6 +117,22 @@ impl Config {
self.application_name = Some(name.to_string());
}

/// Sets the TDS packet size for the connection.
///
/// Larger packet sizes can improve bulk insert performance by reducing
/// the number of network round-trips. Valid values are 512 to 32767.
/// The server may negotiate a different size.
///
/// - Defaults to 4096 bytes.
pub fn packet_size(&mut self, size: u32) {
self.packet_size = Some(size);
}

/// Gets the configured packet size, if set.
pub fn get_packet_size(&self) -> Option<u32> {
self.packet_size
}

/// Set the preferred encryption level.
///
/// - With `tls` feature, defaults to `Required`.
Expand Down
61 changes: 56 additions & 5 deletions src/client/connection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ use asynchronous_codec::Framed;
use bytes::BytesMut;
#[cfg(any(windows, feature = "integrated-auth-gssapi"))]
use codec::TokenSspi;
use futures_util::io::{AsyncRead, AsyncWrite};
use futures_util::io::{AsyncRead, AsyncWrite, AsyncWriteExt};
use futures_util::ready;
use futures_util::sink::SinkExt;
use futures_util::stream::{Stream, TryStream, TryStreamExt};
Expand All @@ -39,6 +39,7 @@ use task::Poll;
use tracing::{event, Level};
#[cfg(all(windows, feature = "winauth"))]
use winauth::{windows::NtlmSspiBuilder, NextBytes};
use zeroize::{Zeroize, Zeroizing};

/// A `Connection` is an abstraction between the [`Client`] and the server. It
/// can be used as a `Stream` to fetch [`Packet`]s from and to `send` packets
Expand Down Expand Up @@ -106,6 +107,7 @@ impl<S: AsyncRead + AsyncWrite + Unpin + Send> Connection<S> {
config.host,
config.application_name,
config.readonly,
config.packet_size,
prelogin,
)
.await?;
Expand Down Expand Up @@ -196,6 +198,45 @@ impl<S: AsyncRead + AsyncWrite + Unpin + Send> Connection<S> {
Ok(())
}

async fn send_sensitive_login(
&mut self,
mut header: PacketHeader,
mut payload: Zeroizing<Vec<u8>>,
) -> crate::Result<()> {
self.flushed = false;
let packet_size = (self.context.packet_size() as usize) - HEADER_BYTES;
let mut offset = 0;

while offset < payload.len() {
let end = cmp::min(payload.len(), offset + packet_size);

if end == payload.len() {
header.set_status(PacketStatus::EndOfMessage);
} else {
header.set_status(PacketStatus::NormalMessage);
}

let mut frame = Zeroizing::new(Vec::with_capacity(HEADER_BYTES + end - offset));
header.encode(&mut *frame)?;
frame.extend_from_slice(&payload[offset..end]);

let size = (frame.len() as u16).to_be_bytes();
frame[2] = size[0];
frame[3] = size[1];

event!(Level::TRACE, "Sending a packet ({} bytes)", frame.len(),);

self.transport.write_all(frame.as_slice()).await?;
frame.zeroize();
payload[offset..end].zeroize();
offset = end;
}

self.transport.flush().await?;

Ok(())
}

/// Sends a packet of data to the database.
///
/// # Warning
Expand Down Expand Up @@ -293,6 +334,7 @@ impl<S: AsyncRead + AsyncWrite + Unpin + Send> Connection<S> {
server_name: Option<String>,
application_name: Option<String>,
readonly: bool,
packet_size: Option<u32>,
prelogin: PreloginMessage,
) -> crate::Result<Self> {
let mut login_message = LoginMessage::new();
Expand All @@ -311,6 +353,10 @@ impl<S: AsyncRead + AsyncWrite + Unpin + Send> Connection<S> {

login_message.readonly(readonly);

if let Some(size) = packet_size {
login_message.packet_size(size);
}

match auth {
#[cfg(all(windows, feature = "winauth"))]
AuthMethod::Integrated => {
Expand All @@ -332,7 +378,7 @@ impl<S: AsyncRead + AsyncWrite + Unpin + Send> Connection<S> {
event!(Level::TRACE, sspi_response_len = sspi_response.len());

let id = self.context.next_packet_id();
let header = PacketHeader::login(id);
let header = PacketHeader::sspi(id);

let token = TokenSspi::new(sspi_response);
self.send(header, token).await?;
Expand Down Expand Up @@ -415,11 +461,16 @@ impl<S: AsyncRead + AsyncWrite + Unpin + Send> Connection<S> {
self = self.post_login_encryption(encryption);
}
AuthMethod::SqlServer(auth) => {
login_message.user_name(auth.user());
login_message.password(auth.password());
let (user, mut password) = auth.into_credentials();

login_message.user_name(user);
login_message.password(password.as_str());
let payload = login_message.encode_to_vec()?;
password.zeroize();

let id = self.context.next_packet_id();
self.send(PacketHeader::login(id), login_message).await?;
self.send_sensitive_login(PacketHeader::login(id), payload)
.await?;
self = self.post_login_encryption(encryption);
}
AuthMethod::AADToken(token) => {
Expand Down
Loading
Loading