diff --git a/src/chat.rs b/src/chat.rs index 69007e2..5acd4ec 100644 --- a/src/chat.rs +++ b/src/chat.rs @@ -34,6 +34,11 @@ impl ChatStorage { self.packet_size.store(packet_size, Ordering::SeqCst); *self.messages.write().unwrap() = messages; } + + pub fn append(&self, messages: Vec, packet_size: usize) { + self.packet_size.store(packet_size, Ordering::SeqCst); + self.messages.write().unwrap().append(&mut messages.clone()); + } } @@ -67,7 +72,13 @@ Press enter to close")?; let message = format!("Checking ping... {:X}", random::()); send_message(&mut connect(&ctx.host, ctx.enable_ssl)?, &message)?; loop { - let data = read_messages(&mut connect(&ctx.host, ctx.enable_ssl)?, ctx.max_messages, before, !ctx.enable_ssl).ok().flatten(); + let data = read_messages( + &mut connect(&ctx.host, ctx.enable_ssl)?, + ctx.max_messages, + before, + !ctx.enable_ssl, + ctx.enable_chunked + ).ok().flatten(); if let Some((data, size)) = data { if let Some(last) = data.iter().rev().find(|o| o.contains(&message)) { @@ -95,6 +106,7 @@ pub fn print_console(ctx: Arc, messages: Vec, input: &str) -> R let mut messages = messages .into_iter() .flat_map(|o| string_chunks(&o, width as usize - 1)) + .map(|o| (o.0.white().blink().to_string(), o.1)) .collect::>(); let messages_size = if messages.len() >= height { @@ -478,15 +490,27 @@ fn poll_events(ctx: Arc) -> Result<(), Box> { } pub fn recv_tick(ctx: Arc) -> Result<(), Box> { - match read_messages(&mut connect(&ctx.host, ctx.enable_ssl)?, ctx.max_messages, ctx.messages.packet_size(), !ctx.enable_ssl) { + match read_messages( + &mut connect(&ctx.host, ctx.enable_ssl)?, + ctx.max_messages, + ctx.messages.packet_size(), + !ctx.enable_ssl, + ctx.enable_chunked + ) { Ok(Some((messages, size))) => { let messages: Vec = if ctx.disable_formatting { messages } else { messages.into_iter().flat_map(|o| format_message(ctx.clone(), o)).collect() }; - ctx.messages.update(messages.clone(), size); - print_console(ctx.clone(), messages, &ctx.input.read().unwrap())?; + + if ctx.enable_chunked { + ctx.messages.append(messages.clone(), size); + print_console(ctx.clone(), ctx.messages.messages(), &ctx.input.read().unwrap())?; + } else { + ctx.messages.update(messages.clone(), size); + print_console(ctx.clone(), messages, &ctx.input.read().unwrap())?; + } } Err(e) => { println!("{:?}", e); diff --git a/src/config.rs b/src/config.rs index aded370..f7523bb 100644 --- a/src/config.rs +++ b/src/config.rs @@ -33,6 +33,8 @@ pub struct Config { pub enable_auth: bool, #[serde(default)] pub enable_ssl: bool, + #[serde(default)] + pub enable_chunked: bool, } fn default_max_messages() -> usize { 200 } @@ -75,6 +77,7 @@ pub fn configure(path: PathBuf) -> Config { let disable_ip_hiding = ask_bool("Enable your IP viewing?", false); let enable_auth = ask_bool("Enable auth-mode?", false); let enable_ssl = ask_bool("Enable SSL?", false); + let enable_chunked = ask_bool("Enable chunked reading?", false); let config = Config { host, @@ -85,7 +88,8 @@ pub fn configure(path: PathBuf) -> Config { enable_ip_viewing, disable_ip_hiding, enable_auth, - enable_ssl + enable_ssl, + enable_chunked }; let config_text = serde_yml::to_string(&config).expect("Config save error"); @@ -188,6 +192,10 @@ pub struct Args { /// Enable SSL #[arg(short='S', long)] pub enable_ssl: bool, + + /// Enable chunked reading + #[arg(short='u', long)] + pub enable_chunked: bool, } pub struct Context { @@ -205,6 +213,7 @@ pub struct Context { pub scroll: Arc, pub enable_auth: bool, pub enable_ssl: bool, + pub enable_chunked: bool, } impl Context { @@ -223,7 +232,8 @@ impl Context { enable_ip_viewing: args.enable_users_ip_viewing || config.enable_ip_viewing, scroll: Arc::new(AtomicUsize::new(0)), enable_auth: args.enable_auth || config.enable_auth, - enable_ssl: args.enable_ssl || config.enable_ssl + enable_ssl: args.enable_ssl || config.enable_ssl, + enable_chunked: args.enable_chunked || config.enable_chunked, } } } \ No newline at end of file diff --git a/src/main.rs b/src/main.rs index 7424856..19b613c 100644 --- a/src/main.rs +++ b/src/main.rs @@ -52,7 +52,8 @@ fn main() { &mut stream, ctx.max_messages, 0, - !ctx.enable_ssl + !ctx.enable_ssl, + false ) .ok().flatten() .expect("Error reading messages").0.join("\n") diff --git a/src/proto.rs b/src/proto.rs index 9972a1e..948baf6 100644 --- a/src/proto.rs +++ b/src/proto.rs @@ -66,11 +66,17 @@ fn skip_null(stream: &mut impl Read) -> Result, Box> { } } -pub fn read_messages(stream: &mut (impl Read + Write), max_messages: usize, last_size: usize, skip_null_: bool) -> Result, usize)>, Box> { +pub fn read_messages( + stream: &mut (impl Read + Write), + max_messages: usize, + last_size: usize, + start_null: bool, + chunked: bool +) -> Result, usize)>, Box> { stream.write_all(&[0x00])?; let packet_size = { - let data = if skip_null_ { + let data = if start_null { let mut data = skip_null(stream)?; loop { @@ -99,20 +105,25 @@ pub fn read_messages(stream: &mut (impl Read + Write), max_messages: usize, last return Ok(None); } - stream.write_all(&[0x01])?; + let to_read = if !chunked || last_size == 0 { + stream.write_all(&[0x01])?; + packet_size + } else { + stream.write_all(format!("\x02{}", last_size).as_bytes())?; + packet_size - last_size + }; let packet_data = { - let data = if skip_null_ { + let data = if start_null { let mut data = skip_null(stream)?; - while data.len() < packet_size { - let mut buf = vec![0; packet_size - data.len()]; - let read_bytes = stream.read(&mut buf)?; - buf.truncate(read_bytes); + while data.len() < to_read { + let mut buf = vec![0; to_read - data.len()]; + stream.read_exact(&mut buf)?; data.append(&mut buf); } data } else { - let mut data = vec![0; packet_size]; + let mut data = vec![0; to_read]; stream.read_exact(&mut data)?; data };