1pub mod actions;
12
13use crate::llm::action_helper::call_llm;
14use crate::llm::actions::protocol_trait::ActionResult;
15use crate::llm::ollama_client::OllamaClient;
16use crate::logging::emit::Log;
17use crate::protocol::Event;
18use crate::server::connection::ConnectionId;
19use crate::state::app_state::AppState;
20use crate::state::server::{ConnectionState, ConnectionStatus, ProtocolConnectionInfo};
21use crate::utils::clock::Instant;
22use actions::*;
23use anyhow::{Context, Result};
24use bytes::{Buf, BytesMut};
25use cassandra_protocol::compression::Compression;
26use cassandra_protocol::frame::{Direction, Envelope, Flags, Opcode, Version};
27use serde_json::json;
28use std::collections::HashMap;
29use std::net::SocketAddr;
30use std::sync::Arc;
31use tokio::io::{AsyncReadExt, AsyncWriteExt};
32use tokio::net::{TcpListener, TcpStream};
33use tokio::sync::{mpsc, Mutex};
34use tracing::{debug, error, trace, warn};
35
36pub const MAX_FRAME_BODY_BYTES: usize = 256 * 1024 * 1024;
38
39const FIRST_FRAME_READ_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);
46
47const IDLE_BETWEEN_FRAMES_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(900);
60
61const MAX_CONNECTIONS: usize = crate::server::accept_bounded::DEFAULT_MAX_CONNECTIONS;
66
67const CONNECTION_CAP_REFUSAL: &[u8] = b"";
77
78const MAX_PREPARED_STATEMENTS: usize = 1024;
85
86const CASSANDRA_ERROR_SERVER_ERROR: u32 = 0x0000;
88
89const CASSANDRA_ERROR_OVERLOADED: u32 = 0x1001;
92
93const CASSANDRA_ERROR_PROTOCOL: u32 = 0x000A;
96
97const CASSANDRA_ERROR_UNPREPARED: u32 = 0x2500;
103
104fn consistency_name(code: u16) -> &'static str {
110 match code {
111 0x0000 => "ANY",
112 0x0001 => "ONE",
113 0x0002 => "TWO",
114 0x0003 => "THREE",
115 0x0004 => "QUORUM",
116 0x0005 => "ALL",
117 0x0006 => "LOCAL_QUORUM",
118 0x0007 => "EACH_QUORUM",
119 0x0008 => "SERIAL",
120 0x0009 => "LOCAL_SERIAL",
121 0x000A => "LOCAL_ONE",
122 _ => "UNKNOWN",
123 }
124}
125
126pub struct CassandraServer {
128 llm_client: OllamaClient,
129 app_state: Arc<AppState>,
130 _status_tx: mpsc::UnboundedSender<String>,
131 server_id: Option<crate::state::ServerId>,
132}
133
134struct CassandraConnectionState {
136 ready: bool,
137 protocol_version: u8,
138 prepared_statements: HashMap<Vec<u8>, (String, usize)>,
140 authenticated: bool,
142 username: Option<String>,
144}
145
146struct ConnWrite {
154 write_half: Arc<Mutex<tokio::io::WriteHalf<TcpStream>>>,
155 app_state: Arc<AppState>,
156 server_id: Option<crate::state::ServerId>,
157 connection_id: ConnectionId,
158}
159
160impl ConnWrite {
161 async fn write(&self, data: &[u8]) -> Result<()> {
164 {
165 let mut writer = self.write_half.lock().await;
166 writer.write_all(data).await?;
167 writer.flush().await?;
168 }
169 if let Some(server_id) = self.server_id {
170 self.app_state
171 .update_connection_stats(
172 server_id,
173 self.connection_id,
174 None,
175 Some(data.len() as u64),
176 None,
177 Some(1),
178 )
179 .await;
180 }
181 Ok(())
182 }
183}
184
185impl CassandraServer {
186 pub fn new(
188 llm_client: OllamaClient,
189 app_state: Arc<AppState>,
190 status_tx: mpsc::UnboundedSender<String>,
191 server_id: Option<crate::state::ServerId>,
192 ) -> Self {
193 Self {
194 llm_client,
195 app_state,
196 _status_tx: status_tx,
197 server_id,
198 }
199 }
200
201 pub async fn spawn_with_llm_actions(
203 listen_addr: SocketAddr,
204 llm_client: OllamaClient,
205 app_state: Arc<AppState>,
206 status_tx: mpsc::UnboundedSender<String>,
207 server_id: crate::state::ServerId,
208 ) -> Result<SocketAddr> {
209 let listener = TcpListener::bind(listen_addr).await?;
210 let actual_addr = listener.local_addr()?;
211
212 Log::new(Some(&status_tx))
213 .info(format!("Cassandra/CQL server listening on {}", actual_addr));
214
215 let server = Arc::new(CassandraServer::new(
216 llm_client,
217 app_state.clone(),
218 status_tx.clone(),
219 Some(server_id),
220 ));
221
222 let limiter = crate::server::accept_bounded::ConnectionLimiter::new(MAX_CONNECTIONS);
224 let loop_state = app_state.clone();
225 let accept_handle = tokio::spawn(async move {
226 loop {
227 match crate::server::accept_bounded::accept_bounded(
228 &listener,
229 &limiter,
230 CONNECTION_CAP_REFUSAL,
231 "Cassandra",
232 Some(&status_tx),
233 )
234 .await
235 {
236 Ok((stream, addr, permit)) => {
237 Log::new(Some(&status_tx))
238 .debug(format!("Cassandra connection from {}", addr));
239
240 let server_clone = server.clone();
241 let status_tx_clone = status_tx.clone();
242
243 let task_owner = loop_state.clone();
245 task_owner
246 .spawn_server_task(server_id, async move {
247 let _permit = permit;
250 if let Err(e) = server_clone
251 .handle_connection(stream, addr, status_tx_clone)
252 .await
253 {
254 error!("Cassandra connection error: {}", e);
255 }
256 })
257 .await;
258 }
259 Err(e) => {
260 Log::new(Some(&status_tx))
264 .error(format!("Cassandra accept failed, listener stopped: {}", e));
265 break;
266 }
267 }
268 }
269 });
270
271 app_state
273 .register_server_task(server_id, accept_handle)
274 .await;
275
276 Ok(actual_addr)
277 }
278
279 async fn handle_connection(
281 &self,
282 stream: TcpStream,
283 addr: SocketAddr,
284 status_tx: mpsc::UnboundedSender<String>,
285 ) -> Result<()> {
286 let connection_id = ConnectionId::new(self.app_state.get_next_unified_id().await);
287 let local_addr = stream.local_addr().unwrap_or(addr);
288
289 if let Some(server_id) = self.server_id {
291 let now = Instant::now();
292 let conn_state = ConnectionState {
293 id: connection_id,
294 remote_addr: addr,
295 local_addr,
296 bytes_sent: 0,
297 bytes_received: 0,
298 packets_sent: 0,
299 packets_received: 0,
300 last_activity: now,
301 status: ConnectionStatus::Active,
302 status_changed_at: now,
303 protocol_info: ProtocolConnectionInfo::empty(),
304 };
305
306 self.app_state
307 .add_connection_to_server(server_id, conn_state)
308 .await;
309 }
310
311 let (reader, write_half) = tokio::io::split(stream);
313 let write_half = Arc::new(Mutex::new(write_half));
314 let out = ConnWrite {
315 write_half: write_half.clone(),
316 app_state: self.app_state.clone(),
317 server_id: self.server_id,
318 connection_id,
319 };
320
321 if let Some(server_id) = self.server_id {
326 let peer_rx = crate::server::peer_support::register_peer_channel(
327 &self.app_state,
328 server_id,
329 connection_id.as_u32(),
330 )
331 .await;
332 let protocol = Arc::new(CassandraProtocol::new(
333 connection_id,
334 self.app_state.clone(),
335 status_tx.clone(),
336 ));
337 crate::server::peer_support::spawn_peer_command_task(
338 peer_rx,
339 protocol,
340 self.app_state.clone(),
341 server_id,
342 connection_id.as_u32(),
343 write_half.clone(),
344 status_tx.clone(),
345 );
346 }
347
348 let result = self
349 .run_session(reader, &out, addr, connection_id, &status_tx)
350 .await;
351
352 if let Some(server_id) = self.server_id {
357 self.app_state
358 .remove_peer_handle(server_id, connection_id.as_u32())
359 .await;
360 }
361 let _ = write_half.lock().await.shutdown().await;
362
363 if let Some(server_id) = self.server_id {
365 self.app_state
366 .close_connection_on_server(server_id, connection_id)
367 .await;
368 }
369
370 result
371 }
372
373 async fn run_session(
375 &self,
376 mut reader: tokio::io::ReadHalf<TcpStream>,
377 out: &ConnWrite,
378 addr: SocketAddr,
379 connection_id: ConnectionId,
380 status_tx: &mpsc::UnboundedSender<String>,
381 ) -> Result<()> {
382 let mut conn_state = CassandraConnectionState {
383 ready: false,
384 protocol_version: 4,
385 prepared_statements: HashMap::new(),
386 authenticated: false,
387 username: None,
388 };
389
390 let mut buffer = BytesMut::with_capacity(4096);
391 let mut closing = false;
395
396 loop {
397 let read_timeout = if conn_state.ready {
404 IDLE_BETWEEN_FRAMES_TIMEOUT
405 } else {
406 FIRST_FRAME_READ_TIMEOUT
407 };
408
409 let read = match tokio::time::timeout(read_timeout, reader.read_buf(&mut buffer)).await
411 {
412 Ok(read) => read,
413 Err(_) => {
414 Log::new(Some(status_tx)).debug(format!(
415 "Cassandra client {} sent nothing for {}s; closing idle connection",
416 addr,
417 read_timeout.as_secs()
418 ));
419 break;
420 }
421 };
422 let n = match read {
423 Ok(0) => {
424 Log::new(Some(status_tx))
425 .debug(format!("Cassandra client {} disconnected", addr));
426 break;
427 }
428 Ok(n) => n,
429 Err(e) => {
430 error!("Read error from {}: {}", addr, e);
431 break;
432 }
433 };
434
435 if let Some(server_id) = self.server_id {
438 self.app_state
439 .update_connection_stats(
440 server_id,
441 connection_id,
442 Some(n as u64),
443 None,
444 Some(1),
445 None,
446 )
447 .await;
448 }
449
450 trace!("Read {} bytes from Cassandra client {}", n, addr);
451
452 while buffer.remaining() >= 9 {
454 let frame_start = buffer.as_ref();
456 if frame_start.len() < 9 {
457 break;
458 }
459
460 let length = u32::from_be_bytes([
462 frame_start[5],
463 frame_start[6],
464 frame_start[7],
465 frame_start[8],
466 ]) as usize;
467
468 if length > MAX_FRAME_BODY_BYTES {
473 Log::new(Some(status_tx)).error(format!(
474 "Cassandra frame too large ({} bytes, limit {}), closing {}",
475 length, MAX_FRAME_BODY_BYTES, addr
476 ));
477 closing = true;
478 break;
479 }
480
481 if buffer.remaining() < 9 + length {
483 trace!(
484 "Waiting for complete frame: have {}, need {}",
485 buffer.remaining(),
486 9 + length
487 );
488 break;
489 }
490
491 let frame_bytes = buffer.split_to(9 + length);
493
494 match self
495 .handle_frame(&frame_bytes, &mut conn_state, out, connection_id, status_tx)
496 .await
497 {
498 Ok(should_continue) => {
499 if !should_continue {
500 debug!("Closing Cassandra connection to {}", addr);
501 closing = true;
502 break;
503 }
504 }
505 Err(e) => {
506 Log::new(Some(status_tx)).error(format!(
518 "Cassandra connection {} decision=fail_closed_protocol_error: {}",
519 connection_id, e
520 ));
521 let stream_id = i16::from_be_bytes([frame_bytes[2], frame_bytes[3]]);
522 if let Err(send_err) = self
523 .send_error(
524 stream_id,
525 CASSANDRA_ERROR_PROTOCOL,
526 "netget: malformed CQL frame",
527 out,
528 status_tx,
529 )
530 .await
531 {
532 debug!(
533 "Cassandra could not report the protocol error to {}: {}",
534 addr, send_err
535 );
536 }
537 closing = true;
538 break;
539 }
540 }
541 }
542
543 if closing {
544 break;
545 }
546 }
547
548 Ok(())
549 }
550
551 async fn handle_frame(
553 &self,
554 frame_bytes: &[u8],
555 conn_state: &mut CassandraConnectionState,
556 out: &ConnWrite,
557 connection_id: ConnectionId,
558 status_tx: &mpsc::UnboundedSender<String>,
559 ) -> Result<bool> {
560 let parsed = Envelope::from_buffer(frame_bytes, Compression::None)
562 .context("Failed to parse Cassandra frame")?;
563 let frame = parsed.envelope;
564
565 Log::new(Some(status_tx)).trace(format!(
566 "Cassandra ← {:?} stream={}",
567 frame.opcode, frame.stream_id
568 ));
569
570 match frame.opcode {
571 Opcode::Startup => {
572 self.handle_startup(frame, conn_state, out, connection_id, status_tx)
573 .await
574 }
575 Opcode::Options => {
576 self.handle_options(frame, out, connection_id, status_tx)
577 .await
578 }
579 Opcode::Query => {
580 self.handle_query(frame, out, connection_id, status_tx)
581 .await
582 }
583 Opcode::Prepare => {
584 self.handle_prepare(frame, conn_state, out, connection_id, status_tx)
585 .await
586 }
587 Opcode::Execute => {
588 self.handle_execute(frame, conn_state, out, connection_id, status_tx)
589 .await
590 }
591 Opcode::AuthResponse => {
592 self.handle_auth_response(frame, conn_state, out, connection_id, status_tx)
593 .await
594 }
595 Opcode::Register => {
596 Log::new(Some(status_tx))
599 .debug("Cassandra: Client registered for events (not supported, no-op)");
600 self.send_ready(frame.stream_id, out, status_tx).await?;
601 Ok(true)
602 }
603 _ => {
604 Log::new(Some(status_tx))
605 .warn(format!("Unsupported Cassandra opcode: {:?}", frame.opcode));
606 self.send_error(
608 frame.stream_id,
609 0x000A,
610 "Unsupported operation",
611 out,
612 status_tx,
613 )
614 .await?;
615 Ok(true)
616 }
617 }
618 }
619
620 async fn handle_startup(
622 &self,
623 frame: Envelope,
624 conn_state: &mut CassandraConnectionState,
625 out: &ConnWrite,
626 connection_id: ConnectionId,
627 status_tx: &mpsc::UnboundedSender<String>,
628 ) -> Result<bool> {
629 debug!("Handling STARTUP from connection {}", connection_id);
630
631 let protocol =
636 CassandraProtocol::new(connection_id, self.app_state.clone(), status_tx.clone());
637
638 let event = Event {
639 event_type: &CASSANDRA_STARTUP_EVENT,
640 data: json!({
641 "protocol_version": conn_state.protocol_version,
642 "options": {"CQL_VERSION": "3.0.0"}
643 }),
644 };
645
646 let server_id = self.server_id.context("Server ID not set")?;
647
648 let execution_result = match call_llm(
649 &self.llm_client,
650 &self.app_state,
651 server_id,
652 Some(connection_id),
653 &event,
654 &protocol,
655 )
656 .await
657 {
658 Ok(result) => result,
659 Err(e) => {
660 self.send_llm_failure_error(
661 frame.stream_id,
662 "STARTUP",
663 &e,
664 out,
665 connection_id,
666 status_tx,
667 )
668 .await?;
669 return Ok(true);
670 }
671 };
672
673 for message in &execution_result.messages {
675 Log::new(Some(status_tx)).info(format!("{}", message));
676 }
677
678 for action_result in execution_result.protocol_results {
680 match action_result {
681 ActionResult::Custom { name, data } => match name.as_str() {
682 "cassandra_ready" => {
683 conn_state.ready = true;
684 self.send_ready(frame.stream_id, out, status_tx).await?;
685 return Ok(true);
686 }
687 "cassandra_authenticate" => {
692 let authenticator = data
693 .get("authenticator")
694 .and_then(|v| v.as_str())
695 .unwrap_or("org.apache.cassandra.auth.PasswordAuthenticator");
696 self.send_authenticate(frame.stream_id, authenticator, out, status_tx)
697 .await?;
698 return Ok(true);
699 }
700 "cassandra_error" => {
701 let error_code = data
702 .get("error_code")
703 .and_then(|v| v.as_u64())
704 .unwrap_or(0x0000) as u32;
705 let message = data
706 .get("message")
707 .and_then(|v| v.as_str())
708 .unwrap_or("Unknown error");
709 self.send_model_error(
710 frame.stream_id,
711 "STARTUP",
712 error_code,
713 message,
714 out,
715 connection_id,
716 status_tx,
717 )
718 .await?;
719 return Ok(true);
720 }
721 _ => {}
722 },
723 ActionResult::CloseConnection => {
724 return Ok(false);
725 }
726 _ => {
727 warn!("Unexpected action result for STARTUP");
728 }
729 }
730 }
731
732 self.send_no_answer_error(frame.stream_id, "STARTUP", out, connection_id, status_tx)
737 .await?;
738 Ok(false)
739 }
740
741 async fn handle_options(
743 &self,
744 frame: Envelope,
745 out: &ConnWrite,
746 connection_id: ConnectionId,
747 status_tx: &mpsc::UnboundedSender<String>,
748 ) -> Result<bool> {
749 debug!("Handling OPTIONS from connection {}", connection_id);
750
751 let protocol =
752 CassandraProtocol::new(connection_id, self.app_state.clone(), status_tx.clone());
753
754 let event = Event {
755 event_type: &CASSANDRA_OPTIONS_EVENT,
756 data: json!({}),
757 };
758
759 let server_id = self.server_id.context("Server ID not set")?;
760
761 let execution_result = match call_llm(
762 &self.llm_client,
763 &self.app_state,
764 server_id,
765 Some(connection_id),
766 &event,
767 &protocol,
768 )
769 .await
770 {
771 Ok(result) => result,
772 Err(e) => {
773 self.send_llm_failure_error(
774 frame.stream_id,
775 "OPTIONS",
776 &e,
777 out,
778 connection_id,
779 status_tx,
780 )
781 .await?;
782 return Ok(true);
783 }
784 };
785
786 for message in &execution_result.messages {
788 Log::new(Some(status_tx)).info(format!("{}", message));
789 }
790
791 for action_result in execution_result.protocol_results {
793 match action_result {
794 ActionResult::Custom { name, data } => match name.as_str() {
795 "cassandra_supported" => {
796 let options = data
797 .get("options")
798 .and_then(|v| v.as_object())
799 .cloned()
800 .unwrap_or_default();
801 self.send_supported(frame.stream_id, options, out, status_tx)
802 .await?;
803 return Ok(true);
804 }
805 "cassandra_error" => {
806 let error_code = data
807 .get("error_code")
808 .and_then(|v| v.as_u64())
809 .unwrap_or(0x0000) as u32;
810 let message = data
811 .get("message")
812 .and_then(|v| v.as_str())
813 .unwrap_or("Unknown error");
814 self.send_model_error(
815 frame.stream_id,
816 "OPTIONS",
817 error_code,
818 message,
819 out,
820 connection_id,
821 status_tx,
822 )
823 .await?;
824 return Ok(true);
825 }
826 _ => {}
827 },
828 ActionResult::CloseConnection => {
829 return Ok(false);
830 }
831 _ => {
832 warn!("Unexpected action result for OPTIONS");
833 }
834 }
835 }
836
837 self.send_supported(frame.stream_id, serde_json::Map::new(), out, status_tx)
839 .await?;
840 Ok(true)
841 }
842
843 async fn handle_query(
845 &self,
846 frame: Envelope,
847 out: &ConnWrite,
848 connection_id: ConnectionId,
849 status_tx: &mpsc::UnboundedSender<String>,
850 ) -> Result<bool> {
851 let (query_str, consistency) = self.parse_query(&frame)?;
853
854 Log::new(Some(status_tx)).debug(format!(
855 "Cassandra ← Query from connection {}: {}",
856 connection_id, query_str
857 ));
858
859 let protocol =
860 CassandraProtocol::new(connection_id, self.app_state.clone(), status_tx.clone());
861
862 let event = Event {
863 event_type: &CASSANDRA_QUERY_EVENT,
864 data: json!({
865 "query": query_str,
866 "consistency": consistency
867 }),
868 };
869
870 let server_id = self.server_id.context("Server ID not set")?;
871
872 let execution_result = match call_llm(
873 &self.llm_client,
874 &self.app_state,
875 server_id,
876 Some(connection_id),
877 &event,
878 &protocol,
879 )
880 .await
881 {
882 Ok(result) => result,
883 Err(e) => {
884 self.send_llm_failure_error(
885 frame.stream_id,
886 "QUERY",
887 &e,
888 out,
889 connection_id,
890 status_tx,
891 )
892 .await?;
893 return Ok(true);
894 }
895 };
896
897 for message in &execution_result.messages {
899 Log::new(Some(status_tx)).info(format!("{}", message));
900 }
901
902 for action_result in execution_result.protocol_results {
904 match action_result {
905 ActionResult::Custom { name, data } => match name.as_str() {
906 "cassandra_result_rows" => {
907 let columns = data
908 .get("columns")
909 .and_then(|v| v.as_array())
910 .cloned()
911 .unwrap_or_default();
912 let rows = data
913 .get("rows")
914 .and_then(|v| v.as_array())
915 .cloned()
916 .unwrap_or_default();
917 self.send_result_rows(frame.stream_id, columns, rows, out, status_tx)
918 .await?;
919 return Ok(true);
920 }
921 "cassandra_error" => {
922 let error_code = data
923 .get("error_code")
924 .and_then(|v| v.as_u64())
925 .unwrap_or(0x0000) as u32;
926 let message = data
927 .get("message")
928 .and_then(|v| v.as_str())
929 .unwrap_or("Unknown error");
930 self.send_model_error(
931 frame.stream_id,
932 "QUERY",
933 error_code,
934 message,
935 out,
936 connection_id,
937 status_tx,
938 )
939 .await?;
940 return Ok(true);
941 }
942 _ => {}
943 },
944 ActionResult::CloseConnection => {
945 return Ok(false);
946 }
947 _ => {
948 warn!("Unexpected action result for QUERY");
949 }
950 }
951 }
952
953 self.send_no_answer_error(frame.stream_id, "QUERY", out, connection_id, status_tx)
956 .await?;
957 Ok(true)
958 }
959
960 fn parse_query(&self, frame: &Envelope) -> Result<(String, &'static str)> {
971 let body = &frame.body;
972 if body.len() < 4 {
973 return Err(anyhow::anyhow!("Query frame too short"));
974 }
975
976 let query_len = u32::from_be_bytes([body[0], body[1], body[2], body[3]]) as usize;
978
979 if body.len().saturating_sub(4) < query_len {
981 return Err(anyhow::anyhow!("Query frame truncated"));
982 }
983
984 let query_bytes = &body[4..4 + query_len];
985 let query_str = String::from_utf8_lossy(query_bytes).to_string();
986
987 let consistency = body
991 .get(4 + query_len..6 + query_len)
992 .map(|b| consistency_name(u16::from_be_bytes([b[0], b[1]])))
993 .unwrap_or("UNKNOWN");
994
995 Ok((query_str, consistency))
996 }
997
998 async fn send_ready(
1000 &self,
1001 stream_id: i16,
1002 out: &ConnWrite,
1003 status_tx: &mpsc::UnboundedSender<String>,
1004 ) -> Result<()> {
1005 let response = Envelope {
1006 version: Version::V4,
1007 direction: Direction::Response,
1008 flags: Flags::empty(),
1009 stream_id: stream_id,
1010 opcode: Opcode::Ready,
1011 body: vec![],
1012 tracing_id: None,
1013 warnings: vec![],
1014 };
1015
1016 let bytes = response.encode_with(Compression::None)?;
1017 out.write(&bytes).await?;
1018
1019 Log::new(Some(status_tx)).trace(format!("Cassandra → READY ({} bytes)", bytes.len()));
1020
1021 Ok(())
1022 }
1023
1024 async fn send_supported(
1026 &self,
1027 stream_id: i16,
1028 options: serde_json::Map<String, serde_json::Value>,
1029 out: &ConnWrite,
1030 status_tx: &mpsc::UnboundedSender<String>,
1031 ) -> Result<()> {
1032 let mut body = Vec::new();
1035
1036 body.extend_from_slice(&(options.len() as u16).to_be_bytes());
1038
1039 for (key, value) in options.iter() {
1040 let key_bytes = key.as_bytes();
1042 body.extend_from_slice(&(key_bytes.len() as u16).to_be_bytes());
1043 body.extend_from_slice(key_bytes);
1044
1045 if let Some(arr) = value.as_array() {
1047 body.extend_from_slice(&(arr.len() as u16).to_be_bytes());
1048 for item in arr {
1049 if let Some(s) = item.as_str() {
1050 let s_bytes = s.as_bytes();
1051 body.extend_from_slice(&(s_bytes.len() as u16).to_be_bytes());
1052 body.extend_from_slice(s_bytes);
1053 }
1054 }
1055 } else {
1056 body.extend_from_slice(&0u16.to_be_bytes());
1058 }
1059 }
1060
1061 let response = Envelope {
1062 version: Version::V4,
1063 direction: Direction::Response,
1064 flags: Flags::empty(),
1065 stream_id: stream_id,
1066 opcode: Opcode::Supported,
1067 body,
1068 tracing_id: None,
1069 warnings: vec![],
1070 };
1071
1072 let bytes = response.encode_with(Compression::None)?;
1073 out.write(&bytes).await?;
1074
1075 Log::new(Some(status_tx)).trace(format!("Cassandra → SUPPORTED ({} bytes)", bytes.len()));
1076
1077 Ok(())
1078 }
1079
1080 async fn send_result_rows(
1082 &self,
1083 stream_id: i16,
1084 columns: Vec<serde_json::Value>,
1085 rows: Vec<serde_json::Value>,
1086 out: &ConnWrite,
1087 status_tx: &mpsc::UnboundedSender<String>,
1088 ) -> Result<()> {
1089 let mut body = Vec::new();
1092
1093 body.extend_from_slice(&0x00000002u32.to_be_bytes());
1095
1096 body.extend_from_slice(&0x00000001u32.to_be_bytes());
1099
1100 body.extend_from_slice(&(columns.len() as u32).to_be_bytes());
1102
1103 let keyspace = b"system";
1105 let table = b"local";
1106 body.extend_from_slice(&(keyspace.len() as u16).to_be_bytes());
1107 body.extend_from_slice(keyspace);
1108 body.extend_from_slice(&(table.len() as u16).to_be_bytes());
1109 body.extend_from_slice(table);
1110
1111 for col in &columns {
1113 let name = col.get("name").and_then(|v| v.as_str()).unwrap_or("col");
1114 let col_type = col
1115 .get("type")
1116 .and_then(|v| v.as_str())
1117 .unwrap_or("varchar");
1118
1119 body.extend_from_slice(&(name.len() as u16).to_be_bytes());
1121 body.extend_from_slice(name.as_bytes());
1122
1123 let type_code: u16 = match col_type {
1125 "int" => 0x0009,
1126 "boolean" => 0x0004,
1127 _ => 0x000D, };
1129 body.extend_from_slice(&type_code.to_be_bytes());
1130 }
1131
1132 body.extend_from_slice(&(rows.len() as u32).to_be_bytes());
1134
1135 for row in &rows {
1143 let empty = Vec::new();
1144 let row_arr = row.as_array().unwrap_or(&empty);
1145 if row_arr.len() != columns.len() {
1146 warn!(
1147 "Cassandra row has {} cell(s) but {} column(s) were declared; padding/truncating",
1148 row_arr.len(),
1149 columns.len()
1150 );
1151 }
1152 for i in 0..columns.len() {
1153 let col_type = columns[i].get("type").and_then(|v| v.as_str());
1154 let cell_bytes = row_arr
1155 .get(i)
1156 .and_then(|cell| self.serialize_cell_value(cell, col_type));
1157
1158 if let Some(bytes) = cell_bytes {
1159 body.extend_from_slice(&(bytes.len() as i32).to_be_bytes());
1160 body.extend_from_slice(&bytes);
1161 } else {
1162 body.extend_from_slice(&(-1i32).to_be_bytes());
1164 }
1165 }
1166 }
1167
1168 let response = Envelope {
1169 version: Version::V4,
1170 direction: Direction::Response,
1171 flags: Flags::empty(),
1172 stream_id: stream_id,
1173 opcode: Opcode::Result,
1174 body,
1175 tracing_id: None,
1176 warnings: vec![],
1177 };
1178
1179 let bytes = response.encode_with(Compression::None)?;
1180 out.write(&bytes).await?;
1181
1182 Log::new(Some(status_tx)).trace(format!(
1183 "Cassandra → RESULT ({} rows, {} bytes)",
1184 rows.len(),
1185 bytes.len()
1186 ));
1187
1188 Ok(())
1189 }
1190
1191 fn serialize_cell_value(
1203 &self,
1204 value: &serde_json::Value,
1205 col_type: Option<&str>,
1206 ) -> Option<Vec<u8>> {
1207 if value.is_null() {
1208 return None;
1209 }
1210
1211 match col_type.unwrap_or("varchar") {
1212 "int" => {
1213 let n = match value {
1214 serde_json::Value::Number(n) => n.as_i64(),
1215 serde_json::Value::String(s) => s.trim().parse::<i64>().ok(),
1216 serde_json::Value::Bool(b) => Some(if *b { 1 } else { 0 }),
1217 _ => None,
1218 };
1219 match n {
1220 Some(n) if n >= i32::MIN as i64 && n <= i32::MAX as i64 => {
1223 Some((n as i32).to_be_bytes().to_vec())
1224 }
1225 _ => {
1226 warn!("Cassandra int column got {:?}, sending NULL", value);
1227 None
1228 }
1229 }
1230 }
1231 "boolean" => {
1232 let b = match value {
1233 serde_json::Value::Bool(b) => Some(*b),
1234 serde_json::Value::Number(n) => n.as_i64().map(|n| n != 0),
1235 serde_json::Value::String(s) => match s.trim().to_ascii_lowercase().as_str() {
1236 "true" | "1" | "yes" => Some(true),
1237 "false" | "0" | "no" => Some(false),
1238 _ => None,
1239 },
1240 _ => None,
1241 };
1242 match b {
1243 Some(b) => Some(vec![u8::from(b)]),
1244 None => {
1245 warn!("Cassandra boolean column got {:?}, sending NULL", value);
1246 None
1247 }
1248 }
1249 }
1250 _ => match value {
1254 serde_json::Value::String(s) => Some(s.as_bytes().to_vec()),
1255 other => Some(other.to_string().into_bytes()),
1256 },
1257 }
1258 }
1259
1260 async fn handle_prepare(
1262 &self,
1263 frame: Envelope,
1264 conn_state: &mut CassandraConnectionState,
1265 out: &ConnWrite,
1266 connection_id: ConnectionId,
1267 status_tx: &mpsc::UnboundedSender<String>,
1268 ) -> Result<bool> {
1269 debug!("Handling PREPARE from connection {}", connection_id);
1270
1271 let (query, _consistency) = self.parse_query(&frame)?;
1275 trace!("PREPARE query: {}", query);
1276
1277 use std::collections::hash_map::DefaultHasher;
1279 use std::hash::{Hash, Hasher};
1280 let mut hasher = DefaultHasher::new();
1281 query.hash(&mut hasher);
1282 let hash = hasher.finish();
1283 let statement_id = hash.to_be_bytes().to_vec();
1284
1285 let param_count = query.matches('?').count();
1287
1288 if conn_state.prepared_statements.len() >= MAX_PREPARED_STATEMENTS
1290 && !conn_state.prepared_statements.contains_key(&statement_id)
1291 {
1292 warn!(
1293 "Cassandra connection {} reached {} prepared statements; rejecting PREPARE",
1294 connection_id, MAX_PREPARED_STATEMENTS
1295 );
1296 self.send_error(
1297 frame.stream_id,
1298 0x2200,
1299 "Too many prepared statements on this connection",
1300 out,
1301 status_tx,
1302 )
1303 .await?;
1304 return Ok(true);
1305 }
1306 conn_state
1307 .prepared_statements
1308 .insert(statement_id.clone(), (query.clone(), param_count));
1309
1310 debug!(
1311 "Prepared statement ID {:?} with {} params",
1312 statement_id, param_count
1313 );
1314
1315 let protocol =
1317 CassandraProtocol::new(connection_id, self.app_state.clone(), status_tx.clone());
1318
1319 let event = Event {
1320 event_type: &CASSANDRA_PREPARE_EVENT,
1321 data: json!({
1322 "query": query,
1323 "statement_id": hex::encode(&statement_id),
1324 "param_count": param_count,
1325 }),
1326 };
1327
1328 let server_id = self.server_id.context("Server ID not set")?;
1329
1330 let execution_result = match call_llm(
1331 &self.llm_client,
1332 &self.app_state,
1333 server_id,
1334 Some(connection_id),
1335 &event,
1336 &protocol,
1337 )
1338 .await
1339 {
1340 Ok(result) => result,
1341 Err(e) => {
1342 self.send_llm_failure_error(
1343 frame.stream_id,
1344 "PREPARE",
1345 &e,
1346 out,
1347 connection_id,
1348 status_tx,
1349 )
1350 .await?;
1351 return Ok(true);
1352 }
1353 };
1354
1355 for message in &execution_result.messages {
1357 Log::new(Some(status_tx)).info(format!("{}", message));
1358 }
1359
1360 for action_result in execution_result.protocol_results {
1362 match action_result {
1363 ActionResult::Custom { name, data } => match name.as_str() {
1364 "cassandra_prepared" => {
1365 let columns = data
1366 .get("columns")
1367 .and_then(|v| v.as_array())
1368 .cloned()
1369 .unwrap_or_default();
1370 let params = data.get("params").and_then(|v| v.as_array()).cloned();
1371 self.send_prepared(
1372 frame.stream_id,
1373 statement_id,
1374 columns,
1375 params,
1376 param_count,
1377 out,
1378 status_tx,
1379 )
1380 .await?;
1381 return Ok(true);
1382 }
1383 "cassandra_error" => {
1384 let error_code = data
1385 .get("error_code")
1386 .and_then(|v| v.as_u64())
1387 .unwrap_or(0x0000) as u32;
1388 let message = data
1389 .get("message")
1390 .and_then(|v| v.as_str())
1391 .unwrap_or("Unknown error");
1392 self.send_model_error(
1393 frame.stream_id,
1394 "PREPARE",
1395 error_code,
1396 message,
1397 out,
1398 connection_id,
1399 status_tx,
1400 )
1401 .await?;
1402 return Ok(true);
1403 }
1404 _ => {}
1405 },
1406 ActionResult::CloseConnection => {
1407 return Ok(false);
1408 }
1409 _ => {
1410 warn!("Unexpected action result for PREPARE");
1411 }
1412 }
1413 }
1414
1415 self.send_no_answer_error(frame.stream_id, "PREPARE", out, connection_id, status_tx)
1418 .await?;
1419 Ok(true)
1420 }
1421
1422 async fn handle_execute(
1424 &self,
1425 frame: Envelope,
1426 conn_state: &mut CassandraConnectionState,
1427 out: &ConnWrite,
1428 connection_id: ConnectionId,
1429 status_tx: &mpsc::UnboundedSender<String>,
1430 ) -> Result<bool> {
1431 debug!("Handling EXECUTE from connection {}", connection_id);
1432
1433 let (statement_id, params, consistency) = self.parse_execute(&frame)?;
1435
1436 let Some((query, expected_param_count)) = conn_state
1446 .prepared_statements
1447 .get(&statement_id)
1448 .map(|(q, n)| (q.clone(), *n))
1449 else {
1450 Log::new(Some(status_tx)).warn(format!(
1451 "Cassandra connection {} decision=unprepared stage=EXECUTE: statement id {} is \
1452 not prepared on this connection; answering ERROR 0x{:04X} so the driver \
1453 re-prepares",
1454 connection_id,
1455 hex::encode(&statement_id),
1456 CASSANDRA_ERROR_UNPREPARED
1457 ));
1458 self.send_unprepared_error(frame.stream_id, &statement_id, out, status_tx)
1459 .await?;
1460 return Ok(true);
1461 };
1462
1463 trace!("EXECUTE statement: {} with {} params", query, params.len());
1464
1465 if params.len() != expected_param_count {
1467 let err_msg = format!(
1468 "Expected {} parameters, got {}",
1469 expected_param_count,
1470 params.len()
1471 );
1472 self.send_error(frame.stream_id, 0x2200, &err_msg, out, status_tx)
1473 .await?;
1474 return Ok(true);
1475 }
1476
1477 let protocol =
1479 CassandraProtocol::new(connection_id, self.app_state.clone(), status_tx.clone());
1480
1481 let event = Event {
1482 event_type: &CASSANDRA_EXECUTE_EVENT,
1483 data: json!({
1484 "query": query,
1485 "statement_id": hex::encode(&statement_id),
1486 "parameters": params,
1487 "consistency": consistency,
1488 }),
1489 };
1490
1491 let server_id = self.server_id.context("Server ID not set")?;
1492
1493 let execution_result = match call_llm(
1494 &self.llm_client,
1495 &self.app_state,
1496 server_id,
1497 Some(connection_id),
1498 &event,
1499 &protocol,
1500 )
1501 .await
1502 {
1503 Ok(result) => result,
1504 Err(e) => {
1505 self.send_llm_failure_error(
1506 frame.stream_id,
1507 "EXECUTE",
1508 &e,
1509 out,
1510 connection_id,
1511 status_tx,
1512 )
1513 .await?;
1514 return Ok(true);
1515 }
1516 };
1517
1518 for message in &execution_result.messages {
1520 Log::new(Some(status_tx)).info(format!("{}", message));
1521 }
1522
1523 for action_result in execution_result.protocol_results {
1525 match action_result {
1526 ActionResult::Custom { name, data } => match name.as_str() {
1527 "cassandra_result_rows" => {
1528 let columns = data
1529 .get("columns")
1530 .and_then(|v| v.as_array())
1531 .cloned()
1532 .unwrap_or_default();
1533 let rows = data
1534 .get("rows")
1535 .and_then(|v| v.as_array())
1536 .cloned()
1537 .unwrap_or_default();
1538 self.send_result_rows(frame.stream_id, columns, rows, out, status_tx)
1539 .await?;
1540 return Ok(true);
1541 }
1542 "cassandra_error" => {
1543 let error_code = data
1544 .get("error_code")
1545 .and_then(|v| v.as_u64())
1546 .unwrap_or(0x0000) as u32;
1547 let message = data
1548 .get("message")
1549 .and_then(|v| v.as_str())
1550 .unwrap_or("Unknown error");
1551 self.send_model_error(
1552 frame.stream_id,
1553 "EXECUTE",
1554 error_code,
1555 message,
1556 out,
1557 connection_id,
1558 status_tx,
1559 )
1560 .await?;
1561 return Ok(true);
1562 }
1563 _ => {}
1564 },
1565 ActionResult::CloseConnection => {
1566 return Ok(false);
1567 }
1568 _ => {
1569 warn!("Unexpected action result for EXECUTE");
1570 }
1571 }
1572 }
1573
1574 self.send_no_answer_error(frame.stream_id, "EXECUTE", out, connection_id, status_tx)
1577 .await?;
1578 Ok(true)
1579 }
1580
1581 fn parse_execute(
1583 &self,
1584 frame: &Envelope,
1585 ) -> Result<(Vec<u8>, Vec<serde_json::Value>, &'static str)> {
1586 let body = &frame.body;
1587 if body.len() < 2 {
1588 return Err(anyhow::anyhow!("EXECUTE frame too short"));
1589 }
1590
1591 let id_len = u16::from_be_bytes([body[0], body[1]]) as usize;
1593 if body.len() < 2 + id_len {
1594 return Err(anyhow::anyhow!("EXECUTE frame truncated (statement ID)"));
1595 }
1596
1597 let statement_id = body[2..2 + id_len].to_vec();
1598 let mut offset = 2 + id_len;
1599
1600 if body.len() < offset + 2 {
1604 return Err(anyhow::anyhow!("EXECUTE frame truncated (consistency)"));
1605 }
1606 let consistency = consistency_name(u16::from_be_bytes([body[offset], body[offset + 1]]));
1607 offset += 2;
1608
1609 if body.len() < offset + 1 {
1611 return Err(anyhow::anyhow!("EXECUTE frame truncated (flags)"));
1612 }
1613 let flags = body[offset];
1614 offset += 1;
1615
1616 let mut params = Vec::new();
1617
1618 if flags & 0x01 != 0 {
1620 if body.len() < offset + 2 {
1622 return Err(anyhow::anyhow!("EXECUTE frame truncated (param count)"));
1623 }
1624 let param_count = u16::from_be_bytes([body[offset], body[offset + 1]]) as usize;
1625 offset += 2;
1626
1627 for _ in 0..param_count {
1629 if body.len() < offset + 4 {
1630 return Err(anyhow::anyhow!("EXECUTE frame truncated (param length)"));
1631 }
1632
1633 let param_len = i32::from_be_bytes([
1634 body[offset],
1635 body[offset + 1],
1636 body[offset + 2],
1637 body[offset + 3],
1638 ]);
1639 offset += 4;
1640
1641 if param_len < 0 {
1642 params.push(serde_json::Value::Null);
1644 } else {
1645 let param_len = param_len as usize;
1646 if body.len() < offset + param_len {
1647 return Err(anyhow::anyhow!("EXECUTE frame truncated (param data)"));
1648 }
1649
1650 let param_bytes = &body[offset..offset + param_len];
1651 let param_str = String::from_utf8_lossy(param_bytes).to_string();
1653 params.push(json!(param_str));
1654 offset += param_len;
1655 }
1656 }
1657 }
1658
1659 Ok((statement_id, params, consistency))
1660 }
1661
1662 async fn send_prepared(
1664 &self,
1665 stream_id: i16,
1666 statement_id: Vec<u8>,
1667 columns: Vec<serde_json::Value>,
1668 params: Option<Vec<serde_json::Value>>,
1669 param_count: usize,
1670 out: &ConnWrite,
1671 status_tx: &mpsc::UnboundedSender<String>,
1672 ) -> Result<()> {
1673 let mut body = Vec::new();
1674
1675 body.extend_from_slice(&0x00000004u32.to_be_bytes());
1677
1678 body.extend_from_slice(&(statement_id.len() as u16).to_be_bytes());
1680 body.extend_from_slice(&statement_id);
1681
1682 let keyspace = b"netget";
1684 let table = b"data";
1685
1686 body.extend_from_slice(&0x00000001u32.to_be_bytes());
1688
1689 body.extend_from_slice(&(param_count as u32).to_be_bytes());
1691
1692 body.extend_from_slice(&0u32.to_be_bytes()); body.extend_from_slice(&(keyspace.len() as u16).to_be_bytes());
1698 body.extend_from_slice(keyspace);
1699 body.extend_from_slice(&(table.len() as u16).to_be_bytes());
1700 body.extend_from_slice(table);
1701
1702 for i in 0..param_count {
1704 let param_name = format!("param{}", i);
1705 let param_bytes = param_name.as_bytes();
1706 body.extend_from_slice(&(param_bytes.len() as u16).to_be_bytes());
1707 body.extend_from_slice(param_bytes);
1708
1709 let param_type = params
1711 .as_ref()
1712 .and_then(|p| p.get(i))
1713 .and_then(|param| param.get("type"))
1714 .and_then(|v| v.as_str())
1715 .unwrap_or("varchar");
1716
1717 let type_code: u16 = match param_type {
1719 "int" => 0x0009,
1720 "varchar" | "text" => 0x000D,
1721 "boolean" => 0x0004,
1722 _ => 0x000D, };
1724 body.extend_from_slice(&type_code.to_be_bytes());
1725 }
1726
1727 body.extend_from_slice(&0x00000001u32.to_be_bytes());
1731
1732 body.extend_from_slice(&(columns.len() as u32).to_be_bytes());
1734
1735 body.extend_from_slice(&(keyspace.len() as u16).to_be_bytes());
1739 body.extend_from_slice(keyspace);
1740 body.extend_from_slice(&(table.len() as u16).to_be_bytes());
1741 body.extend_from_slice(table);
1742
1743 for col in columns {
1745 let name = col.get("name").and_then(|v| v.as_str()).unwrap_or("col");
1746 let col_type = col
1747 .get("type")
1748 .and_then(|v| v.as_str())
1749 .unwrap_or("varchar");
1750
1751 let name_bytes = name.as_bytes();
1753 body.extend_from_slice(&(name_bytes.len() as u16).to_be_bytes());
1754 body.extend_from_slice(name_bytes);
1755
1756 let type_code: u16 = match col_type {
1758 "int" => 0x0009,
1759 "varchar" | "text" => 0x000D,
1760 "boolean" => 0x0004,
1761 _ => 0x000D, };
1763 body.extend_from_slice(&type_code.to_be_bytes());
1764 }
1765
1766 trace!(
1767 "PREPARED response body hex (first 200 bytes): {}",
1768 hex::encode(&body[..std::cmp::min(200, body.len())])
1769 );
1770
1771 let response = Envelope {
1772 version: Version::V4,
1773 direction: Direction::Response,
1774 flags: Flags::empty(),
1775 stream_id: stream_id,
1776 opcode: Opcode::Result,
1777 body,
1778 tracing_id: None,
1779 warnings: vec![],
1780 };
1781
1782 let bytes = response.encode_with(Compression::None)?;
1783 out.write(&bytes).await?;
1784
1785 Log::new(Some(status_tx)).trace(format!(
1786 "Cassandra → RESULT (Prepared: {} params, {} bytes)",
1787 param_count,
1788 bytes.len()
1789 ));
1790
1791 Ok(())
1792 }
1793
1794 async fn send_llm_failure_error(
1811 &self,
1812 stream_id: i16,
1813 stage: &str,
1814 err: &anyhow::Error,
1815 out: &ConnWrite,
1816 connection_id: ConnectionId,
1817 status_tx: &mpsc::UnboundedSender<String>,
1818 ) -> Result<()> {
1819 let overloaded = crate::llm::is_overload_error(err);
1820 let (code, label) = if overloaded {
1821 (CASSANDRA_ERROR_OVERLOADED, "Overloaded")
1822 } else {
1823 (CASSANDRA_ERROR_SERVER_ERROR, "Server error")
1824 };
1825 let message = crate::utils::WireFailure::classify(err).prefixed_text();
1826 Log::new(Some(status_tx)).warn(format!(
1827 "Cassandra connection {} answering {} with ERROR 0x{:04X} ({}): {}",
1828 connection_id, stage, code, label, message
1829 ));
1830 self.send_error(stream_id, code, &message, out, status_tx)
1831 .await
1832 }
1833
1834 async fn send_no_answer_error(
1850 &self,
1851 stream_id: i16,
1852 stage: &str,
1853 out: &ConnWrite,
1854 connection_id: ConnectionId,
1855 status_tx: &mpsc::UnboundedSender<String>,
1856 ) -> Result<()> {
1857 Log::new(Some(status_tx)).warn(format!(
1858 "Cassandra connection {} decision=fail_closed_no_answer stage={}: handler produced \
1859 no recognised action, answering ERROR 0x{:04X} (Server error)",
1860 connection_id, stage, CASSANDRA_ERROR_SERVER_ERROR
1861 ));
1862 self.send_error(
1863 stream_id,
1864 CASSANDRA_ERROR_SERVER_ERROR,
1865 "netget: no handler answer for this request",
1866 out,
1867 status_tx,
1868 )
1869 .await
1870 }
1871
1872 async fn send_model_error(
1877 &self,
1878 stream_id: i16,
1879 stage: &str,
1880 error_code: u32,
1881 message: &str,
1882 out: &ConnWrite,
1883 connection_id: ConnectionId,
1884 status_tx: &mpsc::UnboundedSender<String>,
1885 ) -> Result<()> {
1886 Log::new(Some(status_tx)).info(format!(
1887 "Cassandra connection {} decision=model_reject stage={}: ERROR 0x{:04X}",
1888 connection_id, stage, error_code
1889 ));
1890 self.send_error(stream_id, error_code, message, out, status_tx)
1891 .await
1892 }
1893
1894 async fn send_unprepared_error(
1901 &self,
1902 stream_id: i16,
1903 statement_id: &[u8],
1904 out: &ConnWrite,
1905 status_tx: &mpsc::UnboundedSender<String>,
1906 ) -> Result<()> {
1907 let message = "netget: statement not prepared on this connection";
1908 let mut body = Vec::new();
1909 body.extend_from_slice(&CASSANDRA_ERROR_UNPREPARED.to_be_bytes());
1910 body.extend_from_slice(&(message.len() as u16).to_be_bytes());
1911 body.extend_from_slice(message.as_bytes());
1912 let id_len = u16::try_from(statement_id.len()).unwrap_or(u16::MAX) as usize;
1916 body.extend_from_slice(&(id_len as u16).to_be_bytes());
1917 body.extend_from_slice(&statement_id[..id_len]);
1918
1919 let response = Envelope {
1920 version: Version::V4,
1921 direction: Direction::Response,
1922 flags: Flags::empty(),
1923 stream_id,
1924 opcode: Opcode::Error,
1925 body,
1926 tracing_id: None,
1927 warnings: vec![],
1928 };
1929
1930 let bytes = response.encode_with(Compression::None)?;
1931 out.write(&bytes).await?;
1932
1933 Log::new(Some(status_tx)).trace(format!(
1934 "Cassandra → ERROR 0x{:04X} UNPREPARED id={}",
1935 CASSANDRA_ERROR_UNPREPARED,
1936 hex::encode(statement_id)
1937 ));
1938
1939 Ok(())
1940 }
1941
1942 async fn send_error(
1944 &self,
1945 stream_id: i16,
1946 error_code: u32,
1947 message: &str,
1948 out: &ConnWrite,
1949 status_tx: &mpsc::UnboundedSender<String>,
1950 ) -> Result<()> {
1951 let mut body = Vec::new();
1952
1953 body.extend_from_slice(&error_code.to_be_bytes());
1955
1956 let msg_bytes = message.as_bytes();
1958 body.extend_from_slice(&(msg_bytes.len() as u16).to_be_bytes());
1959 body.extend_from_slice(msg_bytes);
1960
1961 let response = Envelope {
1962 version: Version::V4,
1963 direction: Direction::Response,
1964 flags: Flags::empty(),
1965 stream_id: stream_id,
1966 opcode: Opcode::Error,
1967 body,
1968 tracing_id: None,
1969 warnings: vec![],
1970 };
1971
1972 let bytes = response.encode_with(Compression::None)?;
1973 out.write(&bytes).await?;
1974
1975 Log::new(Some(status_tx)).trace(format!(
1976 "Cassandra → ERROR 0x{:04X} - {}",
1977 error_code, message
1978 ));
1979
1980 Ok(())
1981 }
1982
1983 async fn handle_auth_response(
1985 &self,
1986 frame: Envelope,
1987 conn_state: &mut CassandraConnectionState,
1988 out: &ConnWrite,
1989 connection_id: ConnectionId,
1990 status_tx: &mpsc::UnboundedSender<String>,
1991 ) -> Result<bool> {
1992 debug!("Handling AUTH_RESPONSE from connection {}", connection_id);
1993
1994 let body = &frame.body;
1996
1997 let (username, password) = self.parse_sasl_plain(body)?;
1999
2000 trace!("AUTH_RESPONSE: username={}", username);
2001
2002 let protocol =
2004 CassandraProtocol::new(connection_id, self.app_state.clone(), status_tx.clone());
2005
2006 let event = Event {
2007 event_type: &CASSANDRA_AUTH_EVENT,
2008 data: json!({
2009 "username": username,
2010 "password": password,
2011 }),
2012 };
2013
2014 let server_id = self.server_id.context("Server ID not set")?;
2015
2016 let execution_result = match call_llm(
2017 &self.llm_client,
2018 &self.app_state,
2019 server_id,
2020 Some(connection_id),
2021 &event,
2022 &protocol,
2023 )
2024 .await
2025 {
2026 Ok(result) => result,
2027 Err(e) => {
2028 self.send_llm_failure_error(
2029 frame.stream_id,
2030 "AUTH_RESPONSE",
2031 &e,
2032 out,
2033 connection_id,
2034 status_tx,
2035 )
2036 .await?;
2037 return Ok(true);
2038 }
2039 };
2040
2041 for message in &execution_result.messages {
2043 Log::new(Some(status_tx)).info(format!("{}", message));
2044 }
2045
2046 for action_result in execution_result.protocol_results {
2048 match action_result {
2049 ActionResult::Custom { name, data } => {
2050 match name.as_str() {
2051 "cassandra_auth_success" => {
2052 conn_state.authenticated = true;
2053 conn_state.username = Some(username);
2054 self.send_auth_success(frame.stream_id, out, status_tx)
2055 .await?;
2056 return Ok(true);
2057 }
2058 "cassandra_error" => {
2059 let error_code =
2060 data.get("error_code")
2061 .and_then(|v| v.as_u64())
2062 .unwrap_or(0x0000) as u32;
2063 let message = data
2064 .get("message")
2065 .and_then(|v| v.as_str())
2066 .unwrap_or("Unknown error");
2067 self.send_model_error(
2068 frame.stream_id,
2069 "AUTH_RESPONSE",
2070 error_code,
2071 message,
2072 out,
2073 connection_id,
2074 status_tx,
2075 )
2076 .await?;
2077 return Ok(false); }
2079 _ => {}
2080 }
2081 }
2082 ActionResult::CloseConnection => {
2083 return Ok(false);
2084 }
2085 _ => {
2086 warn!("Unexpected action result for AUTH_RESPONSE");
2087 }
2088 }
2089 }
2090
2091 self.send_error(
2093 frame.stream_id,
2094 0x0100,
2095 "Authentication failed",
2096 out,
2097 status_tx,
2098 )
2099 .await?;
2100 Ok(false)
2101 }
2102
2103 fn parse_sasl_plain(&self, body: &[u8]) -> Result<(String, String)> {
2105 if body.is_empty() {
2106 return Err(anyhow::anyhow!("Empty AUTH_RESPONSE body"));
2107 }
2108
2109 let mut idx = 0;
2111 while idx < body.len() && body[idx] != 0 {
2112 idx += 1;
2113 }
2114 idx += 1; if idx >= body.len() {
2117 return Err(anyhow::anyhow!("Invalid SASL PLAIN format"));
2118 }
2119
2120 let username_start = idx;
2122 while idx < body.len() && body[idx] != 0 {
2123 idx += 1;
2124 }
2125 let username = String::from_utf8_lossy(&body[username_start..idx]).to_string();
2126 idx += 1; if idx >= body.len() {
2129 return Err(anyhow::anyhow!(
2130 "Invalid SASL PLAIN format - missing password"
2131 ));
2132 }
2133
2134 let password = String::from_utf8_lossy(&body[idx..]).to_string();
2136
2137 Ok((username, password))
2138 }
2139
2140 async fn send_authenticate(
2142 &self,
2143 stream_id: i16,
2144 authenticator: &str,
2145 out: &ConnWrite,
2146 status_tx: &mpsc::UnboundedSender<String>,
2147 ) -> Result<()> {
2148 let mut body = Vec::new();
2149
2150 let auth_bytes = authenticator.as_bytes();
2152 body.extend_from_slice(&(auth_bytes.len() as u16).to_be_bytes());
2153 body.extend_from_slice(auth_bytes);
2154
2155 let response = Envelope {
2156 version: Version::V4,
2157 direction: Direction::Response,
2158 flags: Flags::empty(),
2159 stream_id: stream_id,
2160 opcode: Opcode::Authenticate,
2161 body,
2162 tracing_id: None,
2163 warnings: vec![],
2164 };
2165
2166 let bytes = response.encode_with(Compression::None)?;
2167 out.write(&bytes).await?;
2168
2169 Log::new(Some(status_tx)).trace(format!("Cassandra → AUTHENTICATE ({})", authenticator));
2170
2171 Ok(())
2172 }
2173
2174 async fn send_auth_success(
2176 &self,
2177 stream_id: i16,
2178 out: &ConnWrite,
2179 status_tx: &mpsc::UnboundedSender<String>,
2180 ) -> Result<()> {
2181 let body = vec![];
2183
2184 let response = Envelope {
2185 version: Version::V4,
2186 direction: Direction::Response,
2187 flags: Flags::empty(),
2188 stream_id: stream_id,
2189 opcode: Opcode::AuthSuccess,
2190 body,
2191 tracing_id: None,
2192 warnings: vec![],
2193 };
2194
2195 let bytes = response.encode_with(Compression::None)?;
2196 out.write(&bytes).await?;
2197
2198 Log::new(Some(status_tx)).trace("Cassandra → AUTH_SUCCESS");
2199
2200 Ok(())
2201 }
2202}