pub mod actions;
use crate::llm::action_helper::call_llm;
use crate::llm::actions::protocol_trait::ActionResult;
use crate::llm::ollama_client::OllamaClient;
use crate::logging::emit::Log;
use crate::protocol::Event;
use crate::server::connection::ConnectionId;
use crate::state::app_state::AppState;
use actions::{MssqlProtocol, MSSQL_QUERY_EVENT};
use anyhow::Result;
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::io::{AsyncReadExt, AsyncWriteExt, ReadHalf};
use tokio::net::{TcpListener, TcpStream};
use tokio::sync::{mpsc, Mutex};
use tracing::{debug, error, info, trace, warn};
const TDS_PACKET_SIZE: usize = 4096;
type SharedWrite = Arc<Mutex<tokio::io::WriteHalf<TcpStream>>>;
const FIRST_PACKET_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
const IDLE_BETWEEN_PACKETS_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(600);
const MAX_CONNECTIONS: usize = crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS;
const CONNECTION_CAP_REFUSAL: &[u8] = b"";
const LOGIN_FAILED_ERROR: u32 = 18456;
const MSSQL_ERROR_GENERIC: u32 = 50000;
const MSSQL_ERROR_NOT_ENOUGH_RESOURCES: u32 = 49918;
pub struct MssqlServer {
llm_client: OllamaClient,
app_state: Arc<AppState>,
_status_tx: mpsc::UnboundedSender<String>,
server_id: Option<crate::state::ServerId>,
}
impl MssqlServer {
pub fn new(
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
server_id: Option<crate::state::ServerId>,
) -> Self {
Self {
llm_client,
app_state,
_status_tx: status_tx,
server_id,
}
}
pub async fn spawn_with_llm_actions(
listen_addr: SocketAddr,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
_send_first: bool,
server_id: crate::state::ServerId,
) -> Result<SocketAddr> {
let listener = TcpListener::bind(listen_addr).await?;
let actual_addr = listener.local_addr()?;
Log::new(Some(&status_tx)).info(format!("MSSQL server listening on {}", actual_addr));
let server = Arc::new(MssqlServer::new(
llm_client,
app_state.clone(),
status_tx.clone(),
Some(server_id),
));
let status_tx_clone = status_tx.clone();
let task_registrar = app_state.clone();
let limiter = crate::server::accept_bounded::ConnectionLimiter::new(MAX_CONNECTIONS);
let accept_handle = tokio::spawn(async move {
loop {
match crate::server::accept_bounded::accept_bounded(
&listener,
&limiter,
CONNECTION_CAP_REFUSAL,
"MSSQL",
Some(&status_tx),
)
.await
{
Ok((stream, addr, permit)) => {
Log::new(Some(&status_tx)).info(format!("MSSQL connection from {}", addr));
let connection_id =
ConnectionId::new(app_state.get_next_unified_id().await);
let local_addr_conn = stream.local_addr().unwrap_or(actual_addr);
if let Some(server_id) = server.server_id {
use crate::state::server::{
ConnectionState as ServerConnectionState, ConnectionStatus,
ProtocolConnectionInfo,
};
let now = crate::utils::clock::Instant::now();
let conn_state = ServerConnectionState {
id: connection_id,
remote_addr: addr,
local_addr: local_addr_conn,
bytes_sent: 0,
bytes_received: 0,
packets_sent: 0,
packets_received: 0,
last_activity: now,
status: ConnectionStatus::Active,
status_changed_at: now,
protocol_info: ProtocolConnectionInfo::empty(),
};
server
.app_state
.add_connection_to_server(server_id, conn_state)
.await;
}
let handler = MssqlHandler::new(
connection_id,
server.llm_client.clone(),
server.app_state.clone(),
status_tx.clone(),
server.server_id,
addr,
);
let task_owner = app_state.clone();
task_owner
.spawn_server_task(server_id, async move {
let _permit = permit;
if let Err(e) = handler.handle_connection(stream).await {
error!("MSSQL connection error: {:?}", e);
}
})
.await;
}
Err(e) => {
Log::new(Some(&status_tx)).error(format!("MSSQL accept error: {}", e));
}
}
}
});
task_registrar
.register_server_task(server_id, accept_handle)
.await;
let _ = status_tx_clone.send("__UPDATE_UI__".to_string());
Ok(actual_addr)
}
}
pub struct MssqlHandler {
connection_id: ConnectionId,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
#[allow(dead_code)]
server_id: Option<crate::state::ServerId>,
#[allow(dead_code)]
remote_addr: SocketAddr,
protocol: Arc<MssqlProtocol>,
}
impl MssqlHandler {
pub fn new(
connection_id: ConnectionId,
llm_client: OllamaClient,
app_state: Arc<AppState>,
status_tx: mpsc::UnboundedSender<String>,
server_id: Option<crate::state::ServerId>,
remote_addr: SocketAddr,
) -> Self {
let protocol = Arc::new(MssqlProtocol::new(
connection_id,
app_state.clone(),
status_tx.clone(),
));
Self {
connection_id,
llm_client,
app_state,
status_tx,
server_id,
remote_addr,
protocol,
}
}
async fn handle_connection(self, stream: TcpStream) -> Result<()> {
let server_id = self.server_id;
let connection_id = self.connection_id;
let app_state = self.app_state.clone();
let (mut reader, write_half) = tokio::io::split(stream);
let write_half: SharedWrite = Arc::new(Mutex::new(write_half));
if let Some(server_id) = server_id {
let peer_rx = crate::server::peer_support::register_peer_channel(
&app_state,
server_id,
connection_id.as_u32(),
)
.await;
crate::server::peer_support::spawn_peer_command_task(
peer_rx,
self.protocol.clone(),
app_state.clone(),
server_id,
connection_id.as_u32(),
write_half.clone(),
self.status_tx.clone(),
);
}
let result = self.run_session(&mut reader, &write_half).await;
if let Some(server_id) = server_id {
app_state
.remove_peer_handle(server_id, connection_id.as_u32())
.await;
}
let _ = write_half.lock().await.shutdown().await;
if let Some(server_id) = server_id {
app_state
.close_connection_on_server(server_id, connection_id)
.await;
}
result
}
async fn run_session(
&self,
reader: &mut ReadHalf<TcpStream>,
write_half: &SharedWrite,
) -> Result<()> {
info!("MSSQL connection established");
let mut authenticated = false;
loop {
let read_timeout = if authenticated {
IDLE_BETWEEN_PACKETS_TIMEOUT
} else {
FIRST_PACKET_READ_TIMEOUT
};
let header = match self.read_tds_header(reader, read_timeout).await {
Ok(h) => h,
Err(e) => {
debug!("Error reading TDS header: {}", e);
break;
}
};
if header.length < 8 {
debug!("Invalid TDS packet length: {}", header.length);
break;
}
let data_len = header.length - 8;
let mut data = vec![0u8; data_len as usize];
match tokio::time::timeout(read_timeout, reader.read_exact(&mut data)).await {
Ok(read) => read?,
Err(_) => {
debug!(
"MSSQL connection {} announced a {}-byte packet and stalled for {}s \
without delivering it; closing",
self.connection_id,
header.length,
read_timeout.as_secs()
);
break;
}
};
if let Some(server_id) = self.server_id {
self.app_state
.update_connection_stats(
server_id,
self.connection_id,
Some(u64::from(header.length)),
None,
Some(1),
None,
)
.await;
}
trace!(
"TDS packet type: 0x{:02x}, length: {}",
header.packet_type,
header.length
);
match header.packet_type {
0x12 => {
debug!("Received Pre-Login packet (length: {})", header.length);
self.send_prelogin_response(write_half).await?;
debug!("Pre-Login response sent");
}
0x10 => {
debug!("Received Login packet");
if !self.handle_login(write_half, &data).await? {
break;
}
authenticated = true;
}
0x01 | 0x03 if !authenticated => {
Log::new(Some(&self.status_tx)).warn(
"MSSQL statement refused (decision=fail_closed_not_logged_in): the \
client sent a batch or RPC before any LOGIN7"
.to_string(),
);
self.send_error(
write_half,
LOGIN_FAILED_ERROR,
"Login failed. The login is from an untrusted domain and cannot be \
used with Windows authentication.",
14,
)
.await?;
break;
}
0x01 => {
debug!("Received SQL Batch packet");
let query = self.parse_sql_batch(&data)?;
debug!("SQL Query: {}", query);
if self.handle_query(write_half, &query).await? {
break;
}
}
0x03 => {
debug!("Received RPC Request");
let query = self.parse_rpc_request(&data)?;
if !query.is_empty() {
debug!("RPC Query: {}", query);
if self.handle_query(write_half, &query).await? {
break;
}
} else {
debug!("RPC call without extractable query (ignoring)");
self.send_empty_result(write_half).await?;
}
}
0x0E => {
debug!("Received Bulk Load (not implemented)");
self.send_error(write_half, 40002, "Bulk load not supported", 16)
.await?;
}
0x07 => {
debug!("Received Attention signal");
break;
}
_ => {
debug!("Unknown TDS packet type: 0x{:02x}", header.packet_type);
self.send_error(write_half, 40002, "Unknown packet type", 16)
.await?;
}
}
}
Ok(())
}
async fn read_tds_header(
&self,
reader: &mut ReadHalf<TcpStream>,
read_timeout: std::time::Duration,
) -> Result<TdsHeader> {
let mut header_bytes = [0u8; 8];
match tokio::time::timeout(read_timeout, reader.read_exact(&mut header_bytes)).await {
Ok(read) => read?,
Err(_) => {
anyhow::bail!(
"peer sent nothing for {}s; closing idle connection",
read_timeout.as_secs()
)
}
};
Ok(TdsHeader {
packet_type: header_bytes[0],
status: header_bytes[1],
length: u16::from_be_bytes([header_bytes[2], header_bytes[3]]),
spid: u16::from_be_bytes([header_bytes[4], header_bytes[5]]),
packet_id: header_bytes[6],
window: header_bytes[7],
})
}
async fn send_prelogin_response(&self, write_half: &SharedWrite) -> Result<()> {
let mut response = Vec::new();
let header_size = 16u16;
let version_offset = header_size; let version_length = 6u16;
let encryption_offset = version_offset + version_length; let encryption_length = 1u16;
let threadid_offset = encryption_offset + encryption_length; let threadid_length = 4u16;
response.push(0x00);
response.extend_from_slice(&version_offset.to_be_bytes()); response.extend_from_slice(&version_length.to_be_bytes());
response.push(0x01);
response.extend_from_slice(&encryption_offset.to_be_bytes()); response.extend_from_slice(&encryption_length.to_be_bytes());
response.push(0x03);
response.extend_from_slice(&threadid_offset.to_be_bytes()); response.extend_from_slice(&threadid_length.to_be_bytes());
response.push(0xFF);
response.extend_from_slice(&[0x10, 0x00, 0x00, 0x00, 0x00, 0x00]);
response.push(0x02);
response.extend_from_slice(&[0x00, 0x00, 0x00, 0x00]);
self.send_tds_packet(write_half, 0x04, &response).await
}
async fn handle_login(&self, write_half: &SharedWrite, data: &[u8]) -> Result<bool> {
let (username, database, app_name) = parse_login7(data);
let event = Event::new(
&actions::MSSQL_LOGIN_EVENT,
serde_json::json!({
"username": username,
"database": database,
"app_name": app_name,
}),
);
let server_id = self
.server_id
.unwrap_or_else(|| crate::state::ServerId::new(0));
let llm_result = call_llm(
&self.llm_client,
&self.app_state,
server_id,
Some(self.connection_id),
&event,
self.protocol.as_ref(),
)
.await;
match llm_result {
Ok(result) => {
for r in &result.protocol_results {
if let ActionResult::Custom { name, data } = r {
match name.as_str() {
"mssql_login_ack" => {
let db = data
.get("database")
.and_then(|v| v.as_str())
.unwrap_or("master");
Log::new(Some(&self.status_tx)).info(format!(
"MSSQL login accepted for user '{}' (decision=model_accept)",
username
));
self.send_login_response_with_database(write_half, db)
.await?;
return Ok(true);
}
"mssql_error" => {
let number = data
.get("error_number")
.and_then(|v| v.as_u64())
.unwrap_or(LOGIN_FAILED_ERROR as u64)
as u32;
let message = data
.get("message")
.and_then(|v| v.as_str())
.unwrap_or("Login failed.");
Log::new(Some(&self.status_tx)).info(format!(
"MSSQL login refused for user '{}' (decision=model_reject): {}",
username, message
));
self.send_error(write_half, number, message, 14).await?;
return Ok(false);
}
_ => {}
}
}
}
Log::new(Some(&self.status_tx)).warn(format!(
"MSSQL login refused for user '{}' (decision=fail_closed_no_action): the \
handler produced no mssql_login_ack",
username
));
self.send_error(
write_half,
LOGIN_FAILED_ERROR,
&format!("Login failed for user '{}'.", username),
14,
)
.await?;
Ok(false)
}
Err(e) => {
let failure = crate::utils::WireFailure::classify(&e);
error!(
"MSSQL login for user '{}' decision=fail_closed_llm_error category={:?}: {:#}",
username, failure, e
);
Log::new(Some(&self.status_tx)).warn(format!(
"MSSQL login refused for user '{}' (decision=fail_closed_llm_error)",
username
));
self.send_error(write_half, LOGIN_FAILED_ERROR, failure.text(), 14)
.await?;
Ok(false)
}
}
}
async fn send_login_response_with_database(
&self,
write_half: &SharedWrite,
db_name: &str,
) -> Result<()> {
let mut response = Vec::new();
let db_name_utf16: Vec<u8> = db_name
.encode_utf16()
.flat_map(|c| c.to_le_bytes())
.collect();
let mut envchange_db = Vec::new();
envchange_db.push(0xE3); let envchange_db_len = 1 + 1 + db_name_utf16.len() + 1 + db_name_utf16.len();
envchange_db.extend_from_slice(&(envchange_db_len as u16).to_le_bytes());
envchange_db.push(0x01); envchange_db.push(db_name.len() as u8); envchange_db.extend_from_slice(&db_name_utf16);
envchange_db.push(db_name.len() as u8); envchange_db.extend_from_slice(&db_name_utf16);
response.extend_from_slice(&envchange_db);
let lang = "us_english";
let lang_utf16: Vec<u8> = lang.encode_utf16().flat_map(|c| c.to_le_bytes()).collect();
let mut envchange_lang = Vec::new();
envchange_lang.push(0xE3); let envchange_lang_len = 1 + 1 + lang_utf16.len() + 1 + 0; envchange_lang.extend_from_slice(&(envchange_lang_len as u16).to_le_bytes());
envchange_lang.push(0x02); envchange_lang.push(lang.len() as u8); envchange_lang.extend_from_slice(&lang_utf16);
envchange_lang.push(0x00); response.extend_from_slice(&envchange_lang);
let pkt_size_new = "4096";
let pkt_size_old = "512"; let pkt_size_new_utf16: Vec<u8> = pkt_size_new
.encode_utf16()
.flat_map(|c| c.to_le_bytes())
.collect();
let pkt_size_old_utf16: Vec<u8> = pkt_size_old
.encode_utf16()
.flat_map(|c| c.to_le_bytes())
.collect();
let mut envchange_pkt = Vec::new();
envchange_pkt.push(0xE3); let envchange_pkt_len = 1 + 1 + pkt_size_new_utf16.len() + 1 + pkt_size_old_utf16.len();
envchange_pkt.extend_from_slice(&(envchange_pkt_len as u16).to_le_bytes());
envchange_pkt.push(0x04); envchange_pkt.push(pkt_size_new.len() as u8); envchange_pkt.extend_from_slice(&pkt_size_new_utf16);
envchange_pkt.push(pkt_size_old.len() as u8); envchange_pkt.extend_from_slice(&pkt_size_old_utf16);
response.extend_from_slice(&envchange_pkt);
let msg = "Login succeeded";
let msg_utf16: Vec<u8> = msg.encode_utf16().flat_map(|c| c.to_le_bytes()).collect();
response.push(0xAB); let info_len = 4 + 1 + 1 + 2 + msg_utf16.len() + 1 + 1 + 4;
response.extend_from_slice(&(info_len as u16).to_le_bytes());
response.extend_from_slice(&[0x01, 0x00, 0x00, 0x00]); response.push(0x01); response.push(0x00); response.extend_from_slice(&(msg.len() as u16).to_le_bytes()); response.extend_from_slice(&msg_utf16);
response.push(0x00); response.push(0x00); response.extend_from_slice(&[0x00, 0x00, 0x00, 0x00]);
response.push(0xFD); response.extend_from_slice(&[0x00, 0x00]); response.extend_from_slice(&[0x00, 0x00]); response.extend_from_slice(&[0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]);
self.send_tds_packet(write_half, 0x04, &response).await
}
fn parse_sql_batch(&self, data: &[u8]) -> Result<String> {
if data.len() < 22 {
return Ok(String::new());
}
let sql_bytes = &data[22..];
let sql_u16: Vec<u16> = sql_bytes
.chunks_exact(2)
.map(|chunk| u16::from_le_bytes([chunk[0], chunk[1]]))
.collect();
Ok(String::from_utf16_lossy(&sql_u16).trim().to_string())
}
fn parse_rpc_request(&self, data: &[u8]) -> Result<String> {
const SQL_KEYWORDS: [&str; 7] = [
"SELECT", "INSERT", "UPDATE", "DELETE", "CREATE", "DROP", "ALTER",
];
if data.len() < 10 {
return Ok(String::new());
}
let preview_len = std::cmp::min(data.len(), 200);
let hex_preview: String = data[..preview_len]
.iter()
.map(|b| format!("{:02x}", b))
.collect::<Vec<_>>()
.join(" ");
debug!(
"RPC data preview (first {} bytes): {}",
preview_len, hex_preview
);
let units: Vec<u16> = data
.chunks_exact(2)
.map(|c| u16::from_le_bytes([c[0], c[1]]))
.collect();
let text = String::from_utf16_lossy(&units);
let upper = text.to_ascii_uppercase();
let Some(pos) = SQL_KEYWORDS.iter().filter_map(|kw| upper.find(kw)).min() else {
return Ok(String::new());
};
let sql: String = text[pos..]
.chars()
.take_while(|c| {
c.is_ascii() && (*c == '\n' || *c == '\r' || *c == '\t' || !c.is_ascii_control())
})
.collect();
let sql = sql.trim().to_string();
if sql.is_empty() {
return Ok(String::new());
}
debug!("Extracted SQL from RPC at offset {}: {}", pos, sql);
Ok(sql)
}
async fn handle_query(&self, write_half: &SharedWrite, query: &str) -> Result<bool> {
trace!("Calling LLM for MSSQL query: {}", query);
let event = Event::new(
&MSSQL_QUERY_EVENT,
serde_json::json!({
"query": query,
}),
);
let server_id = self
.server_id
.unwrap_or_else(|| crate::state::ServerId::new(0));
let llm_result = call_llm(
&self.llm_client,
&self.app_state,
server_id,
Some(self.connection_id),
&event,
self.protocol.as_ref(),
)
.await;
match llm_result {
Ok(execution_result) => {
let close_requested = execution_result
.protocol_results
.iter()
.any(|r| matches!(r, ActionResult::CloseConnection));
let mut responded = false;
for result in execution_result.protocol_results {
if responded {
if let ActionResult::Custom { name, .. } = &result {
warn!(
"MSSQL: query already answered, ignoring extra action result '{}'",
name
);
}
continue;
}
match result {
ActionResult::Custom { name, data } => match name.as_str() {
"mssql_query_response" => {
let columns = data
.get("columns")
.and_then(|v| v.as_array())
.cloned()
.unwrap_or_default();
let rows = data
.get("rows")
.and_then(|v| v.as_array())
.cloned()
.unwrap_or_default();
self.send_result_set(write_half, columns, rows).await?;
responded = true;
}
"mssql_error" => {
let error_number = data
.get("error_number")
.and_then(|v| v.as_u64())
.unwrap_or(50000)
as u32;
let message = data
.get("message")
.and_then(|v| v.as_str())
.unwrap_or("Unknown error");
let severity =
data.get("severity").and_then(|v| v.as_u64()).unwrap_or(16)
as u8;
Log::new(Some(&self.status_tx)).info(format!(
"MSSQL query decision=model_reject error={} severity={}",
error_number, severity
));
self.send_error(write_half, error_number, message, severity)
.await?;
responded = true;
}
"mssql_ok" => {
let rows_affected = data
.get("rows_affected")
.and_then(|v| v.as_u64())
.unwrap_or(0);
self.send_done(write_half, rows_affected).await?;
responded = true;
}
other => {
warn!(
"MSSQL: no wire encoding for action result '{}', ignoring",
other
);
}
},
_ => {}
}
}
if !responded {
if close_requested {
Log::new(Some(&self.status_tx))
.info("MSSQL query decision=model_close (closing without a result)");
self.send_done(write_half, 0).await?;
} else {
let message = crate::utils::WireFailure::Unavailable.prefixed_text();
Log::new(Some(&self.status_tx)).warn(format!(
"MSSQL query decision=fail_closed_no_action; replying error {}: {}",
MSSQL_ERROR_GENERIC, message
));
warn!("MSSQL: no response action produced for query {:?}", query);
self.send_error(write_half, MSSQL_ERROR_GENERIC, message, 16)
.await?;
}
}
Ok(close_requested)
}
Err(e) => {
let overloaded = crate::llm::is_overload_error(&e);
let number = if overloaded {
MSSQL_ERROR_NOT_ENOUGH_RESOURCES
} else {
MSSQL_ERROR_GENERIC
};
let message = crate::utils::WireFailure::classify(&e).prefixed_text();
error!("MSSQL: LLM error answering query {:?}: {:#}", query, e);
Log::new(Some(&self.status_tx)).warn(format!(
"MSSQL query decision=fail_closed_llm_error ({}); replying error {number}: {message}",
if overloaded { "overloaded" } else { "unavailable" }
));
self.send_error(write_half, number, message, 16).await?;
Ok(false)
}
}
}
async fn send_result_set(
&self,
write_half: &SharedWrite,
columns: Vec<serde_json::Value>,
rows: Vec<serde_json::Value>,
) -> Result<()> {
debug!(
"send_result_set called: {} columns, {} rows",
columns.len(),
rows.len()
);
let col_types: Vec<TdsColumnType> = columns
.iter()
.map(|col| {
TdsColumnType::from_name(
col.get("type")
.and_then(|v| v.as_str())
.unwrap_or("NVARCHAR"),
)
})
.collect();
let mut response = Vec::new();
response.push(0x81);
response.extend_from_slice(&(columns.len() as u16).to_le_bytes());
for (col, col_type) in columns.iter().zip(&col_types) {
let col_name = col.get("name").and_then(|v| v.as_str()).unwrap_or("column");
response.extend_from_slice(&[0x00, 0x00, 0x00, 0x00]); response.extend_from_slice(&[0x01, 0x00]); col_type.write_type_info(&mut response);
let name_units: Vec<u16> = col_name.encode_utf16().take(128).collect();
response.push(name_units.len() as u8);
for unit in &name_units {
response.extend_from_slice(&unit.to_le_bytes());
}
}
for row in &rows {
response.push(0xD1);
let row_values = row.as_array().cloned().unwrap_or_default();
for (idx, col_type) in col_types.iter().enumerate() {
let value = row_values.get(idx).unwrap_or(&serde_json::Value::Null);
col_type.write_value(&mut response, value);
}
}
response.push(0xFD);
response.extend_from_slice(&0x0010u16.to_le_bytes()); response.extend_from_slice(&0x00C1u16.to_le_bytes()); response.extend_from_slice(&(rows.len() as u64).to_le_bytes());
debug!("Sending result set: {} bytes", response.len());
self.send_tds_packet(write_half, 0x04, &response).await
}
async fn send_empty_result(&self, write_half: &SharedWrite) -> Result<()> {
let mut response = Vec::new();
response.push(0xFD);
response.extend_from_slice(&[0x00, 0x00]); response.extend_from_slice(&[0x00, 0x00]); response.extend_from_slice(&[0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]);
self.send_tds_packet(write_half, 0x04, &response).await
}
async fn send_error(
&self,
write_half: &SharedWrite,
error_number: u32,
message: &str,
severity: u8,
) -> Result<()> {
let mut response = Vec::new();
response.push(0xAA);
let msg_u16: Vec<u16> = message.encode_utf16().collect();
let msg_bytes: Vec<u8> = msg_u16.iter().flat_map(|c| c.to_le_bytes()).collect();
let token_len = 4 + 1 + 1 + 2 + msg_bytes.len() + 1 + 1 + 4;
response.extend_from_slice(&(token_len as u16).to_le_bytes());
response.extend_from_slice(&error_number.to_le_bytes());
response.push(0x01); response.push(severity); response.extend_from_slice(&(msg_u16.len() as u16).to_le_bytes());
response.extend_from_slice(&msg_bytes);
response.push(0x00); response.push(0x00); response.extend_from_slice(&[0x00, 0x00, 0x00, 0x00]);
response.push(0xFD);
response.extend_from_slice(&[0x00, 0x00]); response.extend_from_slice(&[0x00, 0x00]); response.extend_from_slice(&[0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00]);
self.send_tds_packet(write_half, 0x04, &response).await
}
async fn send_done(&self, write_half: &SharedWrite, rows_affected: u64) -> Result<()> {
let mut response = Vec::new();
response.push(0xFD); response.extend_from_slice(&0x0010u16.to_le_bytes()); response.extend_from_slice(&0x00C1u16.to_le_bytes()); response.extend_from_slice(&rows_affected.to_le_bytes());
self.send_tds_packet(write_half, 0x04, &response).await
}
async fn send_tds_packet(
&self,
write_half: &SharedWrite,
packet_type: u8,
data: &[u8],
) -> Result<()> {
let max_payload = TDS_PACKET_SIZE - 8;
let mut packet_id: u8 = 1;
let mut offset = 0;
let mut framed: Vec<u8> = Vec::with_capacity(data.len() + 8);
let mut packets: u64 = 0;
loop {
let end = std::cmp::min(offset + max_payload, data.len());
let chunk = &data[offset..end];
let is_last = end == data.len();
framed.push(packet_type); framed.push(if is_last { 0x01 } else { 0x00 }); framed.extend_from_slice(&((8 + chunk.len()) as u16).to_be_bytes()); framed.extend_from_slice(&[0x00, 0x00]); framed.push(packet_id); framed.push(0x00); framed.extend_from_slice(chunk);
packets += 1;
if is_last {
break;
}
offset = end;
packet_id = packet_id.wrapping_add(1).max(1);
}
write_counted(
write_half,
&framed,
packets,
&self.app_state,
self.server_id,
self.connection_id,
)
.await?;
Ok(())
}
}
async fn write_counted(
write_half: &SharedWrite,
data: &[u8],
packets: u64,
app_state: &AppState,
server_id: Option<crate::state::ServerId>,
connection_id: ConnectionId,
) -> std::io::Result<()> {
{
let mut writer = write_half.lock().await;
writer.write_all(data).await?;
writer.flush().await?;
}
if let Some(server_id) = server_id {
app_state
.update_connection_stats(
server_id,
connection_id,
None,
Some(data.len() as u64),
None,
Some(packets),
)
.await;
}
Ok(())
}
#[allow(dead_code)]
struct TdsHeader {
packet_type: u8,
#[allow(dead_code)]
status: u8,
length: u16,
#[allow(dead_code)]
spid: u16,
#[allow(dead_code)]
packet_id: u8,
#[allow(dead_code)]
window: u8,
}
const NVARCHAR_MAX_BYTES: usize = 8000;
#[derive(Clone, Copy, PartialEq, Eq)]
enum TdsColumnType {
Int(u8),
Bit,
Float,
NVarChar,
}
impl TdsColumnType {
fn from_name(type_name: &str) -> Self {
match type_name.trim().to_uppercase().as_str() {
"TINYINT" => TdsColumnType::Int(1),
"SMALLINT" => TdsColumnType::Int(2),
"INT" | "INTEGER" => TdsColumnType::Int(4),
"BIGINT" => TdsColumnType::Int(8),
"BIT" | "BOOL" | "BOOLEAN" => TdsColumnType::Bit,
"FLOAT" | "REAL" | "DOUBLE" | "DECIMAL" | "NUMERIC" | "MONEY" => TdsColumnType::Float,
_ => TdsColumnType::NVarChar,
}
}
fn write_type_info(&self, out: &mut Vec<u8>) {
match self {
TdsColumnType::Int(width) => {
out.push(0x26); out.push(*width);
}
TdsColumnType::Bit => {
out.push(0x68); out.push(0x01);
}
TdsColumnType::Float => {
out.push(0x6D); out.push(0x08);
}
TdsColumnType::NVarChar => {
out.push(0xE7); out.extend_from_slice(&(NVARCHAR_MAX_BYTES as u16).to_le_bytes());
out.extend_from_slice(&[0x00, 0x00, 0x00, 0x00, 0x00]); }
}
}
fn write_value(&self, out: &mut Vec<u8>, value: &serde_json::Value) {
match self {
TdsColumnType::Int(width) => {
match value_as_i64(value) {
Some(v) => {
out.push(*width);
match width {
1 => out.push(v as u8),
2 => out.extend_from_slice(&(v as i16).to_le_bytes()),
4 => out.extend_from_slice(&(v as i32).to_le_bytes()),
_ => out.extend_from_slice(&v.to_le_bytes()),
}
}
None => out.push(0x00),
}
}
TdsColumnType::Bit => match value_as_bool(value) {
Some(b) => {
out.push(0x01);
out.push(u8::from(b));
}
None => out.push(0x00),
},
TdsColumnType::Float => match value_as_f64(value) {
Some(f) => {
out.push(0x08);
out.extend_from_slice(&f.to_le_bytes());
}
None => out.push(0x00),
},
TdsColumnType::NVarChar => {
if value.is_null() {
out.extend_from_slice(&0xFFFFu16.to_le_bytes());
return;
}
let text = json_to_string(value);
let mut bytes: Vec<u8> = Vec::new();
for unit in text.encode_utf16() {
if bytes.len() + 2 > NVARCHAR_MAX_BYTES {
break;
}
bytes.extend_from_slice(&unit.to_le_bytes());
}
out.extend_from_slice(&(bytes.len() as u16).to_le_bytes());
out.extend_from_slice(&bytes);
}
}
}
}
fn value_as_i64(value: &serde_json::Value) -> Option<i64> {
match value {
serde_json::Value::Number(n) => n.as_i64().or_else(|| n.as_f64().map(|f| f as i64)),
serde_json::Value::Bool(b) => Some(i64::from(*b)),
serde_json::Value::String(s) => s.parse().ok(),
_ => None,
}
}
fn value_as_f64(value: &serde_json::Value) -> Option<f64> {
match value {
serde_json::Value::Number(n) => n.as_f64(),
serde_json::Value::Bool(b) => Some(if *b { 1.0 } else { 0.0 }),
serde_json::Value::String(s) => s.parse().ok(),
_ => None,
}
}
fn value_as_bool(value: &serde_json::Value) -> Option<bool> {
match value {
serde_json::Value::Bool(b) => Some(*b),
serde_json::Value::Number(n) => n.as_i64().map(|v| v != 0),
serde_json::Value::String(s) => match s.to_ascii_lowercase().as_str() {
"true" | "1" | "yes" => Some(true),
"false" | "0" | "no" => Some(false),
_ => None,
},
_ => None,
}
}
fn json_to_string(value: &serde_json::Value) -> String {
match value {
serde_json::Value::Null => String::new(),
serde_json::Value::Bool(b) => if *b { "1" } else { "0" }.to_string(),
serde_json::Value::Number(n) => n.to_string(),
serde_json::Value::String(s) => s.clone(),
serde_json::Value::Array(_) | serde_json::Value::Object(_) => value.to_string(),
}
}
#[cfg(not(target_arch = "wasm32"))]
fn parse_login7(data: &[u8]) -> (String, String, String) {
const IB_USERNAME: usize = 40;
const IB_APPNAME: usize = 48;
const IB_DATABASE: usize = 68;
let field = |pair_at: usize| -> String {
let Some(bytes) = data.get(pair_at..pair_at + 4) else {
return String::new();
};
let offset = u16::from_le_bytes([bytes[0], bytes[1]]) as usize;
let chars = u16::from_le_bytes([bytes[2], bytes[3]]) as usize;
if chars == 0 {
return String::new();
}
let Some(raw) = data.get(offset..offset + chars * 2) else {
return String::new();
};
let units: Vec<u16> = raw
.chunks_exact(2)
.map(|c| u16::from_le_bytes([c[0], c[1]]))
.collect();
String::from_utf16_lossy(&units)
};
(field(IB_USERNAME), field(IB_DATABASE), field(IB_APPNAME))
}