src/main.rs

   1 use std::cmp::min;
   2 use std::collections::HashMap;
   3 use std::ffi::{OsStr, OsString};
   4 use std::io::{ErrorKind, Write};
   5 use std::os::unix::io::{AsRawFd, FromRawFd};
   6 use std::os::unix::process::CommandExt;
   7 use std::process::Command;
   8 use std::time::{Duration, Instant};
   9
  10 use anyhow::{bail, format_err, Result};
  11 use clap::{App, AppSettings, Arg};
  12 use mio::net::{TcpListener, TcpStream};
  13 use mio::unix::SourceFd;
  14 use mio::{Events, Interest, Poll, Token};
  15
  16 use proxmox_io::ByteBuffer;
  17 use proxmox_lang::error::io_err_other;
  18 use proxmox_sys::linux::pty::{make_controlling_terminal, PTY};
  19
  20 const MSG_TYPE_DATA: u8 = 0;
  21 const MSG_TYPE_RESIZE: u8 = 1;
  22 //const MSG_TYPE_PING: u8 = 2;
  23
  24 fn remove_number(buf: &mut ByteBuffer) -> Option<usize> {
  25     loop {
  26         if let Some(pos) = &buf.iter().position(|&x| x == b':') {
  27             let data = buf.remove_data(*pos);
  28             buf.consume(1); // the ':'
  29             let len = match std::str::from_utf8(&data) {
  30                 Ok(lenstring) => match lenstring.parse() {
  31                     Ok(len) => len,
  32                     Err(err) => {
  33                         eprintln!("error parsing number: '{}'", err);
  34                         break;
  35                     }
  36                 },
  37                 Err(err) => {
  38                     eprintln!("error parsing number: '{}'", err);
  39                     break;
  40                 }
  41             };
  42             return Some(len);
  43         } else if buf.len() > 20 {
  44             buf.consume(20);
  45         } else {
  46             break;
  47         }
  48     }
  49     None
  50 }
  51
  52 fn process_queue(buf: &mut ByteBuffer, pty: &mut PTY) -> Option<usize> {
  53     if buf.is_empty() {
  54         return None;
  55     }
  56
  57     loop {
  58         if buf.len() < 2 {
  59             break;
  60         }
  61
  62         let msgtype = buf[0] - b'0';
  63
  64         if msgtype == MSG_TYPE_DATA {
  65             buf.consume(2);
  66             if let Some(len) = remove_number(buf) {
  67                 return Some(len);
  68             }
  69         } else if msgtype == MSG_TYPE_RESIZE {
  70             buf.consume(2);
  71             if let Some(cols) = remove_number(buf) {
  72                 if let Some(rows) = remove_number(buf) {
  73                     pty.set_size(cols as u16, rows as u16).ok()?;
  74                 }
  75             }
  76         // ignore incomplete messages
  77         } else {
  78             buf.consume(1);
  79             // ignore invalid or ping (msgtype 2)
  80         }
  81     }
  82
  83     None
  84 }
  85
  86 type TicketResult = Result<(Box<[u8]>, Box<[u8]>)>;
  87
  88 /// Reads from the stream and returns the first line and the rest
  89 fn read_ticket_line(
  90     stream: &mut TcpStream,
  91     buf: &mut ByteBuffer,
  92     timeout: Duration,
  93 ) -> TicketResult {
  94     let mut poll = Poll::new()?;
  95     poll.registry()
  96         .register(stream, Token(0), Interest::READABLE)?;
  97     let mut events = Events::with_capacity(1);
  98
  99     let now = Instant::now();
 100     let mut elapsed = Duration::new(0, 0);
 101
 102     loop {
 103         poll.poll(&mut events, Some(timeout - elapsed))?;
 104         if !events.is_empty() {
 105             match buf.read_from(stream) {
 106                 Ok(n) => {
 107                     if n == 0 {
 108                         bail!("connection closed before authentication");
 109                     }
 110                 }
 111                 Err(err) if err.kind() == ErrorKind::WouldBlock => {}
 112                 Err(err) => return Err(err.into()),
 113             }
 114
 115             if buf[..].contains(&b'\n') {
 116                 break;
 117             }
 118
 119             if buf.is_full() {
 120                 bail!("authentication data is incomplete: {:?}", &buf[..]);
 121             }
 122         }
 123
 124         elapsed = now.elapsed();
 125         if elapsed > timeout {
 126             bail!("timed out");
 127         }
 128     }
 129
 130     let newline_idx = &buf[..].iter().position(|&x| x == b'\n').unwrap();
 131
 132     let line = buf.remove_data(*newline_idx);
 133     buf.consume(1); // discard newline
 134
 135     match line.iter().position(|&b| b == b':') {
 136         Some(pos) => {
 137             let (username, ticket) = line.split_at(pos);
 138             Ok((username.into(), ticket[1..].into()))
 139         }
 140         None => bail!("authentication data is invalid"),
 141     }
 142 }
 143
 144 fn authenticate(
 145     username: &[u8],
 146     ticket: &[u8],
 147     path: &str,
 148     perm: Option<&str>,
 149     authport: u16,
 150     port: Option<u16>,
 151 ) -> Result<()> {
 152     let mut post_fields: Vec<(&str, &str)> = Vec::with_capacity(5);
 153     post_fields.push(("username", std::str::from_utf8(username)?));
 154     post_fields.push(("password", std::str::from_utf8(ticket)?));
 155     post_fields.push(("path", path));
 156     if let Some(perm) = perm {
 157         post_fields.push(("privs", perm));
 158     }
 159     let port_str;
 160     if let Some(port) = port {
 161         port_str = port.to_string();
 162         post_fields.push(("port", &port_str));
 163     }
 164
 165     let url = format!("http://localhost:{}/api2/json/access/ticket", authport);
 166
 167     match ureq::post(&url).send_form(&post_fields[..]) {
 168         Ok(res) if res.status() == 200 => Ok(()),
 169         Ok(res) | Err(ureq::Error::Status(_, res)) => {
 170             let code = res.status();
 171             bail!("invalid authentication - {} {}", code, res.status_text())
 172         }
 173         Err(err) => bail!("authentication request failed - {}", err),
 174     }
 175 }
 176
 177 fn listen_and_accept(
 178     hostname: &str,
 179     port: u64,
 180     port_as_fd: bool,
 181     timeout: Duration,
 182 ) -> Result<(TcpStream, u16)> {
 183     let listener = if port_as_fd {
 184         unsafe { std::net::TcpListener::from_raw_fd(port as i32) }
 185     } else {
 186         std::net::TcpListener::bind((hostname, port as u16))?
 187     };
 188     let port = listener.local_addr()?.port();
 189     let mut listener = TcpListener::from_std(listener);
 190     let mut poll = Poll::new()?;
 191
 192     poll.registry()
 193         .register(&mut listener, Token(0), Interest::READABLE)?;
 194
 195     let mut events = Events::with_capacity(1);
 196
 197     let now = Instant::now();
 198     let mut elapsed = Duration::new(0, 0);
 199
 200     loop {
 201         poll.poll(&mut events, Some(timeout - elapsed))?;
 202         if !events.is_empty() {
 203             let (stream, client) = listener.accept()?;
 204             println!("client connection: {:?}", client);
 205             return Ok((stream, port));
 206         }
 207
 208         elapsed = now.elapsed();
 209         if elapsed > timeout {
 210             bail!("timed out");
 211         }
 212     }
 213 }
 214
 215 fn run_pty(cmd: &OsStr, params: clap::OsValues) -> Result<PTY> {
 216     let (mut pty, secondary_name) = PTY::new().map_err(io_err_other)?;
 217
 218     let mut filtered_env: HashMap<OsString, OsString> = std::env::vars_os()
 219         .filter(|&(ref k, _)| {
 220             k == "PATH"
 221                 || k == "USER"
 222                 || k == "HOME"
 223                 || k == "LANG"
 224                 || k == "LANGUAGE"
 225                 || k.to_string_lossy().starts_with("LC_")
 226         })
 227         .collect();
 228     filtered_env.insert("TERM".into(), "xterm-256color".into());
 229
 230     let mut command = Command::new(cmd);
 231
 232     command.args(params).env_clear().envs(&filtered_env);
 233
 234     unsafe {
 235         command.pre_exec(move || {
 236             make_controlling_terminal(&secondary_name).map_err(io_err_other)?;
 237             Ok(())
 238         });
 239     }
 240
 241     command.spawn()?;
 242
 243     pty.set_size(80, 20)?;
 244     Ok(pty)
 245 }
 246
 247 const TCP: Token = Token(0);
 248 const PTY: Token = Token(1);
 249
 250 fn do_main() -> Result<()> {
 251     let matches = App::new("termproxy")
 252         .setting(AppSettings::TrailingVarArg)
 253         .arg(Arg::with_name("port").takes_value(true).required(true))
 254         .arg(
 255             Arg::with_name("authport")
 256                 .takes_value(true)
 257                 .long("authport"),
 258         )
 259         .arg(Arg::with_name("use-port-as-fd").long("port-as-fd"))
 260         .arg(
 261             Arg::with_name("path")
 262                 .takes_value(true)
 263                 .long("path")
 264                 .required(true),
 265         )
 266         .arg(Arg::with_name("perm").takes_value(true).long("perm"))
 267         .arg(
 268             Arg::with_name("cmd")
 269                 .allow_invalid_utf8(true)
 270                 .multiple(true)
 271                 .required(true),
 272         )
 273         .get_matches();
 274
 275     let port: u64 = matches
 276         .value_of("port")
 277         .unwrap()
 278         .parse()
 279         .map_err(io_err_other)?;
 280     let path = matches.value_of("path").unwrap();
 281     let perm: Option<&str> = matches.value_of("perm");
 282     let mut cmdparams = matches.values_of_os("cmd").unwrap();
 283     let cmd = cmdparams.next().unwrap();
 284     let authport: u16 = matches
 285         .value_of("authport")
 286         .unwrap_or("85")
 287         .parse()
 288         .map_err(io_err_other)?;
 289     let mut pty_buf = ByteBuffer::new();
 290     let mut tcp_buf = ByteBuffer::new();
 291
 292     let use_port_as_fd = matches.is_present("use-port-as-fd");
 293
 294     if use_port_as_fd && port > u16::MAX as u64 {
 295         return Err(format_err!("port too big"));
 296     } else if port > i32::MAX as u64 {
 297         return Err(format_err!("Invalid FD number"));
 298     }
 299
 300     let (mut tcp_handle, port) =
 301         listen_and_accept("localhost", port, use_port_as_fd, Duration::new(10, 0))
 302             .map_err(|err| format_err!("failed waiting for client: {}", err))?;
 303
 304     let (username, ticket) = read_ticket_line(&mut tcp_handle, &mut pty_buf, Duration::new(10, 0))
 305         .map_err(|err| format_err!("failed reading ticket: {}", err))?;
 306     let port = if use_port_as_fd { Some(port) } else { None };
 307     authenticate(&username, &ticket, path, perm, authport, port)?;
 308     tcp_handle.write_all(b"OK").expect("error writing response");
 309
 310     let mut poll = Poll::new()?;
 311     let mut events = Events::with_capacity(128);
 312
 313     let mut pty = run_pty(cmd, cmdparams)?;
 314
 315     poll.registry().register(
 316         &mut tcp_handle,
 317         TCP,
 318         Interest::READABLE | Interest::WRITABLE,
 319     )?;
 320     poll.registry().register(
 321         &mut SourceFd(&pty.as_raw_fd()),
 322         PTY,
 323         Interest::READABLE | Interest::WRITABLE,
 324     )?;
 325
 326     let mut tcp_writable = true;
 327     let mut pty_writable = true;
 328     let mut tcp_readable = true;
 329     let mut pty_readable = true;
 330     let mut remaining = 0;
 331     let mut finished = false;
 332
 333     while !finished {
 334         if tcp_readable && !pty_buf.is_full() || pty_readable && !tcp_buf.is_full() {
 335             poll.poll(&mut events, Some(Duration::new(0, 0)))?;
 336         } else {
 337             poll.poll(&mut events, None)?;
 338         }
 339
 340         for event in &events {
 341             let writable = event.is_writable();
 342             let readable = event.is_readable();
 343             if event.is_read_closed() {
 344                 finished = true;
 345             }
 346             match event.token() {
 347                 TCP => {
 348                     if readable {
 349                         tcp_readable = true;
 350                     }
 351                     if writable {
 352                         tcp_writable = true;
 353                     }
 354                 }
 355                 PTY => {
 356                     if readable {
 357                         pty_readable = true;
 358                     }
 359                     if writable {
 360                         pty_writable = true;
 361                     }
 362                 }
 363                 _ => unreachable!(),
 364             }
 365         }
 366
 367         while tcp_readable && !pty_buf.is_full() {
 368             let bytes = match pty_buf.read_from(&mut tcp_handle) {
 369                 Ok(bytes) => bytes,
 370                 Err(err) if err.kind() == ErrorKind::WouldBlock => {
 371                     tcp_readable = false;
 372                     break;
 373                 }
 374                 Err(err) => {
 375                     if !finished {
 376                         return Err(format_err!("error reading from tcp: {}", err));
 377                     }
 378                     break;
 379                 }
 380             };
 381             if bytes == 0 {
 382                 finished = true;
 383                 break;
 384             }
 385         }
 386
 387         while pty_readable && !tcp_buf.is_full() {
 388             let bytes = match tcp_buf.read_from(&mut pty) {
 389                 Ok(bytes) => bytes,
 390                 Err(err) if err.kind() == ErrorKind::WouldBlock => {
 391                     pty_readable = false;
 392                     break;
 393                 }
 394                 Err(err) => {
 395                     if !finished {
 396                         return Err(format_err!("error reading from pty: {}", err));
 397                     }
 398                     break;
 399                 }
 400             };
 401             if bytes == 0 {
 402                 finished = true;
 403                 break;
 404             }
 405         }
 406
 407         while !tcp_buf.is_empty() && tcp_writable {
 408             let bytes = match tcp_handle.write(&tcp_buf[..]) {
 409                 Ok(bytes) => bytes,
 410                 Err(err) if err.kind() == ErrorKind::WouldBlock => {
 411                     tcp_writable = false;
 412                     break;
 413                 }
 414                 Err(err) => {
 415                     if !finished {
 416                         return Err(format_err!("error writing to tcp : {}", err));
 417                     }
 418                     break;
 419                 }
 420             };
 421             tcp_buf.consume(bytes);
 422         }
 423
 424         while !pty_buf.is_empty() && pty_writable {
 425             if remaining == 0 {
 426                 remaining = match process_queue(&mut pty_buf, &mut pty) {
 427                     Some(val) => val,
 428                     None => break,
 429                 };
 430             }
 431             let len = min(remaining, pty_buf.len());
 432             let bytes = match pty.write(&pty_buf[..len]) {
 433                 Ok(bytes) => bytes,
 434                 Err(err) if err.kind() == ErrorKind::WouldBlock => {
 435                     pty_writable = false;
 436                     break;
 437                 }
 438                 Err(err) => {
 439                     if !finished {
 440                         return Err(format_err!("error writing to pty : {}", err));
 441                     }
 442                     break;
 443                 }
 444             };
 445             remaining -= bytes;
 446             pty_buf.consume(bytes);
 447         }
 448     }
 449
 450     Ok(())
 451 }
 452
 453 fn main() {
 454     std::process::exit(match do_main() {
 455         Ok(_) => 0,
 456         Err(err) => {
 457             eprintln!("{}", err);
 458             1
 459         }
 460     });
 461 }