ratel_rust/

shell.rs

use anyhow::{Context, Result};
use std::sync::Arc;
use tokio::io::AsyncWriteExt;
use tokio::sync::Mutex;

use crate::model::AuthInfo;
use crate::network::{Connection, NetType, TcpConnection, WebSocketConnection};
use crate::util::CLEAN_LINE;
use tokio::sync::mpsc;

const IS_START: &str = "INTERACTIVE_SIGNAL_START";
const IS_STOP: &str = "INTERACTIVE_SIGNAL_STOP";

pub enum ShellConnection {
    Tcp(TcpConnection),
    WebSocket(WebSocketConnection),
}

impl ShellConnection {
    async fn send(&mut self, data: &[u8]) -> Result<()> {
        match self {
            ShellConnection::Tcp(conn) => conn.send(data).await,
            ShellConnection::WebSocket(conn) => conn.send(data).await,
        }
    }

    async fn receive(&mut self) -> Result<Option<Vec<u8>>> {
        match self {
            ShellConnection::Tcp(conn) => conn.receive().await,
            ShellConnection::WebSocket(conn) => conn.receive().await,
        }
    }
}

pub struct Shell {
    addr: String,
    name: String,
    id: i64,
    score: i64,
    conn: Option<ShellConnection>,
    is_active: Arc<Mutex<bool>>,
}

impl Shell {
    pub fn new(addr: String, name: String) -> Self {
        Self {
            addr,
            name,
            id: std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .unwrap()
                .as_nanos() as i64,
            score: 100,
            conn: None,
            is_active: Arc::new(Mutex::new(false)),
        }
    }

    fn detect_net_type(&self) -> NetType {
        if self.addr.ends_with("9998") {
            NetType::WebSocket
        } else {
            NetType::Tcp
        }
    }

    async fn connect(&mut self) -> Result<()> {
        let net_type = self.detect_net_type();

        self.conn = Some(match net_type {
            NetType::Tcp => {
                let conn = TcpConnection::connect(&self.addr).await?;
                ShellConnection::Tcp(conn)
            }
            NetType::WebSocket => {
                let conn = WebSocketConnection::connect(&self.addr).await?;
                ShellConnection::WebSocket(conn)
            }
        });

        Ok(())
    }

    async fn auth(&mut self) -> Result<()> {
        let auth_info = AuthInfo {
            id: self.id,
            name: self.name.clone(),
            score: self.score,
        };

        // Send auth info directly as JSON with 4-byte big-endian length prefix
        let json_data = serde_json::to_vec(&auth_info)?;
        let len_bytes = (json_data.len() as u32).to_be_bytes();

        let mut data = Vec::new();
        data.extend_from_slice(&len_bytes);
        data.extend_from_slice(&json_data);

        if let Some(conn) = self.conn.as_mut() {
            conn.send(&data).await?;
        }

        Ok(())
    }

    fn get_prompt(&self) -> String {
        format!(
            "{}[{}@ratel {}]# ",
            CLEAN_LINE,
            self.name.to_lowercase(),
            "~"
        )
    }

    async fn print(text: &str) {
        print!("{}", text);
        use std::io::{self, Write};
        io::stdout().flush().ok();
    }

    pub async fn start(&mut self) -> Result<()> {
        self.connect().await.context("Failed to connect")?;
        self.auth().await.context("Failed to authenticate")?;

        let is_active = self.is_active.clone();
        let conn = Arc::new(Mutex::new(self.conn.take()));
        let prompt = self.get_prompt();

        // Create channel for stdin -> network communication
        let (tx, mut rx) = mpsc::unbounded_channel::<Vec<u8>>();

        let is_active_stdin = is_active.clone();
        let prompt_stdin = prompt.clone();

        // Spawn stdin handler
        tokio::spawn(async move {
            use tokio::io::{AsyncBufReadExt, BufReader};

            let stdin = tokio::io::stdin();
            let mut reader = BufReader::new(stdin);

            loop {
                let mut line = String::new();
                match reader.read_line(&mut line).await {
                    Ok(0) => break,
                    Ok(_) => {
                        let line = line.trim();
                        let active = *is_active_stdin.lock().await;

                        if active {
                            print!("{}{}", CLEAN_LINE, prompt_stdin);
                            tokio::io::stdout().flush().await.ok();

                            // Create packet with 4-byte big-endian length prefix
                            let data = line.as_bytes();
                            let len_bytes = (data.len() as u32).to_be_bytes();

                            let mut packet = Vec::new();
                            packet.extend_from_slice(&len_bytes);
                            packet.extend_from_slice(data);

                            // Send to channel
                            let _ = tx.send(packet);
                        }
                    }
                    Err(_) => break,
                }
            }
        });

        // Main network loop: handle both receiving and sending
        if let Some(c) = conn.lock().await.as_mut() {
            loop {
                tokio::select! {
                    // Handle incoming messages from stdin
                    Some(packet) = rx.recv() => {
                        if let Err(e) = c.send(&packet).await {
                            eprintln!("Send error: {:?}", e);
                            break;
                        }
                    }
                    // Handle incoming messages from server
                    result = c.receive() => {
                        match result {
                            Ok(Some(data)) => {
                                let text = String::from_utf8_lossy(&data);

                                let mut active = is_active.lock().await;

                                if text == IS_START {
                                    if !*active {
                                        Self::print(&prompt).await;
                                    }
                                    *active = true;
                                } else if text == IS_STOP {
                                    if *active {
                                        Self::print(CLEAN_LINE).await;
                                    }
                                    *active = false;
                                } else if *active {
                                    let output = format!("{}{}{}", CLEAN_LINE, text, prompt);
                                    Self::print(&output).await;
                                } else {
                                    Self::print(&text).await;
                                }
                            }
                            Ok(None) => {
                                break;
                            }
                            Err(e) => {
                                eprintln!("Error receiving: {:?}", e);
                                break;
                            }
                        }
                    }
                }
            }
        }

        Ok(())
    }
}