rubot/src/irc/connection.rs

use std::io::{self, BufRead, BufReader, Read, Write};
use std::net::TcpStream;
use std::time::{Duration, Instant};

use native_tls::TlsConnector;

use super::Message;

const MAX_LINE: usize = 500;
const SEND_THROTTLE: Duration = Duration::from_millis(300);

trait Stream: Read + Write {}
impl<T: Read + Write> Stream for T {}

pub struct Connection {
    inner: BufReader<Box<dyn Stream>>,
    log_prefix: String,
    verbose: bool,
    throttle: Duration,
    last_send: Option<Instant>,
}

impl Connection {
    pub fn connect(host: &str, port: u16, use_tls: bool, verbose: bool) -> io::Result<Connection> {
        let tcp = TcpStream::connect((host, port))?;
        let stream: Box<dyn Stream> = if use_tls {
            let connector =
                TlsConnector::new().map_err(|e| io::Error::other(format!("TLS init: {e}")))?;
            let tls = connector
                .connect(host, tcp)
                .map_err(|e| io::Error::other(format!("TLS handshake: {e}")))?;
            Box::new(tls)
        } else {
            Box::new(tcp)
        };
        Ok(Connection {
            inner: BufReader::new(stream),
            log_prefix: String::new(),
            verbose,
            throttle: SEND_THROTTLE,
            last_send: None,
        })
    }

    pub fn from_stream<S: Read + Write + 'static>(stream: S) -> Connection {
        let boxed: Box<dyn Stream> = Box::new(stream);
        Connection {
            inner: BufReader::new(boxed),
            log_prefix: String::new(),
            verbose: false,
            throttle: Duration::ZERO,
            last_send: None,
        }
    }

    pub fn set_label(&mut self, label: &str) {
        self.log_prefix = if label.is_empty() {
            String::new()
        } else {
            format!("[{label}] ")
        };
    }

    pub fn read_message(&mut self) -> io::Result<Option<Message>> {
        loop {
            let mut line = String::new();
            if self.inner.read_line(&mut line)? == 0 {
                return Ok(None);
            }
            let trimmed = line.trim_end_matches(['\r', '\n']);
            if trimmed.is_empty() {
                continue;
            }
            if self.verbose {
                eprintln!("{}<< {trimmed}", self.log_prefix);
            }
            if let Some(msg) = Message::parse(trimmed) {
                return Ok(Some(msg));
            }
        }
    }

    pub fn send_raw(&mut self, line: &str) -> io::Result<()> {
        if !self.throttle.is_zero() {
            if let Some(last) = self.last_send {
                let elapsed = last.elapsed();
                if elapsed < self.throttle {
                    std::thread::sleep(self.throttle - elapsed);
                }
            }
            self.last_send = Some(Instant::now());
        }
        let clean = line.replace(['\r', '\n'], " ");
        let clean = truncate_to(&clean, MAX_LINE);
        if self.verbose {
            eprintln!("{}>> {clean}", self.log_prefix);
        }
        let writer = self.inner.get_mut();
        write!(writer, "{clean}\r\n")?;
        writer.flush()
    }

    pub fn pass(&mut self, password: &str) -> io::Result<()> {
        self.send_raw(&format!("PASS {password}"))
    }

    pub fn nick(&mut self, nick: &str) -> io::Result<()> {
        self.send_raw(&format!("NICK {nick}"))
    }

    pub fn user(&mut self, user: &str, realname: &str) -> io::Result<()> {
        self.send_raw(&format!("USER {user} 0 * :{realname}"))
    }

    pub fn join(&mut self, channel: &str) -> io::Result<()> {
        self.send_raw(&format!("JOIN {channel}"))
    }

    pub fn pong(&mut self, token: &str) -> io::Result<()> {
        self.send_raw(&format!("PONG :{token}"))
    }

    pub fn privmsg(&mut self, target: &str, text: &str) -> io::Result<()> {
        self.send_raw(&format!("PRIVMSG {target} :{text}"))
    }

    pub fn notice(&mut self, target: &str, text: &str) -> io::Result<()> {
        self.send_raw(&format!("NOTICE {target} :{text}"))
    }

    pub fn quit(&mut self, reason: &str) -> io::Result<()> {
        self.send_raw(&format!("QUIT :{reason}"))
    }

    pub fn cap_ls(&mut self) -> io::Result<()> {
        self.send_raw("CAP LS 302")
    }

    pub fn cap_req(&mut self, caps: &str) -> io::Result<()> {
        self.send_raw(&format!("CAP REQ :{caps}"))
    }

    pub fn cap_end(&mut self) -> io::Result<()> {
        self.send_raw("CAP END")
    }

    pub fn authenticate(&mut self, payload: &str) -> io::Result<()> {
        self.send_raw(&format!("AUTHENTICATE {payload}"))
    }
}

fn truncate_to(s: &str, max: usize) -> &str {
    if s.len() <= max {
        return s;
    }
    let mut end = max;
    while end > 0 && !s.is_char_boundary(end) {
        end -= 1;
    }
    &s[..end]
}