├── .envrc ├── .gitignore ├── default.nix ├── .github └── workflows │ ├── ci.yml │ └── release.yml ├── flake.nix ├── LICENSE.txt ├── Cargo.toml ├── flake.lock ├── rfbproxy.nix ├── src ├── messages │ ├── mod.rs │ ├── io.rs │ ├── client.rs │ └── server.rs ├── audio │ ├── mp3.rs │ ├── mod.rs │ └── opus.rs ├── client.proto ├── main.rs ├── rfb.rs └── auth.rs └── README.md /.envrc: -------------------------------------------------------------------------------- 1 | use flake 2 | dotenv_if_exists 3 | source_env_if_exists .local.envrc 4 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | /target/ 2 | 3 | /docker/ 4 | /pulse/ 5 | test.py 6 | *.log 7 | .idea 8 | /.direnv/ 9 | 10 | # nix build creates a "result" symlink to the nix store output 11 | result 12 | -------------------------------------------------------------------------------- /default.nix: -------------------------------------------------------------------------------- 1 | # flake-compat makes Nix flakes compatible with the old Nix cli commands, like nix-build and nix-shell. 2 | (import ( 3 | fetchTarball { 4 | url = "https://github.com/edolstra/flake-compat/archive/99f1c2157fba4bfe6211a321fd0ee43199025dbf.tar.gz"; 5 | sha256 = "0x2jn3vrawwv9xp15674wjz9pixwjyj3j771izayl962zziivbx2"; } 6 | ) { 7 | src = ./.; 8 | }).defaultNix 9 | -------------------------------------------------------------------------------- /.github/workflows/ci.yml: -------------------------------------------------------------------------------- 1 | name: CI 2 | 3 | on: 4 | pull_request: {} 5 | push: 6 | branches: 7 | - main 8 | 9 | jobs: 10 | ci: 11 | runs-on: ubuntu-24.04 12 | 13 | steps: 14 | - uses: actions/checkout@v2 15 | - uses: actions-rust-lang/setup-rust-toolchain@v1 16 | with: 17 | rust-version: stable 18 | components: clippy 19 | 20 | - name: Install dependencies 21 | run: sudo apt-get update && sudo apt-get install -y libpulse-dev libopus-dev libmp3lame-dev 22 | 23 | - name: build 24 | run: cargo build 25 | - name: lint 26 | run: cargo clippy -- -D warnings 27 | - name: rustfmt 28 | run: cargo fmt -- --check 29 | - name: test 30 | run: cargo test 31 | - name: doc 32 | run: cargo doc 33 | 34 | nix: 35 | runs-on: ubuntu-24.04 36 | steps: 37 | - uses: actions/checkout@v2 38 | - uses: actions/cache@v4 39 | with: 40 | path: /nix 41 | key: ${{ runner.os }}-cargo-${{ hashFiles('**/Cargo.lock') }} 42 | - uses: DeterminateSystems/nix-installer-action@main 43 | -------------------------------------------------------------------------------- /flake.nix: -------------------------------------------------------------------------------- 1 | { 2 | description = "An RFB proxy that enables WebSockets and audio"; 3 | 4 | inputs = { 5 | nixpkgs.url = "github:NixOS/nixpkgs/nixos-24.11"; 6 | flake-utils.url = "github:numtide/flake-utils"; 7 | }; 8 | 9 | outputs = { self, nixpkgs, flake-utils }: 10 | flake-utils.lib.eachDefaultSystem 11 | (system: 12 | let 13 | pkgs = nixpkgs.legacyPackages.${system}; 14 | rfbproxy = pkgs.callPackage ./rfbproxy.nix { 15 | rev = if self ? rev then "0.0.0-${builtins.substring 0 7 self.rev}" else "0.0.0-dirty"; 16 | }; 17 | in 18 | { 19 | defaultPackage = rfbproxy; 20 | packages = { 21 | inherit rfbproxy; 22 | }; 23 | devShell = pkgs.mkShell { 24 | packages = [ 25 | pkgs.cargo 26 | pkgs.crate2nix 27 | pkgs.lame 28 | pkgs.libiconv 29 | pkgs.libpulseaudio 30 | pkgs.openssl 31 | pkgs.pkg-config 32 | pkgs.rustfmt 33 | ]; 34 | }; 35 | } 36 | ); 37 | } 38 | -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 Replit 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "rfbproxy" 3 | version = "0.1.1" 4 | authors = ["Luis Héctor Chávez "] 5 | edition = "2021" 6 | 7 | [dependencies] 8 | anyhow = "1.0" 9 | base64 = "0.13.0" 10 | byte-slice-cast = "1.0.0" 11 | bytes = "1.0.1" 12 | chrono = "0.4" 13 | cipher = "0.4" 14 | clap = "2.33" 15 | des = "0.8" 16 | env_logger = "0.11" 17 | futures = "0.3.12" 18 | http = "0.2" 19 | hyper = { version = "0.14.12", features = ["server", "tcp", "http1", "http2"] } 20 | hyper-staticfile = "0.6" 21 | hyper-tungstenite = "0.11" 22 | lame-sys = "0.1.2" 23 | log = "*" 24 | opus = "0.2.1" 25 | paseto = { version = "2.0.2+1.0.3", features = ["v2"] } 26 | path-clean = "0.1" 27 | prost = "0.8" 28 | prost-types = "0.8" 29 | psimple = { package = "libpulse-simple-binding", version = "2.20.1" } 30 | pulse = { package = "libpulse-binding", version = "2.20.0" } 31 | rand = "0.8.2" 32 | ring = "0.16" 33 | serde_json = "1.0.64" 34 | tokio = { version = "1.18", features = ["full"] } 35 | tokio-tungstenite = "0.20" 36 | 37 | [dev-dependencies] 38 | tempfile = "3" 39 | matroska = "0.5.5" 40 | tokio-test = "0.4.0" 41 | 42 | [build-dependencies] 43 | prost-build = { version = "0.7" } 44 | -------------------------------------------------------------------------------- /flake.lock: -------------------------------------------------------------------------------- 1 | { 2 | "nodes": { 3 | "flake-utils": { 4 | "locked": { 5 | "lastModified": 1623875721, 6 | "narHash": "sha256-A8BU7bjS5GirpAUv4QA+QnJ4CceLHkcXdRp4xITDB0s=", 7 | "owner": "numtide", 8 | "repo": "flake-utils", 9 | "rev": "f7e004a55b120c02ecb6219596820fcd32ca8772", 10 | "type": "github" 11 | }, 12 | "original": { 13 | "owner": "numtide", 14 | "repo": "flake-utils", 15 | "type": "github" 16 | } 17 | }, 18 | "nixpkgs": { 19 | "locked": { 20 | "lastModified": 1739758141, 21 | "narHash": "sha256-uq6A2L7o1/tR6VfmYhZWoVAwb3gTy7j4Jx30MIrH0rE=", 22 | "owner": "NixOS", 23 | "repo": "nixpkgs", 24 | "rev": "c618e28f70257593de75a7044438efc1c1fc0791", 25 | "type": "github" 26 | }, 27 | "original": { 28 | "owner": "NixOS", 29 | "ref": "nixos-24.11", 30 | "repo": "nixpkgs", 31 | "type": "github" 32 | } 33 | }, 34 | "root": { 35 | "inputs": { 36 | "flake-utils": "flake-utils", 37 | "nixpkgs": "nixpkgs" 38 | } 39 | } 40 | }, 41 | "root": "root", 42 | "version": 7 43 | } 44 | -------------------------------------------------------------------------------- /rfbproxy.nix: -------------------------------------------------------------------------------- 1 | { stdenv, openssl, libpulseaudio, protobuf, lame, libopus, git, runCommand 2 | , copyPathToStore, rev, pkg-config, lib, defaultCrateOverrides, buildRustCrate 3 | , buildPackages, fetchurl }@pkgs: 4 | let 5 | generatedBuild = import ./Cargo.nix { 6 | inherit pkgs; 7 | buildRustCrateForPkgs = pkgs: 8 | pkgs.buildRustCrate.override { 9 | defaultCrateOverrides = pkgs.defaultCrateOverrides // { 10 | "opus-sys" = attrs: { 11 | nativeBuildInputs = [ pkg-config ]; 12 | buildInputs = [ libopus ]; 13 | }; 14 | "libpulse-sys" = attrs: { 15 | nativeBuildInputs = [ pkg-config ]; 16 | buildInputs = [ libpulseaudio ]; 17 | }; 18 | "libpulse-simple-sys" = attrs: { 19 | nativeBuildInputs = [ pkg-config ]; 20 | buildInputs = [ libpulseaudio ]; 21 | }; 22 | rfbproxy = attrs: { 23 | buildInputs = [ openssl protobuf lame ]; 24 | nativeBuildInputs = [ pkg-config ]; 25 | 26 | # needed for internal protobuf c wrapper library 27 | PROTOC = "${protobuf}/bin/protoc"; 28 | PROTOC_INCLUDE = "${protobuf}/include"; 29 | }; 30 | }; 31 | }; 32 | }; 33 | 34 | crate2nix = generatedBuild.rootCrate.build; 35 | 36 | in stdenv.mkDerivation { 37 | pname = "rfbproxy"; 38 | version = rev; 39 | 40 | src = crate2nix; 41 | 42 | installPhase = '' 43 | cp -r ${crate2nix} $out 44 | ''; 45 | } 46 | -------------------------------------------------------------------------------- /.github/workflows/release.yml: -------------------------------------------------------------------------------- 1 | name: Release 2 | 3 | on: 4 | push: 5 | branches: 6 | - main 7 | tags: 8 | - v* 9 | 10 | jobs: 11 | 12 | release: 13 | # Should match the version reported by /etc/lsb-release on Repls. 14 | runs-on: ubuntu-24.04 15 | 16 | steps: 17 | - uses: actions/checkout@v2 18 | with: 19 | submodules: recursive 20 | fetch-depth: 0 21 | 22 | - name: Install dependencies 23 | run: sudo apt-get install -y libpulse-dev libopus-dev libmp3lame-dev xz-utils 24 | 25 | - name: Bump version and push tag 26 | id: bump-version 27 | uses: anothrNick/github-tag-action@9aaabdb5e989894e95288328d8b17a6347217ae3 28 | env: 29 | GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} 30 | WITH_V: true 31 | DEFAULT_BUMP: patch 32 | INITIAL_VERSION: 0.1.0 33 | TAG_CONTEXT: repo 34 | 35 | - name: Build 36 | run: | 37 | cargo build --release 38 | tar cJf ./rfbproxy.tar.xz --owner=0:0 --transform=s@target/release@usr/bin@ target/release/rfbproxy 39 | 40 | - name: Create Release 41 | id: create-release 42 | uses: actions/create-release@v1 43 | env: 44 | GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} 45 | with: 46 | tag_name: ${{ steps.bump-version.outputs.tag }} 47 | release_name: ${{ steps.bump-version.outputs.tag }} 48 | draft: false 49 | prerelease: false 50 | 51 | - name: Upload release asset 52 | uses: actions/upload-release-asset@v1 53 | env: 54 | GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} 55 | with: 56 | upload_url: ${{ steps.create-release.outputs.upload_url }} 57 | asset_path: ./rfbproxy.tar.xz 58 | asset_name: rfbproxy.tar.xz 59 | asset_content_type: application/octet-stream 60 | -------------------------------------------------------------------------------- /src/messages/mod.rs: -------------------------------------------------------------------------------- 1 | //! Parsers for RFB messages. 2 | 3 | pub mod client; 4 | pub mod io; 5 | pub mod server; 6 | 7 | use std::fmt::Display; 8 | 9 | use bytes::Buf; 10 | 11 | /// An error for the message parsers. `Incomplete` is used to signal that the buffer from where the 12 | /// message is being parsed does not contain enough data to parse a full message. 13 | #[derive(Debug)] 14 | pub enum Error { 15 | Incomplete, 16 | Other(anyhow::Error), 17 | } 18 | 19 | impl std::error::Error for Error {} 20 | 21 | impl Display for Error { 22 | fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result { 23 | match self { 24 | Error::Incomplete => "stream ended early".fmt(fmt), 25 | Error::Other(err) => err.fmt(fmt), 26 | } 27 | } 28 | } 29 | 30 | /// A structure that represents how pixel values are represented in `FramebufferUpdate` messages. 31 | /// 32 | /// This is documented at 33 | /// 34 | #[derive(Debug)] 35 | #[allow(dead_code)] 36 | pub struct PixelFormat { 37 | pub bits_per_pixel: u8, 38 | depth: u8, 39 | big_endian_flag: u8, 40 | true_colour_flag: u8, 41 | red_max: u16, 42 | green_max: u16, 43 | blue_max: u16, 44 | red_shift: u8, 45 | green_shift: u8, 46 | blue_shift: u8, 47 | padding: [u8; 3], 48 | } 49 | 50 | impl PixelFormat { 51 | pub fn new(buf: &[u8]) -> PixelFormat { 52 | let mut cur = std::io::Cursor::new(buf); 53 | PixelFormat { 54 | bits_per_pixel: cur.get_u8(), 55 | depth: cur.get_u8(), 56 | big_endian_flag: cur.get_u8(), 57 | true_colour_flag: cur.get_u8(), 58 | red_max: cur.get_u16(), 59 | green_max: cur.get_u16(), 60 | blue_max: cur.get_u16(), 61 | red_shift: cur.get_u8(), 62 | green_shift: cur.get_u8(), 63 | blue_shift: cur.get_u8(), 64 | padding: [0, 0, 0], 65 | } 66 | } 67 | } 68 | -------------------------------------------------------------------------------- /src/messages/io.rs: -------------------------------------------------------------------------------- 1 | use super::Error; 2 | 3 | use anyhow::{anyhow, Result}; 4 | use bytes::Buf; 5 | 6 | /// The maximum message length. Anything larger than this will be rejected. 7 | pub(crate) const MAX_MESSAGE_LENGTH: usize = 4 * 1024 * 1024; 8 | 9 | pub(crate) fn check_message_length(length: T) -> Result<(), Error> 10 | where 11 | T: Into, 12 | { 13 | let message_length = length.into(); 14 | if message_length > MAX_MESSAGE_LENGTH { 15 | return Err(Error::Other(anyhow!( 16 | "message size is larger than the allowed limit: {}", 17 | message_length, 18 | ))); 19 | } 20 | Ok(()) 21 | } 22 | 23 | pub(crate) fn peek_u8(src: &mut std::io::Cursor<&[u8]>) -> Result { 24 | if !src.has_remaining() { 25 | return Err(Error::Incomplete); 26 | } 27 | 28 | Ok(src.chunk()[0]) 29 | } 30 | 31 | pub(crate) fn get_u8(src: &mut std::io::Cursor<&[u8]>) -> Result { 32 | if !src.has_remaining() { 33 | return Err(Error::Incomplete); 34 | } 35 | 36 | Ok(src.get_u8()) 37 | } 38 | 39 | pub fn get_u16(src: &mut std::io::Cursor<&[u8]>) -> Result { 40 | if src.remaining() < 2 { 41 | return Err(Error::Incomplete); 42 | } 43 | 44 | Ok(src.get_u16()) 45 | } 46 | 47 | pub fn get_u32(src: &mut std::io::Cursor<&[u8]>) -> Result { 48 | if src.remaining() < 4 { 49 | return Err(Error::Incomplete); 50 | } 51 | 52 | Ok(src.get_u32()) 53 | } 54 | 55 | pub fn get_i32(src: &mut std::io::Cursor<&[u8]>) -> Result { 56 | if src.remaining() < 4 { 57 | return Err(Error::Incomplete); 58 | } 59 | 60 | Ok(src.get_i32()) 61 | } 62 | 63 | pub fn get_compact(src: &mut std::io::Cursor<&[u8]>) -> Result { 64 | let mut l = get_u8(src)? as usize; 65 | let mut len = l & 0x7f; 66 | if (l & 0x80) != 0 { 67 | l = get_u8(src)? as usize; 68 | len |= (l & 0x7f) << 7; 69 | if (l & 0x80) != 0 { 70 | l = get_u8(src)? as usize; 71 | len |= l << 14; 72 | } 73 | } 74 | 75 | Ok(len) 76 | } 77 | 78 | pub fn skip(src: &mut std::io::Cursor<&[u8]>, n: usize) -> Result<(), Error> { 79 | if src.remaining() < n { 80 | return Err(Error::Incomplete); 81 | } 82 | 83 | src.advance(n); 84 | Ok(()) 85 | } 86 | 87 | pub fn read<'a>(src: &mut std::io::Cursor<&'a [u8]>, n: usize) -> Result<&'a [u8], Error> { 88 | if src.remaining() < n { 89 | return Err(Error::Incomplete); 90 | } 91 | 92 | let pos = src.position() as usize; 93 | let result = &src.get_ref()[pos..pos + n]; 94 | src.advance(n); 95 | Ok(result) 96 | } 97 | 98 | pub fn peek(src: &std::io::Cursor<&[u8]>, n: usize) -> Result, Error> { 99 | if src.remaining() < n { 100 | return Err(Error::Incomplete); 101 | } 102 | 103 | Ok(src.chunk()[0..n].to_vec()) 104 | } 105 | 106 | #[cfg(test)] 107 | mod tests { 108 | use super::*; 109 | 110 | #[test] 111 | fn test_get_compact() -> Result<(), Error> { 112 | let data = vec![0x90, 0x4E]; 113 | let mut cur = std::io::Cursor::new(&data[..]); 114 | assert_eq!(get_compact(&mut cur)?, 10000); 115 | Ok(()) 116 | } 117 | } 118 | -------------------------------------------------------------------------------- /src/audio/mp3.rs: -------------------------------------------------------------------------------- 1 | //! An MP3 stream in a MPEG-1 container. 2 | 3 | use super::Encoder; 4 | 5 | use anyhow::{bail, Context, Result}; 6 | 7 | /// Error checker for lame_sys functions. 8 | /// 9 | /// Returns an [`anyhow::Error`] with the stringified error if the result of the function call is negative, 10 | /// and the result of the function as-is otherwise. 11 | fn check_err(num: std::os::raw::c_int) -> Result { 12 | if num == lame_sys::lame_errorcodes_t::LAME_GENERICERROR as i32 { 13 | bail!("lame: generic error"); 14 | } else if num == lame_sys::lame_errorcodes_t::LAME_NOMEM as i32 { 15 | bail!("lame: out of memory"); 16 | } else if num == lame_sys::lame_errorcodes_t::LAME_BADBITRATE as i32 { 17 | bail!("lame: unsupported bitrate"); 18 | } else if num == lame_sys::lame_errorcodes_t::LAME_BADSAMPFREQ as i32 { 19 | bail!("lame: unsupported sample rate"); 20 | } else if num == lame_sys::lame_errorcodes_t::LAME_INTERNALERROR as i32 { 21 | bail!("lame: internal error"); 22 | } else if num == lame_sys::lame_errorcodes_t::FRONTEND_READERROR as i32 { 23 | bail!("lame: frontend read error"); 24 | } else if num == lame_sys::lame_errorcodes_t::FRONTEND_WRITEERROR as i32 { 25 | bail!("lame: frontend write error"); 26 | } else if num == lame_sys::lame_errorcodes_t::FRONTEND_FILETOOLARGE as i32 { 27 | bail!("lame: frontend file too large"); 28 | } else if num < 0 { 29 | bail!("lame: unknown error: {}", num); 30 | } 31 | Ok(num) 32 | } 33 | 34 | /// Mp3Encoder is a MPEG-1 Audio Layer II streaming muxer compliant codec that can be played in a 35 | /// browser using Media Source Extensions. 36 | pub struct Mp3Encoder { 37 | ctx: *mut lame_sys::lame_global_flags, 38 | } 39 | 40 | impl Mp3Encoder { 41 | pub fn new(channels: u8, kbps: u16) -> Result { 42 | let ctx = unsafe { lame_sys::lame_init() }; 43 | if ctx.is_null() { 44 | bail!("could not initialize the lame library"); 45 | } 46 | check_err(unsafe { lame_sys::lame_set_num_channels(ctx, channels as i32) }) 47 | .with_context(|| format!("failed to set channels to {channels}"))?; 48 | check_err(unsafe { lame_sys::lame_set_brate(ctx, kbps as i32) }) 49 | .with_context(|| format!("failed to set kbps to {kbps}"))?; 50 | check_err(unsafe { lame_sys::lame_set_quality(ctx, 7) }) 51 | .context("failed to set quality to 7")?; 52 | check_err(unsafe { lame_sys::lame_init_params(ctx) }) 53 | .context("failed to initialize the lame parameters")?; 54 | Ok(Mp3Encoder { ctx }) 55 | } 56 | } 57 | 58 | impl Drop for Mp3Encoder { 59 | fn drop(&mut self) { 60 | unsafe { lame_sys::lame_close(self.ctx) }; 61 | self.ctx = std::ptr::null_mut(); 62 | } 63 | } 64 | 65 | impl Encoder for Mp3Encoder { 66 | fn sample_rate(&self) -> u32 { 67 | 44100 68 | } 69 | 70 | fn frame_len_ms(&self) -> u64 { 71 | // The MP3 codec uses 576 or 1152 samples (depending on the flags used to initialize the 72 | // codec), but the sample rate is not cleanly divisible by that. The closest would be 20ms, 73 | // but that causes stuttering, so 40ms it is! 74 | 40 75 | } 76 | 77 | fn internal_delay(&mut self) -> Result { 78 | Ok(576) 79 | } 80 | 81 | fn max_frame_len(&self) -> usize { 82 | (1.25 * self.frame_len_ms() as f64 * self.sample_rate() as f64) as usize + 7200 83 | } 84 | 85 | fn encode_frame( 86 | &mut self, 87 | _timestamp: u64, 88 | input: &[i16], 89 | output: &mut [u8], 90 | ) -> Result<(usize, bool)> { 91 | Ok(( 92 | check_err(unsafe { 93 | lame_sys::lame_encode_buffer_interleaved( 94 | self.ctx, 95 | input.as_ptr() as *mut i16, 96 | (input.len() / 2) as std::os::raw::c_int, 97 | output.as_mut_ptr(), 98 | output.len() as std::os::raw::c_int, 99 | ) 100 | })? as usize, 101 | true, // All MP3 frames have a valid MPEG-1 header. 102 | )) 103 | } 104 | } 105 | 106 | unsafe impl Send for Mp3Encoder {} 107 | 108 | #[cfg(test)] 109 | mod tests { 110 | use super::*; 111 | use crate::audio::Encoder; 112 | 113 | use rand::Rng; 114 | 115 | fn init() { 116 | let _ = env_logger::builder().is_test(true).try_init(); 117 | } 118 | 119 | #[test] 120 | fn test_mp3() -> Result<()> { 121 | init(); 122 | 123 | let mut enc = Mp3Encoder::new(2, 32 * 1024).expect("could not create MP3 encoder"); 124 | 125 | let mut chunk = 126 | vec![ 127 | 0i16; 128 | enc.channels() as usize * enc.sample_rate() as usize * enc.frame_len_ms() as usize 129 | / 1000 130 | ]; 131 | rand::thread_rng() 132 | .try_fill(&mut chunk[..]) 133 | .expect("could not generate random payload"); 134 | 135 | log::info!("Read audio chunk"); 136 | 137 | let mut payload = Vec::::with_capacity(enc.max_frame_len()); 138 | payload.resize(enc.max_frame_len(), 0); 139 | payload[..4].copy_from_slice(&[ 140 | 0xFF, // message-type 141 | 0x01, // submessage-type 142 | 0x00, 0x02, // operation 143 | ]); 144 | 145 | log::info!("Encoding audio chunk..."); 146 | let mut payload_size = enc 147 | .encode_frame(0, &chunk, &mut payload[8..]) 148 | .expect("could not encode audio data") 149 | .0; 150 | payload_size += enc 151 | .encode_frame(0, &chunk, &mut payload[8 + payload_size..]) 152 | .expect("could not encode audio data") 153 | .0; 154 | log::info!("Wrote {} audio bytes!", payload_size); 155 | 156 | assert!(payload_size > 0); 157 | 158 | payload[4] = ((payload_size >> 24) & 0xff) as u8; 159 | payload[5] = ((payload_size >> 16) & 0xff) as u8; 160 | payload[6] = ((payload_size >> 8) & 0xff) as u8; 161 | payload[7] = ((payload_size >> 0) & 0xff) as u8; 162 | payload.truncate(payload_size + 8); 163 | assert!(payload.len() > 5); 164 | 165 | Ok(()) 166 | } 167 | } 168 | -------------------------------------------------------------------------------- /src/client.proto: -------------------------------------------------------------------------------- 1 | syntax = "proto3"; 2 | 3 | import "google/protobuf/timestamp.proto"; 4 | 5 | package api; 6 | option go_package = "/api"; 7 | 8 | // This message constitutes the repl metadata and define the repl we're 9 | // connecting to. All fields are required unless otherwise stated. 10 | message Repl { 11 | string id = 1; 12 | string language = 2; 13 | string bucket = 3; 14 | string slug = 4; 15 | string user = 5; 16 | 17 | // (Optional) The replID of a repl to be used as the source filesystem. All 18 | // writes will still go to the actual repl. This is intended to be a 19 | // replacement for guest repls, giving us cheap COW semantics so all 20 | // connections can have a real repl. 21 | // 22 | // One exception: 23 | // 24 | // It's important to note that data is not implicitly copied from src to 25 | // dest. Only what is explicitly written when talking to pid1 (either 26 | // gcsfiles or snapshots) will persist. This makes it slightly different 27 | // than just forking. 28 | // 29 | // It's unclear what the behaviour should be if: 30 | // - the dest and src repl both exist 31 | // - the dest and src are the same 32 | // - we have an src but no dest 33 | // 34 | // consider these unsupported/undefined for now. 35 | string sourceRepl = 6; 36 | } 37 | 38 | // The resource limits that should be applied to the Repl's container. 39 | message ResourceLimits { 40 | // Whether the repl has network access. 41 | bool net = 1; 42 | 43 | // The amount of RAM in bytes that this repl will have. 44 | int64 memory = 2; 45 | 46 | // The number of cores that the container will be allowed to have. 47 | double threads = 3; 48 | 49 | // The Docker container weight factor for the scheduler. Similar to the 50 | // `--cpu-shares` commandline flag. 51 | double shares = 4; 52 | 53 | // The size of the disk in bytes. 54 | int64 disk = 5; 55 | 56 | // Whether these limits are cachable, and if they are, by what facet of the token. 57 | enum Cachability { 58 | // Do not cache these limits. 59 | NONE = 0; 60 | 61 | // These limits can be cached and applied to this and any of the user's 62 | // other repls. 63 | USER = 1; 64 | 65 | // These limits can be cached and applied only to this repl. 66 | REPL = 2; 67 | } 68 | 69 | Cachability cache = 6; 70 | } 71 | 72 | // Permissions allow tokens to perform certain actions. 73 | message Permissions { 74 | // This token has permission to toggle the always on state of a container. 75 | // For a connection to send the AlwaysOn message, it must have this permission. 76 | bool toggleAlwaysOn = 1; 77 | } 78 | 79 | // ReplToken is the expected client options during the handshake. This is encoded 80 | // into the token that is used to connect using WebSocket. 81 | message ReplToken { 82 | // Issue timestamp. Equivalent to JWT's "iat" (Issued At) claim. Tokens with 83 | // no `iat` field will be treated as if they had been issed at the UNIX epoch 84 | // (1970-01-01T00:00:00Z). 85 | google.protobuf.Timestamp iat = 1; 86 | 87 | // Expiration timestamp. Equivalent to JWT's "exp" (Expiration Time) Claim. 88 | // If unset, will default to one hour after `iat`. 89 | google.protobuf.Timestamp exp = 2; 90 | 91 | // An arbitrary string that helps prevent replay attacks by ensuring that all 92 | // tokens are distinct. 93 | string salt = 3; 94 | 95 | // The cluster that a repl is located in. This prevents replay attacks in 96 | // which a user is given a token for one cluster and then presents that same 97 | // token to a conman instance in another token, which could lead to a case 98 | // where multiple containers are associated with a repl. 99 | // 100 | // Conman therefore needs to validate that this parameter matches the 101 | // `-cluster` flag it was started with. 102 | string cluster = 4; 103 | 104 | // Whether to persist filesystem, metadata, or both. 105 | enum Persistence { 106 | // This is the usual mode of operation: both filesystem and metadata will be 107 | // persisted. 108 | PERSISTENT = 0; 109 | 110 | // The ephemeral flag indicates the repl being connected to will have a time 111 | // restriction on stored metadata. This has the consequence that repl will 112 | // be unable to wakeup or serve static traffic once the metadata has timed 113 | // out. This option does NOT affect filesystem and other data persistence. 114 | // 115 | // For context, this value is used on the client when repls are created for: 116 | // - replrun 117 | // - guests 118 | // - anon users 119 | // - temp vnc repls 120 | // - users with non-verified emails 121 | EPHEMERAL = 1; 122 | 123 | // This indicates that the repl being connected does not have the ability to 124 | // persist files or be woken up after the lifetime of this repl expires. 125 | // 126 | // For context, this value is used on the client when repls are created for: 127 | // - replrun 128 | // - guests 129 | // - language pages 130 | NONE = 2; 131 | } 132 | // Whether to persist filesystem, metadata, or both. When connecting to an 133 | // already running/existing repl, its settings will be updated to match this 134 | // mode. 135 | Persistence persistence = 6; 136 | 137 | // Metadata for the classroom. This is deprecated and should be removed 138 | // hopefully soon. 139 | message ClassroomMetadata { 140 | string id = 1; 141 | string language = 2; 142 | } 143 | 144 | // Metadata for a repl that is only identified by its id. 145 | message ReplID { 146 | string id = 1; 147 | 148 | // (Optional) See the comment for Repl.sourceRepl. 149 | string sourceRepl = 2; 150 | } 151 | 152 | // One of the three ways to identify a repl in goval. 153 | oneof metadata { 154 | // This is the standard connection behavior. If the repl doesn't exist it 155 | // will be created. Any future connections with a matching ID will go to 156 | // the same container. If other metadata mismatches besides ID it will be 157 | // rectified (typically by recreating the container to make it match the 158 | // provided value). 159 | Repl repl = 7; 160 | 161 | // The repl must already be known to goval, the connection will proceed 162 | // with the Repl metadata from a previous connection's metadata with the 163 | // same ID. 164 | ReplID id = 8; 165 | 166 | // This is DEPRECATED and only used by the classroom. This will never share 167 | // a container between connections. Please don't use this even for tests, 168 | // we intend to remove it soon. 169 | ClassroomMetadata classroom = 9 [deprecated=true]; 170 | } 171 | 172 | // The resource limits for the container. 173 | ResourceLimits resourceLimits = 10; 174 | 175 | // allows the client to choose a wire format. 176 | enum WireFormat { 177 | // The default wire format: Protobuf-over-WebSocket. 178 | PROTOBUF = 0; 179 | 180 | // Legacy protocol. 181 | JSON = 1 [deprecated=true]; 182 | } 183 | WireFormat format = 12; 184 | 185 | message Presenced { 186 | uint32 bearerID = 1; 187 | string bearerName = 2; 188 | } 189 | Presenced presenced = 13; 190 | 191 | // Flags are handy for passing arbitrary configs along. Mostly used so 192 | // the client can try out new features 193 | repeated string flags = 14; 194 | 195 | Permissions permissions = 15; 196 | } 197 | 198 | // GovalTokenMetadata is information about a goval token, that can be used to 199 | // validate it. It is stored in the footer of the PASETO. 200 | message GovalTokenMetadata { 201 | // The ID of the key that was used to sign the token. 202 | string key_id = 1; 203 | } 204 | -------------------------------------------------------------------------------- /src/audio/mod.rs: -------------------------------------------------------------------------------- 1 | //! A stream encoder that gets the audio samples from PulseAudio. 2 | //! 3 | //! Currently, Opus (in a WebM streamable container) and MP3 (in an MPEG-1 container) are 4 | //! supported, which should cover all of the major browsers. 5 | 6 | mod mp3; 7 | mod opus; 8 | 9 | use anyhow::{bail, Context, Result}; 10 | use byte_slice_cast::*; 11 | use pulse::sample; 12 | use pulse::stream::Direction; 13 | use tokio::sync::{mpsc, oneshot}; 14 | 15 | /// A muxer that can write frames of encoded audio into a container that is appropriate for 16 | /// streaming. 17 | trait StreamMuxer { 18 | /// Writes an encoded packet into a [`std::io::Write`]. 19 | fn write_frame( 20 | &mut self, 21 | w: &mut W, 22 | frame: &[u8], 23 | timestamp: u64, 24 | keyframe: bool, 25 | ) -> Result<()>; 26 | } 27 | 28 | /// An encoder for an audio stream. 29 | trait Encoder { 30 | /// The preferred sample rate for this encoder. 31 | fn sample_rate(&self) -> u32; 32 | 33 | /// The frame length, in milliseconds. 34 | fn frame_len_ms(&self) -> u64; 35 | 36 | /// The number of channels. 37 | #[allow(dead_code)] 38 | fn channels(&self) -> u8 { 39 | 2 40 | } 41 | 42 | /// The internal delay of the codec, in samples. 43 | #[allow(dead_code)] 44 | fn internal_delay(&mut self) -> Result; 45 | 46 | /// The maximum size of an encoded and contained frame. 47 | fn max_frame_len(&self) -> usize; 48 | 49 | /// Encodes a frame into the provided buffer, including the container. Returns the size of the 50 | /// payload, and whether the frame is a keyframe. 51 | fn encode_frame( 52 | &mut self, 53 | timestamp: u64, 54 | input: &[i16], 55 | output: &mut [u8], 56 | ) -> Result<(usize, bool)>; 57 | } 58 | 59 | /// A stream that encodes PulseAudio-provided audio and generates Replit Audio Server AudioData 60 | /// messages. 61 | pub struct Stream { 62 | enc: Box, 63 | pulse: psimple::Simple, 64 | buffer: Vec, 65 | stream_start: Option, 66 | timestamp: u64, 67 | dropped_frames: u64, 68 | } 69 | 70 | impl Stream { 71 | pub fn new(channels: u8, codec: u16, kbps: u16) -> Result { 72 | let enc: Box = match codec { 73 | 0 => Box::new( 74 | opus::OpusEncoder::new(channels, kbps).context("could not create Opus encoder")?, 75 | ), 76 | 1 => Box::new( 77 | mp3::Mp3Encoder::new(channels, kbps).context("could not create MP3 encoder")?, 78 | ), 79 | _ => bail!("unsupported codec: {}", codec), 80 | }; 81 | let frame_len_ms = enc.frame_len_ms(); 82 | let sample_rate = enc.sample_rate(); 83 | 84 | log::debug!("Opened audio pipe"); 85 | 86 | let spec = sample::Spec { 87 | format: sample::Format::S16le, 88 | channels, 89 | rate: enc.sample_rate(), 90 | }; 91 | if !spec.is_valid() { 92 | bail!("invalid channels / sample rate combination"); 93 | } 94 | 95 | let pulse = match psimple::Simple::new( 96 | None, // Use the default server 97 | "audiomuxer", // Our application’s name 98 | Direction::Record, // We want a recording stream 99 | None, // Use the default device 100 | "Music", // Description of our stream 101 | &spec, // Our sample format 102 | None, // Use default channel map 103 | None, // Use default buffering attributes 104 | ) { 105 | Ok(pulse) => pulse, 106 | Err(e) => bail!(e), 107 | }; 108 | 109 | log::debug!("Opened audio pipe"); 110 | 111 | Ok(Stream { 112 | enc, 113 | pulse, 114 | buffer: vec![ 115 | 0i16; 116 | channels as usize * sample_rate as usize * frame_len_ms as usize / 1000 117 | ], 118 | stream_start: None, 119 | dropped_frames: 0, 120 | timestamp: -(frame_len_ms as i64) as u64, 121 | }) 122 | } 123 | 124 | /// Reads a single audio frame into the internal buffer. 125 | fn read_frame(&mut self) -> Result<()> { 126 | if let Err(e) = self.pulse.read(self.buffer.as_mut_byte_slice()) { 127 | bail!("failed to read PulseAudio data: {}", e); 128 | } 129 | self.timestamp = std::cmp::max( 130 | self.timestamp.wrapping_add(self.enc.frame_len_ms()), 131 | self.stream_start 132 | .get_or_insert_with(std::time::Instant::now) 133 | .elapsed() 134 | .as_millis() as u64, 135 | ); 136 | log::debug!("Read audio chunk"); 137 | Ok(()) 138 | } 139 | 140 | /// Encodes the frame from the internal buffer and returns the Replit AudioFrame contents. 141 | fn encode_frame(&mut self) -> Result> { 142 | let mut payload = vec![0; self.enc.max_frame_len()]; 143 | payload[..2].copy_from_slice(&[ 144 | 0xF5, // message-type 145 | 0x01, // submessage-type 146 | ]); 147 | 148 | let (mut payload_size, keyframe) = self 149 | .enc 150 | .encode_frame(self.timestamp, &self.buffer, &mut payload[8..]) 151 | .context("could not encode frame")?; 152 | payload_size += 4; 153 | 154 | if payload_size > 0xFFFF { 155 | bail!("payload exceeds maximum size: {}", payload_size); 156 | } 157 | 158 | // The Most Significant Bit marks whether the frame is a keyframe. 159 | let timestamp = (self.timestamp & 0x7FFFFFFF) | ((keyframe as u64) << 31); 160 | payload[2] = ((payload_size >> 8) & 0xff) as u8; 161 | payload[3] = (payload_size & 0xff) as u8; 162 | payload[4] = ((timestamp >> 24) & 0xff) as u8; 163 | payload[5] = ((timestamp >> 16) & 0xff) as u8; 164 | payload[6] = ((timestamp >> 8) & 0xff) as u8; 165 | payload[7] = (timestamp & 0xff) as u8; 166 | payload.truncate(payload_size + 8); 167 | 168 | log::debug!( 169 | "Sending {} payload bytes at timestamp {}, {} frames dropped", 170 | payload.len(), 171 | self.timestamp, 172 | self.dropped_frames, 173 | ); 174 | 175 | Ok(payload) 176 | } 177 | 178 | /// Runs a PulseAudio thread using the simple (blocking) API. 179 | pub fn run( 180 | mut self, 181 | stop_chan: oneshot::Sender<()>, 182 | audio_message_chan: mpsc::Sender>, 183 | ) { 184 | log::info!("audio thread started"); 185 | 186 | while !stop_chan.is_closed() { 187 | // Always consume the frame from PulseAudio. That way this thread doesn't end up with 188 | // huge jumps or large number of dropped messages. 189 | if let Err(e) = self.read_frame() { 190 | log::error!("failed to read audio frame: {:#}", e); 191 | break; 192 | } 193 | let permit = match audio_message_chan.try_reserve() { 194 | Ok(permit) => permit, 195 | Err(mpsc::error::TrySendError::Full(_)) => { 196 | log::debug!("Dropped message due to backpressure"); 197 | self.dropped_frames += 1; 198 | continue; 199 | } 200 | Err(e) => { 201 | log::error!("failed to send audio data: {:#}", e); 202 | break; 203 | } 204 | }; 205 | let payload = match self.encode_frame() { 206 | Ok(payload) => payload, 207 | Err(e) => { 208 | log::error!("failed to read audio frame: {:#}", e); 209 | break; 210 | } 211 | }; 212 | permit.send(payload); 213 | } 214 | log::info!( 215 | "audio thread finished, {} dropped frames", 216 | self.dropped_frames 217 | ); 218 | } 219 | } 220 | -------------------------------------------------------------------------------- /src/messages/client.rs: -------------------------------------------------------------------------------- 1 | use super::io::*; 2 | use super::Error; 3 | 4 | use std::fmt::Display; 5 | 6 | use anyhow::{anyhow, Result}; 7 | use bytes::Buf; 8 | 9 | /// Represents a message that is sent from the client to the server. 10 | #[derive(Debug)] 11 | #[allow(clippy::enum_variant_names)] 12 | pub enum Message { 13 | SetPixelFormat(Vec), 14 | SetEncodings(Vec), 15 | FramebufferUpdateRequest(Vec), 16 | KeyEvent(Vec), 17 | PointerEvent(Vec), 18 | ClientCutText(Vec), 19 | 20 | // Extensions 21 | EnableContinuousUpdates(Vec), 22 | ClientFence(Vec), 23 | SetDesktopSize(Vec), 24 | QemuClientMessage(Vec), 25 | ReplitClientAudioStartEncoder(Vec, bool, u8, u16, u16), 26 | ReplitClientAudioFrameRequest(Vec), 27 | ReplitClientAudioStartContinuousUpdates(Vec), 28 | } 29 | 30 | impl Display for Message { 31 | fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result { 32 | use Message::*; 33 | 34 | match self { 35 | SetEncodings(payload) => { 36 | let mut cur = std::io::Cursor::new(&payload[4..]); 37 | write!(fmt, "SetEncodings(")?; 38 | while cur.has_remaining() { 39 | write!(fmt, " {}", cur.get_i32())?; 40 | } 41 | write!(fmt, " )") 42 | } 43 | m => { 44 | write!(fmt, "{m:?}") 45 | } 46 | } 47 | } 48 | } 49 | 50 | impl Message { 51 | /// Gets the raw u8 array representation of the message. 52 | pub fn into_data(self) -> Vec { 53 | use Message::*; 54 | 55 | match self { 56 | SetPixelFormat(payload) 57 | | SetEncodings(payload) 58 | | FramebufferUpdateRequest(payload) 59 | | KeyEvent(payload) 60 | | PointerEvent(payload) 61 | | ClientCutText(payload) 62 | | EnableContinuousUpdates(payload) 63 | | ClientFence(payload) 64 | | SetDesktopSize(payload) 65 | | ReplitClientAudioStartEncoder(payload, _, _, _, _) 66 | | ReplitClientAudioFrameRequest(payload) 67 | | ReplitClientAudioStartContinuousUpdates(payload) 68 | | QemuClientMessage(payload) => payload, 69 | } 70 | } 71 | 72 | /// Gets the length of the raw u8 array representation of the message. 73 | pub fn len(&self) -> usize { 74 | use Message::*; 75 | 76 | match self { 77 | SetPixelFormat(payload) 78 | | SetEncodings(payload) 79 | | FramebufferUpdateRequest(payload) 80 | | KeyEvent(payload) 81 | | PointerEvent(payload) 82 | | ClientCutText(payload) 83 | | EnableContinuousUpdates(payload) 84 | | ClientFence(payload) 85 | | SetDesktopSize(payload) 86 | | ReplitClientAudioStartEncoder(payload, _, _, _, _) 87 | | ReplitClientAudioFrameRequest(payload) 88 | | ReplitClientAudioStartContinuousUpdates(payload) 89 | | QemuClientMessage(payload) => payload.len(), 90 | } 91 | } 92 | 93 | /// Tries to parse a message from a [`std::io::Cursor`]. 94 | /// 95 | /// # Errors 96 | /// 97 | /// IF there is not enough data in `src` to parse a complete message, it will return a 98 | /// [`Error::Incomplete`]. 99 | pub(crate) fn parse(src: &mut std::io::Cursor<&[u8]>) -> Result { 100 | match peek_u8(src)? { 101 | 0 => { 102 | let payload = read(src, 20)?; 103 | Ok(Message::SetPixelFormat(payload.to_vec())) 104 | } 105 | 2 => { 106 | skip(src, 2)?; 107 | let number_of_encodings = get_u16(src)? as usize; 108 | src.set_position(0); 109 | let payload = read(src, 4 + 4 * number_of_encodings)?; 110 | Ok(Message::SetEncodings(payload.to_vec())) 111 | } 112 | 3 => { 113 | let payload = read(src, 10)?; 114 | Ok(Message::FramebufferUpdateRequest(payload.to_vec())) 115 | } 116 | 4 => { 117 | let payload = read(src, 8)?; 118 | Ok(Message::KeyEvent(payload.to_vec())) 119 | } 120 | 5 => { 121 | let payload = read(src, 6)?; 122 | Ok(Message::PointerEvent(payload.to_vec())) 123 | } 124 | 6 => { 125 | skip(src, 4)?; // id + padding 126 | let length = get_i32(src)?; 127 | 128 | if length >= 0 { 129 | // Standard message 130 | check_message_length(length as usize)?; 131 | let len = src.position() as usize; 132 | src.set_position(0); 133 | let payload = read(src, len + length as usize)?; 134 | return Ok(Message::ClientCutText(payload.to_vec())); 135 | } 136 | 137 | // Extended message 138 | check_message_length(-length as usize)?; 139 | let flags = get_u32(src)?; 140 | log::debug!("Extended clipboard message: {} {:x}", -length, flags); 141 | 142 | src.set_position(0); 143 | let payload = read(src, 8 + (-length) as usize)?; 144 | Ok(Message::ClientCutText(payload.to_vec())) 145 | } 146 | 150 => { 147 | let payload = read(src, 10)?; 148 | Ok(Message::EnableContinuousUpdates(payload.to_vec())) 149 | } 150 | 245 => { 151 | // Replit Audio Client Message 152 | skip(src, 1)?; // id 153 | let submessage_id = get_u8(src)?; 154 | let length = get_u16(src)? as usize; 155 | match submessage_id { 156 | 0 => { 157 | // Replit Audio Client Message StartEncoder 158 | if length != 6 { 159 | return Err(Error::Other(anyhow!( 160 | "unexpected StartEncoder length: got {}, expected 6", 161 | length 162 | ))); 163 | } 164 | 165 | let enabled = get_u8(src)? != 0; 166 | let channels = get_u8(src)?; 167 | let codec = get_u16(src)?; 168 | let kbps = get_u16(src)?; 169 | src.set_position(0); 170 | let payload = read(src, length + 4)?; 171 | Ok(Message::ReplitClientAudioStartEncoder( 172 | payload.to_vec(), 173 | enabled, 174 | channels, 175 | codec, 176 | kbps, 177 | )) 178 | } 179 | 1 => { 180 | // Replit Audio Client Message FrameRequest 181 | if length != 0 { 182 | return Err(Error::Other(anyhow!( 183 | "unexpected FrameRequest length: got {}, expected 0", 184 | length 185 | ))); 186 | } 187 | 188 | src.set_position(0); 189 | let payload = read(src, length + 4)?; 190 | Ok(Message::ReplitClientAudioFrameRequest(payload.to_vec())) 191 | } 192 | 2 => { 193 | // Replit Audio Client Message EnableContinuousUpdates 194 | if length != 0 { 195 | return Err(Error::Other(anyhow!( 196 | "unexpected EnableContinuousUpdates length: got {}, expected 0", 197 | length 198 | ))); 199 | } 200 | 201 | src.set_position(0); 202 | let payload = read(src, length + 4)?; 203 | Ok(Message::ReplitClientAudioStartContinuousUpdates( 204 | payload.to_vec(), 205 | )) 206 | } 207 | submessage_id => { 208 | src.set_position(0); 209 | Err(Error::Other(anyhow!( 210 | "unsupported Replit Client submessage id: {} {:?}", 211 | submessage_id, 212 | src.get_ref() 213 | ))) 214 | } 215 | } 216 | } 217 | 248 => { 218 | skip(src, 8)?; // id, padding, flags 219 | let length = get_u8(src)? as usize; 220 | 221 | let len = src.position() as usize; 222 | src.set_position(0); 223 | let payload = read(src, len + length)?; 224 | Ok(Message::ClientFence(payload.to_vec())) 225 | } 226 | 251 => { 227 | skip(src, 6)?; // id, padding, width, height 228 | let number_of_screens = get_u8(src)? as usize; 229 | 230 | let len = src.position() as usize; 231 | src.set_position(0); 232 | let payload = read(src, len + 1 + 16 * number_of_screens)?; 233 | Ok(Message::SetDesktopSize(payload.to_vec())) 234 | } 235 | 255 => { 236 | skip(src, 1)?; // id 237 | match get_u8(src)? { 238 | 0 => { 239 | // QEMU Extended Key Event Message. 240 | src.set_position(0); 241 | let payload = read(src, 12)?; 242 | Ok(Message::QemuClientMessage(payload.to_vec())) 243 | } 244 | submessage_id => { 245 | src.set_position(0); 246 | Err(Error::Other(anyhow!( 247 | "unsupported QEMU Client submessage id: {} {:?}", 248 | submessage_id, 249 | src.get_ref() 250 | ))) 251 | } 252 | } 253 | } 254 | message_id => Err(Error::Other(anyhow!( 255 | "unsupported client message id: {} {:?}", 256 | message_id, 257 | src.get_ref(), 258 | ))), 259 | } 260 | } 261 | } 262 | -------------------------------------------------------------------------------- /src/messages/server.rs: -------------------------------------------------------------------------------- 1 | use super::io::*; 2 | use super::Error; 3 | 4 | use anyhow::{anyhow, Result}; 5 | 6 | /// Represents a message that is sent from the server to the client. 7 | #[derive(Debug)] 8 | #[allow(clippy::enum_variant_names)] 9 | pub enum Message { 10 | FramebufferUpdate(Vec), 11 | SetColorMapEntries(Vec), 12 | Bell(Vec), 13 | ServerCutText(Vec), 14 | 15 | // Extensions 16 | EndOfContinuousUpdates(Vec), 17 | ServerFence(Vec), 18 | ReplitAudioServerMessage(Vec), 19 | } 20 | 21 | impl Message { 22 | /// Gets the raw u8 array representation of the message. 23 | pub fn into_data(self) -> Vec { 24 | use Message::*; 25 | 26 | match self { 27 | FramebufferUpdate(payload) 28 | | SetColorMapEntries(payload) 29 | | Bell(payload) 30 | | ServerCutText(payload) 31 | | EndOfContinuousUpdates(payload) 32 | | ReplitAudioServerMessage(payload) 33 | | ServerFence(payload) => payload, 34 | } 35 | } 36 | 37 | /// Gets the length of the raw u8 array representation of the message. 38 | pub fn len(&self) -> usize { 39 | use Message::*; 40 | 41 | match self { 42 | FramebufferUpdate(payload) 43 | | SetColorMapEntries(payload) 44 | | Bell(payload) 45 | | ServerCutText(payload) 46 | | EndOfContinuousUpdates(payload) 47 | | ReplitAudioServerMessage(payload) 48 | | ServerFence(payload) => payload.len(), 49 | } 50 | } 51 | 52 | /// Tries to parse a message from a [`std::io::Cursor`]. 53 | /// 54 | /// # Errors 55 | /// 56 | /// IF there is not enough data in `src` to parse a complete message, it will return a 57 | /// [`Error::Incomplete`]. 58 | pub(crate) fn parse( 59 | src: &mut std::io::Cursor<&[u8]>, 60 | bytes_per_pixel: usize, 61 | ) -> Result { 62 | match peek_u8(src)? { 63 | 0 => { 64 | skip(src, 2)?; // id + padding 65 | let number_of_rectangles = get_u16(src)?; 66 | for _ in 0..number_of_rectangles { 67 | let _x = get_u16(src)? as usize; 68 | let _y = get_u16(src)? as usize; 69 | let width = get_u16(src)? as usize; 70 | let height = get_u16(src)? as usize; 71 | 72 | log::debug!("Rectangle encoding {:?}", peek(src, 4)?); 73 | 74 | // encoding type 75 | (match get_i32(src)? { 76 | // Raw Encoding, width*height*bytesPerPixel 77 | 0 => skip(src, width * height * bytes_per_pixel), 78 | // CopyRect Encoding, 4 bytes 79 | 1 => skip(src, 4), 80 | // Tight Encoding. 81 | 7 => { 82 | let mut ctl = get_u8(src)?; 83 | log::debug!(" Tight control {:x}", ctl); 84 | ctl >>= 4; 85 | if ctl == 0x08 { 86 | skip(src, 3) 87 | } else if ctl == 0x09 { 88 | let len = get_compact(src)?; 89 | skip(src, len) 90 | } else { 91 | let mut filter = 0u8; 92 | if (ctl & 0x4) != 0 { 93 | filter = get_u8(src)?; 94 | } 95 | log::debug!(" Filter {:x}", filter); 96 | 97 | if filter == 0 { 98 | // CopyFilter 99 | let uncompressed_size = width * height * 3; 100 | if uncompressed_size < 12 { 101 | skip(src, uncompressed_size) 102 | } else { 103 | let len = get_compact(src)?; 104 | skip(src, len) 105 | } 106 | } else if filter == 1 { 107 | // PaletteFilter 108 | let number_of_colors = get_u8(src)? as usize + 1; 109 | 110 | log::debug!( 111 | " number of colors: {}, {:?}", 112 | number_of_colors, 113 | peek(src, 3 * number_of_colors)?, 114 | ); 115 | skip(src, 3 * number_of_colors)?; 116 | 117 | let bpp = if number_of_colors <= 2 { 1 } else { 8 }; 118 | let row_size = (width * bpp).div_ceil(8); 119 | let uncompressed_size = row_size * height; 120 | 121 | log::debug!( 122 | " bpp: {}, row_size: {}, uncompressed_size: {}", 123 | bpp, 124 | row_size, 125 | uncompressed_size 126 | ); 127 | 128 | if uncompressed_size < 12 { 129 | skip(src, uncompressed_size) 130 | } else { 131 | let len = get_compact(src)?; 132 | skip(src, len) 133 | } 134 | } else { 135 | Err(Error::Other(anyhow!( 136 | "unsupported Tight filter: {}", 137 | filter 138 | ))) 139 | } 140 | } 141 | } 142 | // ZRLE Encoding, 4 bytes + payload 143 | 16 => { 144 | let length = get_u32(src)? as usize; 145 | check_message_length(length)?; 146 | skip(src, length) 147 | } 148 | // LastRect Pseudo-encoding 149 | -224 => break, 150 | // Cursor Pseudo-encoding 151 | -239 => { 152 | skip(src, width * height * bytes_per_pixel)?; 153 | let row_size = width.div_ceil(8); 154 | skip(src, row_size * height) 155 | } 156 | // QEMU Extended Key Event Pseudo-encoding 157 | -258 => Ok(()), 158 | // ExtendedDesktopSize Pseudo-encoding 159 | -308 => { 160 | let number_of_screens = get_u8(src)? as usize; 161 | skip(src, 3 + number_of_screens * 16) 162 | } 163 | // VMware Cursor Pseudo-encoding 164 | 0x574d5664 => match get_u8(src)? { 165 | 0 => skip(src, 1 + 2 * width * height * bytes_per_pixel), 166 | 1 => skip(src, 1 + width * height * 4), 167 | cursor_type => Err(Error::Other(anyhow!( 168 | "unsupported VMWare Cursor type: {}", 169 | cursor_type 170 | ))), 171 | }, 172 | encoding_type => Err(Error::Other(anyhow!( 173 | "unsupported encoding type: {}", 174 | encoding_type 175 | ))), 176 | })? 177 | } 178 | let len = src.position() as usize; 179 | src.set_position(0); 180 | let payload = read(src, len)?; 181 | Ok(Message::FramebufferUpdate(payload.to_vec())) 182 | } 183 | 1 => { 184 | skip(src, 4)?; // id, padding, first color 185 | let number_of_colors = get_u16(src)? as usize; 186 | 187 | let len = src.position() as usize; 188 | src.set_position(0); 189 | let payload = read(src, len + number_of_colors * 6)?; 190 | Ok(Message::SetColorMapEntries(payload.to_vec())) 191 | } 192 | 2 => { 193 | let payload = read(src, 1)?; 194 | Ok(Message::Bell(payload.to_vec())) 195 | } 196 | 3 => { 197 | skip(src, 4)?; // id + padding 198 | let length = get_i32(src)?; 199 | 200 | if length >= 0 { 201 | // Standard messagee 202 | check_message_length(length as usize)?; 203 | let len = src.position() as usize; 204 | src.set_position(0); 205 | let payload = read(src, len + length as usize)?; 206 | return Ok(Message::ServerCutText(payload.to_vec())); 207 | } 208 | 209 | // Extended message 210 | check_message_length((-length) as usize)?; 211 | let flags = get_u32(src)?; 212 | log::debug!("Extended clipboard message: {} {:x}", -length, flags); 213 | 214 | src.set_position(0); 215 | let payload = read(src, 8 + (-length) as usize)?; 216 | Ok(Message::ServerCutText(payload.to_vec())) 217 | } 218 | 150 => { 219 | let payload = read(src, 1)?; 220 | Ok(Message::EndOfContinuousUpdates(payload.to_vec())) 221 | } 222 | 248 => { 223 | skip(src, 8)?; // id, padding, flags 224 | let length = get_u8(src)? as usize; 225 | 226 | let len = src.position() as usize; 227 | src.set_position(0); 228 | let payload = read(src, len + length)?; 229 | Ok(Message::ServerFence(payload.to_vec())) 230 | } 231 | message_id => Err(Error::Other(anyhow!( 232 | "unsupported server message id: {}", 233 | message_id 234 | ))), 235 | } 236 | } 237 | } 238 | -------------------------------------------------------------------------------- /src/audio/opus.rs: -------------------------------------------------------------------------------- 1 | //! An Opus stream in a WebM streaming container. 2 | 3 | use super::{Encoder, StreamMuxer}; 4 | 5 | use anyhow::{bail, Context, Result}; 6 | 7 | use rand::Rng; 8 | 9 | /// OpusEncoder is a WebM streaming compliant codec using the [`WebmStreamMuxer`], that can be 10 | /// played in a browser using Media Source Extensions. 11 | pub struct OpusEncoder { 12 | enc: opus::Encoder, 13 | muxer: WebmStreamMuxer, 14 | wrote_initialization_segment: bool, 15 | } 16 | 17 | impl OpusEncoder { 18 | const SAMPLE_RATE: u32 = 48000; 19 | const FRAME_LEN_MS: u64 = 40; 20 | 21 | /// Creates a new instance of an OpusEncoder. 22 | pub fn new(channels: u8, kbps: u16) -> Result { 23 | let mut enc = opus::Encoder::new( 24 | Self::SAMPLE_RATE, 25 | match channels { 26 | 1 => opus::Channels::Mono, 27 | 2 => opus::Channels::Stereo, 28 | _ => bail!("invalid channels: {}", channels), 29 | }, 30 | opus::Application::LowDelay, 31 | ) 32 | .context("could not create Opus encoder")?; 33 | enc.set_bitrate(opus::Bitrate::Bits(kbps as i32 * 8)) 34 | .context("could not set Opus bitrate")?; 35 | 36 | let internal_delay = enc 37 | .get_lookahead() 38 | .context("could not get encoder internal delay")? as usize; 39 | 40 | Ok(OpusEncoder { 41 | enc, 42 | muxer: WebmStreamMuxer::new(Self::SAMPLE_RATE, channels, internal_delay), 43 | wrote_initialization_segment: false, 44 | }) 45 | } 46 | } 47 | 48 | impl Encoder for OpusEncoder { 49 | fn sample_rate(&self) -> u32 { 50 | Self::SAMPLE_RATE 51 | } 52 | 53 | fn frame_len_ms(&self) -> u64 { 54 | Self::FRAME_LEN_MS 55 | } 56 | 57 | fn encode_frame( 58 | &mut self, 59 | timestamp: u64, 60 | input: &[i16], 61 | output: &mut [u8], 62 | ) -> Result<(usize, bool)> { 63 | let keyframe = !self.wrote_initialization_segment; 64 | if keyframe { 65 | self.wrote_initialization_segment = true; 66 | } 67 | let mut encoded_chunk = [0u8; 4000]; // Opus' recommended buffer size. 68 | let encoded_chunk_length = self 69 | .enc 70 | .encode(input, &mut encoded_chunk) 71 | .context("could not encode audio data")?; 72 | log::debug!("Encoded {} audio bytes", encoded_chunk_length); 73 | let mut inner_cur = std::io::Cursor::new(output); 74 | assert_eq!(inner_cur.position(), 0); 75 | log::debug!("Encoding audio chunk..."); 76 | self.muxer 77 | .write_frame( 78 | &mut inner_cur, 79 | &encoded_chunk[..encoded_chunk_length], 80 | timestamp, 81 | keyframe, 82 | ) 83 | .context("could not write audio frame")?; 84 | Ok((inner_cur.position() as usize, keyframe)) 85 | } 86 | 87 | fn internal_delay(&mut self) -> Result { 88 | Ok(self.enc.get_lookahead()? as usize) 89 | } 90 | 91 | fn max_frame_len(&self) -> usize { 92 | 4512 93 | } 94 | } 95 | 96 | /// WebmStreamMuxer is a WebM streaming muxer compliant with [mse-byte-stream-format-webm], that can be 97 | /// played in a browser using Media Source Extensions. 98 | /// 99 | /// References: 100 | /// 101 | /// * [mse-byte-stream-format-webm] 102 | /// * 103 | /// * 104 | /// * 105 | /// * 106 | #[derive(Debug)] 107 | pub struct WebmStreamMuxer { 108 | sample_rate: u32, 109 | channels: u8, 110 | internal_delay: usize, 111 | } 112 | 113 | impl WebmStreamMuxer { 114 | pub fn new(sample_rate: u32, channels: u8, internal_delay: usize) -> WebmStreamMuxer { 115 | WebmStreamMuxer { 116 | sample_rate, 117 | channels, 118 | internal_delay, 119 | } 120 | } 121 | 122 | /// Writes the WebM Initialization Segment, which consists of: 123 | /// 124 | /// - The EBML Header 125 | /// - The Segment header (with an unknown size) 126 | /// - The SeekHead element (to comply with the WebM spec and signal the lack of a Cues element) 127 | /// - The Segment Information element 128 | /// - The Tracks element 129 | /// 130 | /// This Initialization Segment can be sent once at the beginning of every connection. 131 | pub fn write_initialization_segment(&self, w: &mut W) -> Result<()> { 132 | let mut track_uid = [0u8; 7]; 133 | rand::thread_rng() 134 | .try_fill(&mut track_uid) 135 | .context("could not generate random payload")?; 136 | 137 | let application_name = "audiomux-v0.0.0".as_bytes(); 138 | assert_eq!(application_name.len(), 15); 139 | 140 | w.write_all(&[ 141 | 0x1A, 0x45, 0xDF, 0xA3, 0x9F, 0x42, 0x86, 0x81, 0x01, 0x42, 0xF7, 0x81, 0x01, 0x42, 142 | 0xF2, 0x81, 0x04, 0x42, 0xF3, 0x81, 0x08, 0x42, 0x82, 0x84, 0x77, 0x65, 0x62, 0x6D, 143 | 0x42, 0x87, 0x81, 0x04, 0x42, 0x85, 0x81, 0x02, 144 | ])?; 145 | 146 | // Segment header, size unknown. 147 | w.write_all(&[ 148 | 0x18, 0x53, 0x80, 0x67, // SegmentHeader 149 | 0x01, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, // size = -1 150 | ])?; 151 | 152 | // SeekHead 153 | w.write_all(&[ 154 | 0x11, 0x4D, 0x9B, 0x74, // SeekHead 155 | 0xAA, // size = 42 156 | 0x4D, 0xBB, // Seek 157 | 0x8B, // size = 11 158 | 0x53, 0xAB, // SeekID 159 | 0x84, // size = 4 160 | 0x15, 0x49, 0xA9, 0x66, // KaxInfo 161 | 0x53, 0xAC, // SeekPosition 162 | 0x81, // size = 1 163 | 0x2F, // 47 164 | 0x4D, 0xBB, // Seek 165 | 0x8B, // size = 11 166 | 0x53, 0xAB, // SeekID 167 | 0x84, // size = 4 168 | 0x16, 0x54, 0xAE, 0x6B, // KaxTracks 169 | 0x53, 0xAC, // SeekPosition 170 | 0x81, // size = 1 171 | 0x5F, // 95 172 | 0x4D, 0xBB, // Seek 173 | 0x8B, // size = 11 174 | 0x53, 0xAB, // SeekID 175 | 0x84, // size = 4 176 | 0x1F, 0x43, 0xB6, 0x75, // KaxCluster 177 | 0x53, 0xAC, // SeekID 178 | 0x81, // size = 1 179 | 0xA5, // 165 180 | ])?; 181 | 182 | // Segment Information 183 | w.write_all(&[ 184 | 0x15, 185 | 0x49, 186 | 0xA9, 187 | 0x66, // Info 188 | 0xAB, // size 189 | 0x2A, 190 | 0xD7, 191 | 0xB1, // TimestampScale 192 | 0x83, // size = 3 193 | 0x0F, 194 | 0x42, 195 | 0x40, // 1_000_000 196 | 0x4D, 197 | 0x80, // Multiplexing app 198 | 0x80 | application_name.len() as u8, // size 199 | ])?; 200 | w.write_all(application_name)?; 201 | w.write_all(&[ 202 | 0x57, 203 | 0x41, // Writing app 204 | 0x80 | application_name.len() as u8, // size 205 | ])?; 206 | w.write_all(application_name)?; 207 | 208 | // Tracks 209 | w.write_all(&[ 210 | 0x16, 211 | 0x54, 212 | 0xAE, 213 | 0x6B, // Tracks 214 | 0xC2, // Size = 66 215 | 0xAE, // TrackEntry 216 | 0xC0, // Size = 64 217 | 0xD7, // TrackNumber 218 | 0x81, // Size = 1 219 | 0x01, // Track = 1 220 | 0x73, 221 | 0xC5, // TrackUID 222 | 0x80 | track_uid.len() as u8, // Size 223 | ])?; 224 | w.write_all(&track_uid)?; 225 | let codec_delay_ns = self.internal_delay * 1_000_000_000 / self.sample_rate as usize; 226 | w.write_all(&[ 227 | 0x83, // TrackType 228 | 0x81, // Size = 1 229 | 0x02, // Audio 230 | 0x86, // CodecID 231 | 0x86, // Size = 6 232 | 0x41, 233 | 0x5F, 234 | 0x4F, 235 | 0x50, 236 | 0x55, 237 | 0x53, // "A_OPUS" 238 | 0x56, 239 | 0xAA, // CodecDelay 240 | 0x84, // Size = 4 241 | ((codec_delay_ns >> 24) & 0xff) as u8, 242 | ((codec_delay_ns >> 16) & 0xff) as u8, 243 | ((codec_delay_ns >> 8) & 0xff) as u8, 244 | (codec_delay_ns & 0xff) as u8, 245 | 0x63, 246 | 0xA2, // CodecPrivate 247 | 0x93, // Size = 19 248 | 0x4F, 249 | 0x70, 250 | 0x75, 251 | 0x73, 252 | 0x48, 253 | 0x65, 254 | 0x61, 255 | 0x64, // "OpusHead" 256 | 0x01, // version 257 | self.channels, 258 | (self.internal_delay & 0xff) as u8, 259 | ((self.internal_delay >> 8) & 0xff) as u8, 260 | (self.sample_rate & 0xff) as u8, 261 | ((self.sample_rate >> 8) & 0xff) as u8, 262 | ((self.sample_rate >> 16) & 0xff) as u8, 263 | ((self.sample_rate >> 24) & 0xff) as u8, 264 | 0x00, 265 | 0x00, // 0 output gain, le16s 266 | 0x00, // Channel Mapping Family 8u 267 | 0xE1, // Audio 268 | 0x89, // Size = 9 269 | 0xB5, // SamplingFrequency 270 | 0x84, // Size = 4 271 | 0x47, 272 | 0x3B, 273 | 0x80, 274 | 0x00, // IEEE-745: 48000.0f 275 | 0x9F, // Channels 276 | 0x81, // Size = 1 277 | self.channels, // Channels = 2 278 | ])?; 279 | 280 | Ok(()) 281 | } 282 | } 283 | 284 | impl StreamMuxer for WebmStreamMuxer { 285 | /// Writes the an Opus-encoded packet in a Cluster. Ideally, we could use a ShortBlock instead 286 | /// of a whole Cluster, but this makes things a bit easier for the time being. 287 | fn write_frame( 288 | &mut self, 289 | w: &mut W, 290 | frame: &[u8], 291 | timestamp: u64, 292 | keyframe: bool, 293 | ) -> Result<()> { 294 | if keyframe { 295 | self.write_initialization_segment(w)?; 296 | } 297 | 298 | let cluster_len = frame.len() + 17; 299 | let simple_block_len = frame.len() + 4; 300 | // Cluster 301 | w.write_all(&[ 302 | 0x1F, 303 | 0x43, 304 | 0xB6, 305 | 0x75, // Cluster 306 | 0x40 | ((cluster_len) >> 8) as u8, 307 | (cluster_len & 0xff) as u8, // Size 308 | 0xE7, // Timestamp 309 | 0x88, // Size = 8 310 | ((timestamp >> 56) & 0xFF) as u8, 311 | ((timestamp >> 48) & 0xFF) as u8, 312 | ((timestamp >> 40) & 0xFF) as u8, 313 | ((timestamp >> 32) & 0xFF) as u8, 314 | ((timestamp >> 24) & 0xFF) as u8, 315 | ((timestamp >> 16) & 0xFF) as u8, 316 | ((timestamp >> 8) & 0xFF) as u8, 317 | (timestamp & 0xFF) as u8, 318 | 0xA3, // SimpleBlock 319 | 0x40 | ((simple_block_len) >> 8) as u8, 320 | (simple_block_len & 0xff) as u8, // size 321 | // SimpleBlock structure 322 | 0x81, // Track = 1 323 | 0x00, 324 | 0x00, // timecode = 0x0000 325 | 0x80, 326 | ])?; 327 | w.write_all(frame)?; 328 | Ok(()) 329 | } 330 | } 331 | 332 | #[cfg(test)] 333 | mod tests { 334 | use super::*; 335 | use crate::audio::Encoder; 336 | 337 | use std::io::Write; 338 | 339 | fn init() { 340 | let _ = env_logger::builder().is_test(true).try_init(); 341 | } 342 | 343 | #[test] 344 | fn test_webm() -> Result<()> { 345 | init(); 346 | 347 | let mut enc = OpusEncoder::new(2, 32 * 1024).expect("could not create Opus encoder"); 348 | 349 | let mut chunk = 350 | vec![ 351 | 0i16; 352 | enc.channels() as usize * enc.sample_rate() as usize * enc.frame_len_ms() as usize 353 | / 1000 354 | ]; 355 | rand::thread_rng() 356 | .try_fill(&mut chunk[..]) 357 | .expect("could not generate random payload"); 358 | 359 | log::info!("Read audio chunk"); 360 | 361 | let mut payload = Vec::::with_capacity(enc.max_frame_len()); 362 | payload.resize(enc.max_frame_len(), 0); 363 | payload[..4].copy_from_slice(&[ 364 | 0xFF, // message-type 365 | 0x01, // submessage-type 366 | 0x00, 0x02, // operation 367 | ]); 368 | 369 | log::info!("Encoding audio chunk..."); 370 | 371 | let mut payload_size = enc 372 | .encode_frame(0, &chunk, &mut payload[8..]) 373 | .expect("could not encode audio data") 374 | .0; 375 | payload_size += enc 376 | .encode_frame(0, &chunk, &mut payload[8 + payload_size..]) 377 | .expect("could not encode audio data") 378 | .0; 379 | log::info!("Wrote {} audio bytes!", payload_size); 380 | 381 | assert!(payload_size > 0); 382 | 383 | payload[4] = ((payload_size >> 24) & 0xff) as u8; 384 | payload[5] = ((payload_size >> 16) & 0xff) as u8; 385 | payload[6] = ((payload_size >> 8) & 0xff) as u8; 386 | payload[7] = ((payload_size >> 0) & 0xff) as u8; 387 | payload.truncate(payload_size + 8); 388 | assert!(payload.len() > 5); 389 | 390 | let mut file = tempfile::NamedTempFile::new()?; 391 | file.write_all(&payload[8..payload_size])?; 392 | 393 | let mkv = 394 | matroska::Matroska::open(file.reopen().context("failed to reopen temporary file")?) 395 | .context("failed to open matroska file")?; 396 | assert_eq!(mkv.info.duration, None); 397 | assert_eq!(mkv.tracks.len(), 1); 398 | assert_eq!(mkv.tracks[0].tracktype, matroska::Tracktype::Audio); 399 | assert_eq!(mkv.tracks[0].codec_id, "A_OPUS"); 400 | if let matroska::Settings::Audio(audio) = &mkv.tracks[0].settings { 401 | assert_eq!(audio.sample_rate, enc.sample_rate() as f64); 402 | assert_eq!(audio.channels, enc.channels() as u64); 403 | } else { 404 | assert!( 405 | false, 406 | "failed to parse audio track: {:?}", 407 | mkv.tracks[0].settings 408 | ); 409 | } 410 | 411 | Ok(()) 412 | } 413 | } 414 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # rfbproxy 2 | 3 | An RFB proxy that enables WebSockets and audio. 4 | 5 | This crate proxies a TCP Remote Framebuffer server connection and exposes a 6 | WebSocket endpoint, translating the connection between them. It can optionally 7 | enable audio using the [Replit Audio RFB extension](#replit-audio-rfb-extension) if the 8 | `--enable-audio` flag is passed or the `VNC_ENABLE_EXPERIMENTAL_AUDIO` 9 | environment variable is set to a non-empty value. 10 | 11 | Note that since this project only supports RFB-over-WebSockets 12 | the only VNC client that supports that particular combination and 13 | the only one that will work with it is [noVNC](https://novnc.com). 14 | 15 | # Running 16 | 17 | Since this is a proxy, you'll need to have an RFB server running 18 | already. [TigerVNC](https://tigervnc.org/) is a good option: 19 | 20 | ```shell 21 | Xvnc --SecurityTypes={None,VNCAuth} --rfbport=5901 --localhost :1 22 | ``` 23 | 24 | Now `rfbproxy` can run: 25 | 26 | ```shell 27 | cargo run -- [--enable-audio] [--address=0.0.0.0:5900] [--rfb-server=127.0.0.1:5901] 28 | ``` 29 | 30 | # Replit Audio RFB extension 31 | 32 | This uses a proposed extension to the [RFB 33 | protocol](https://github.com/rfbproto/rfbproto/blob/master/rfbproto.rst) in 34 | order to negotiate and transmit encoded audio. This is the main difference from 35 | the pre-existing [QEMU Audio 36 | messages](https://github.com/rfbproto/rfbproto/blob/master/rfbproto.rst#qemu-audio-client-message). 37 | 38 | ## Encodings 39 | 40 | This registers the following pseudo-encodings: 41 | 42 | | Number | Name | 43 | |----------- | ------------------------------| 44 | | 0x52706C41 | Replit Audio Pseudo-encoding | 45 | 46 | A client that supports this encoding is indicating that it is able to receive 47 | an encoded audio data stream. If a server wishes to send encoded audio data, it 48 | will send a pseudo-rectangle with the following contents: 49 | 50 | | No. of bytes | Type | Description | 51 | |------------------------|-------------|--------------------| 52 | | 2 | `U16` | _version_ | 53 | | 2 | `U16` | _number-of-codecs_ | 54 | | 2 * _number-of-codecs_ | `U16` array | _codecs_ | 55 | 56 | The supported codecs are as follow: 57 | 58 | | Codec | Description | 59 | |-------|-----------------------------| 60 | | 0 | Opus codec, WebM container | 61 | | 1 | MP3 codec, MPEG-1 container | 62 | 63 | After receiving this notification, clients may optionally use the [Replit Audio 64 | Client Message](#replit-audio-client-client-to-server-messages). 65 | 66 | ## Client to Server Messages 67 | 68 | This registers the following message types: 69 | 70 | | Number | Name | 71 | |--------|-----------------------------| 72 | | 245 | Replit Audio Client Message | 73 | 74 | This message may only be sent if the client has previously received a 75 | _FrameBufferUpdate_ that confirms support for the intended message-type. Every 76 | `Replit Audio Client Message` begins with a standard header 77 | 78 | | No. of bytes | Type | [Value] | Description | 79 | |--------------|-------|---------|-------------------| 80 | | 1 | `U8` | 245 | _message-type_ | 81 | | 1 | `U8` | | _submessage-type_ | 82 | | 2 | `U16` | | _payload-length_ | 83 | 84 | This header is then followed by arbitrary data of length _payload-length_, and 85 | whose format is determined by the _submessage-type_. Possible values for 86 | _submessage-type_ and their associated minimum versions are 87 | 88 | | Submessage Type | Minimum version | Description | 89 | |-----------------|-----------------|-----------------------------------------------------------------------------------| 90 | | 0 | 0 | [Start Encoder](#replit-audio-client-start-encoder-message) | 91 | | 1 | 0 | [Frame Request](#replit-audio-client-frame-request-message) | 92 | | 2 | 0 | [Start Continuous Updates](#replit-audio-client-start-continuous-updates-message) | 93 | 94 | ### Replit Audio Client Start Encoder Message 95 | 96 | This submessage allows the client to request the server to start audio capture 97 | with the provided configuration 98 | 99 | | No. of bytes | Type | [Value] | Description | 100 | |--------------|-------|---------|-------------------| 101 | | 1 | `U8` | 245 | _message-type_ | 102 | | 1 | `U8` | 0 | _submessage-type_ | 103 | | 2 | `U16` | 6 | _payload-length_ | 104 | | 1 | `U8` | | _enabled_ | 105 | | 1 | `U8` | | _channels_ | 106 | | 2 | `U16` | | _codec_ | 107 | | 2 | `U16` | | _kbytes_per_sec_ | 108 | 109 | After invoking this operation, the client will receive a [Replit Audio Server 110 | Start Encoder Message](#replit-audio-server-start-encoder-message) with the 111 | result of the operation. 112 | 113 | Valid values for the _enabled_ field are 0, which disables/stops the audio 114 | encoder, and 1, which starts the audio encoder. Valid values for the _channels_ 115 | field are 1 (Mono audio) and 2 (Stereo audio). Valid values for the _codec_ 116 | field are the ones sent by the server in the [Replit Audio 117 | Pseudo-encoding](#encodings) pseudo-rect. Valid values for the _kbytes_per_sec_ 118 | field are codec-dependent. The Opus codec achieves good performance with 32, 119 | whereas the MP3 codec might require 128 for a comparable experience. 120 | 121 | ### Replit Audio Client Frame Request Message 122 | 123 | This submessage allows the client to request the server for a single audio 124 | frame. The length of an audio frame is codec-dependent, but is typically 125 | between 5 and 40 milliseconds. Each frame is encoded with the parameters chosen 126 | by the [Start Encoder](#replit-audio-client-start-encoder-message) message. 127 | The client MUST send a [Start 128 | Encoder](#replit-audio-client-start-encoder-message) message and have received 129 | acknowledgement from the server that the chosen parameters are valid prior to 130 | sending this message. 131 | 132 | | No. of bytes | Type | [Value] | Description | 133 | |--------------|-------|---------|-------------------| 134 | | 1 | `U8` | 245 | _message-type_ | 135 | | 1 | `U8` | 1 | _submessage-type_ | 136 | | 2 | `U16` | 0 | _payload-length_ | 137 | 138 | After invoking this operation, the client will receive a [Replit Audio Server 139 | Frame Response Message](#replit-audio-server-frame-response-message) with the 140 | encoded audio frame in the corresponding container format. 141 | 142 | ### Replit Audio Client Start Continuous Updates Message 143 | 144 | This submessage allows the client to request the server send audio frames 145 | continuously, which saves bandwidth and reduces audio latency incurred by the 146 | TCP stack by half compared to requesting frames individually. The length of an 147 | audio frame is codec-dependent, but is typically between 5 and 40 milliseconds. 148 | Each frame is encoded with the parameters chosen by the [Start 149 | Encoder](#replit-audio-client-start-encoder-message) message. The client MUST 150 | send a [Start Encoder](#replit-audio-client-start-encoder-message) message and 151 | have received acknowledgement from the server that the chosen parameters are 152 | valid prior to sending this message. 153 | 154 | | No. of bytes | Type | [Value] | Description | 155 | |--------------|-------|---------|-------------------| 156 | | 1 | `U8` | 245 | _message-type_ | 157 | | 1 | `U8` | 1 | _submessage-type_ | 158 | | 2 | `U16` | 0 | _payload-length_ | 159 | 160 | After invoking this operation, the client will receive a [Replit Audio Server 161 | Start Continuous Updates 162 | Message](#replit-audio-server-start-continuous-updates-message) with the result 163 | of the operation. If the operation was successful, that message will be 164 | followed by [Replit Audio Server Frame Response 165 | Message](#replit-audio-server-frame-response-message) messages with and encoded 166 | audio frame in the corresponding container format. 167 | 168 | Once audio frames start being continuously sent, this can be stopped by sending 169 | a [Start Encoder](#replit-audio-client-start-encoder-message) message with the 170 | _enabled_ field set to `0`. Due to inherent race conditions in the protocol, 171 | after disabling the encoder, the client may still receive further [Replit Audio 172 | Server Frame Response Message](#replit-audio-server-frame-response-message) 173 | messages, but once the server acknowledges the receipt of the [Start 174 | Encoder](#replit-audio-client-start-encoder-message) message, no further audio 175 | frames will be sent. 176 | 177 | ## Server to Client Messages 178 | 179 | This registers the following message types: 180 | 181 | | Number | Name | 182 | |--------|-----------------------------| 183 | | 245 | Replit Audio Server Message | 184 | 185 | This message may only be sent if the client has previously sent a [Replit Audio 186 | Client Message](#replit-audio-client-message) that confirms support for the 187 | intended message-type. Every `Replit Audio Server Message` begins with a 188 | standard header 189 | 190 | | No. of bytes | Type | [Value] | Description | 191 | |--------------|-------|---------|-------------------| 192 | | 1 | `U8` | 245 | _message-type_ | 193 | | 1 | `U8` | | _submessage-type_ | 194 | | 2 | `U16` | | _payload-length_ | 195 | 196 | This header is then followed by arbitrary data of length _payload-length_, and 197 | whose format is determined by the _submessage-type_. Possible values for 198 | _submessage-type_ and their associated minimum versions are 199 | 200 | | Submessage Type | Minimum version | Description | 201 | |-----------------|-----------------|-----------------------------------------------------------------------------------| 202 | | 0 | 0 | [Start Encoder](#replit-audio-server-start-encoder-message) | 203 | | 1 | 0 | [Frame Request](#replit-audio-server-frame-request-message) | 204 | | 2 | 0 | [Start Continuous Updates](#replit-audio-server-start-continuous-updates-message) | 205 | 206 | ### Replit Audio Server Start Encoder Message 207 | 208 | This submessage is a response to the [Replit Audio Client Start Encoder 209 | Message](#replit-audio-client-start-encoder-message), and acknowledges the 210 | receipt and/or support for the requested configuration 211 | 212 | | No. of bytes | Type | [Value] | Description | 213 | |--------------|-------|---------|-------------------| 214 | | 1 | `U8` | 245 | _message-type_ | 215 | | 1 | `U8` | 0 | _submessage-type_ | 216 | | 2 | `U16` | 1 | _payload-length_ | 217 | | 1 | `U8` | | _enabled_ | 218 | 219 | If the parameters in the [Replit Audio Client Start Encoder 220 | Message](#replit-audio-client-start-encoder-message) were valid and the server 221 | was able to successfully start an audio capture session, the value of _enabled_ 222 | will be 1. Otherwise it will be 0. 223 | 224 | After receiveing this message with _enabled_ set to 1, the client can send 225 | other [Replit Audio Client Message](#client-to-server-messages) messages. 226 | 227 | ### Replit Audio Server Frame Request Message 228 | 229 | This submessage contains audio data for a single audio frame wrapped in the 230 | container format associated with it. The length of an audio frame is 231 | codec-dependent, but is typically between 5 and 40 milliseconds. The frame is 232 | encoded with the parameters chosen by the [Start 233 | Encoder](#replit-audio-client-start-encoder-message) message. This is a 234 | response to either the [Replit Audio Client Frame Request 235 | Message](#replit-audio-client-frame-request-message) or the [Replit Audio 236 | Client Start Continuous Updates 237 | Message](#replit-audio-client-start-continuous-updates-message). 238 | 239 | | No. of bytes | Type | [Value] | Description | 240 | |----------------|------------|-------------------|-------------------| 241 | | 1 | `U8` | 245 | _message-type_ | 242 | | 1 | `U8` | 1 | _submessage-type_ | 243 | | 2 | `U16` | 4 + _data-length_ | _payload-length_ | 244 | | 4 | `U32` | | _timestamp_ | 245 | | _data-length_ | `U8` array | | _data_ | 246 | 247 | The most significant bit of _timestamp_ denotes whether the audio frame 248 | contains a start-of-stream header or is otherwise a keyframe, which enables 249 | clients to use this information for seeking purposes. Servers SHOULD send 250 | keyframes every few seconds / minutes to allow clients to re-synchronize with 251 | the stream. The 31 least significant bits of _timestamp_ contain the number of 252 | milliseconds from the first audio frame that was captured in the session since 253 | the [Start Encoder](#replit-audio-client-start-encoder-message) message was 254 | acknowledged by the server. _data_ SHOULD be a self-contained audio frame, and 255 | all the audio frames should be concatenable into a valid audio stream. 256 | Furthermore, dropping of a non-keyframe SHOULD not cause the client to 257 | de-synchronize, and SHOULD be recoverable by inserting silence for the duration 258 | of the dropped frame. 259 | 260 | ### Replit Audio Server Start Continuous Updates Message 261 | 262 | This submessage is a response to the [Replit Audio Client Start Continuous 263 | Updates Message](#replit-audio-client-start-continuous-updates-message), and 264 | acknowledges the receipt of it and signals the client that the server will send 265 | [Replit Audio Server Frame Request 266 | Message](#replit-audio-server-frame-request-message) messages continuously. 267 | 268 | | No. of bytes | Type | [Value] | Description | 269 | |--------------|-------|---------|-------------------| 270 | | 1 | `U8` | 245 | _message-type_ | 271 | | 1 | `U8` | 0 | _submessage-type_ | 272 | | 2 | `U16` | 1 | _payload-length_ | 273 | | 1 | `U8` | | _enabled_ | 274 | 275 | _enabled_ will be set to 1 when the stream of [Replit Audio Server Frame 276 | Request Message](#replit-audio-server-frame-request-message) messages will 277 | start. _enabled_ will be set to 0 if client had not sent a [Start 278 | Encoder](#replit-audio-client-start-encoder-message) message beforehand, or if 279 | there was any other problem starting the stream. If there is an error at any 280 | future point, or if the client sent a [Start 281 | Encoder](#replit-audio-client-start-encoder-message) with the _enabled_ field 282 | set to 0, the server will send an additional [Replit Audio Server Start 283 | Continuous Updates 284 | Message](#replit-audio-server-start-continuous-updates-message) with _enabled_ 285 | set to 0 after sending the last audio frame. 286 | -------------------------------------------------------------------------------- /src/main.rs: -------------------------------------------------------------------------------- 1 | //! An RFB proxy that enables WebSockets and audio. 2 | //! 3 | //! This crate proxies a TCP Remote Framebuffer server connection and exposes a WebSocket endpoint, 4 | //! translating the connection between them. It can optionally enable audio using the Replit Audio 5 | //! messages if the `--enable-audio` flag is passed or the `VNC_ENABLE_EXPERIMENTAL_AUDIO` 6 | //! environment variable is set to a non-empty value. 7 | 8 | mod audio; 9 | mod auth; 10 | mod messages; 11 | mod rfb; 12 | 13 | use std::collections::HashMap; 14 | use std::net::SocketAddr; 15 | use std::sync::Arc; 16 | 17 | use anyhow::{bail, Context, Result}; 18 | 19 | use futures::{SinkExt, StreamExt}; 20 | 21 | use hyper::{Body, Request, Response, Server}; 22 | 23 | use path_clean::PathClean; 24 | 25 | use tokio::io::{AsyncReadExt, AsyncWriteExt}; 26 | use tokio::net::{TcpListener, TcpStream}; 27 | use tokio::sync::mpsc; 28 | use tokio_tungstenite::tungstenite; 29 | use tokio_tungstenite::tungstenite::protocol::Message; 30 | 31 | /// The protobuf definitions. 32 | #[allow(clippy::derive_partial_eq_without_eq)] 33 | mod api { 34 | include!(concat!(env!("OUT_DIR"), "/api.rs")); 35 | } 36 | 37 | /// Forwards the data between `socket` and `ws_stream`. Doesn't do anything with the bytes. 38 | async fn forward_streams( 39 | mut socket: TcpStream, 40 | ws_stream: tokio_tungstenite::WebSocketStream, 41 | ) -> Result<()> 42 | where 43 | Stream: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, 44 | { 45 | let (mut wws, mut rws) = ws_stream.split(); 46 | let (mut rs, mut ws) = socket.split(); 47 | 48 | let client_to_server = async move { 49 | while let Some(msg) = rws.next().await { 50 | if let Ok(Message::Binary(payload)) = msg { 51 | if let Err(err) = ws.write_all(&payload).await { 52 | log::error!("failed to write a message to the server: {:#}", err); 53 | break; 54 | } 55 | } 56 | } 57 | 58 | log::info!("client disconnected"); 59 | ws.shutdown().await?; 60 | Ok::<(), anyhow::Error>(()) 61 | }; 62 | 63 | let server_to_client = async move { 64 | let mut buffer = [0u8; 4096]; 65 | loop { 66 | match rs.read(&mut buffer[..]).await { 67 | Ok(0) => { 68 | break; 69 | } 70 | Ok(n) => { 71 | if let Err(err) = wws.send(Message::Binary((buffer[..n]).to_vec())).await { 72 | log::error!("failed to write a message to the client: {:#}", err); 73 | break; 74 | } 75 | } 76 | Err(err) => { 77 | log::error!("failed to read a message from the server: {:#}", err); 78 | break; 79 | } 80 | } 81 | } 82 | 83 | log::info!("server disconnected"); 84 | wws.close().await?; 85 | Ok::<(), anyhow::Error>(()) 86 | }; 87 | 88 | let (cts, stc) = tokio::join!(client_to_server, server_to_client); 89 | cts?; 90 | stc?; 91 | Ok(()) 92 | } 93 | 94 | /// Handles a single WebSocket connection. If `enable_audio` is false, it will just forward the 95 | /// data between them. Otherwise, it will parse and interpret each RFB packet and inject audio 96 | /// data. 97 | async fn handle_connection( 98 | rfb_addr: std::net::SocketAddr, 99 | mut ws_stream: tokio_tungstenite::WebSocketStream, 100 | authentication: &auth::RfbAuthentication, 101 | enable_audio: bool, 102 | ) -> Result<()> 103 | where 104 | Stream: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send, 105 | { 106 | let mut socket = TcpStream::connect(rfb_addr).await?; 107 | auth::authenticate(authentication, &mut socket, &mut ws_stream).await?; 108 | 109 | if !enable_audio { 110 | return forward_streams(socket, ws_stream).await; 111 | } 112 | let (server_tx, mut server_rx) = mpsc::channel(2); 113 | let (client_tx, mut client_rx) = mpsc::channel(2); 114 | let mut conn = rfb::RfbConnection::new(socket, &mut ws_stream, server_tx, client_tx).await?; 115 | 116 | let (mut wws, mut rws) = ws_stream.split(); 117 | let (mut rs, mut ws) = conn.split(); 118 | 119 | let client_to_server = async { 120 | loop { 121 | let payload = tokio::select! { 122 | Some(payload) = server_rx.recv() => Some(payload), 123 | Some(msg) = rws.next() => { 124 | match msg.context("failed to read client-to-server message")? { 125 | Message::Binary(payload) => Some(payload), 126 | Message::Close(_) => break, 127 | msg => { 128 | log::debug!(" ->: Received a message {:?}", msg); 129 | None 130 | } 131 | } 132 | }, 133 | else => break, 134 | }; 135 | 136 | if let Some(payload) = payload { 137 | if let Err(err) = ws.write_all(&payload).await { 138 | log::error!("failed to write message: {:#}", err); 139 | break; 140 | } 141 | } 142 | } 143 | 144 | log::info!("client disconnected"); 145 | ws.shutdown().await?; 146 | Ok::<(), anyhow::Error>(()) 147 | }; 148 | 149 | let server_to_client = async { 150 | loop { 151 | let payload = tokio::select! { 152 | Some(payload) = client_rx.recv() => Some(payload), 153 | message = rs.read_server_message() => { 154 | match message.context("failed to read server-to-client message")? { 155 | None => break, 156 | Some(msg) => { 157 | log::debug!("<-: {:?}", &msg); 158 | Some(msg.into_data()) 159 | } 160 | } 161 | }, 162 | else => break, 163 | }; 164 | 165 | if let Some(payload) = payload { 166 | wws.send(Message::Binary(payload)).await?; 167 | } 168 | } 169 | log::info!("server disconnected"); 170 | wws.close().await?; 171 | Ok::<(), anyhow::Error>(()) 172 | }; 173 | 174 | let (cts, stc) = tokio::join!(client_to_server, server_to_client); 175 | cts?; 176 | stc?; 177 | 178 | Ok(()) 179 | } 180 | 181 | /// Handles HTTP requests. Will serve any files in the current working directory, and the RFB 182 | /// websocket in `/ws`. 183 | async fn handle_request( 184 | rfb_addr: std::net::SocketAddr, 185 | mut req: Request, 186 | remote_addr: SocketAddr, 187 | authentication: Arc, 188 | enable_audio: bool, 189 | ) -> Result> { 190 | // Clean the path so that it can't be used to access files outside the current working 191 | // directory. 192 | *req.uri_mut() = { 193 | let uri = req.uri(); 194 | let clean_path = String::from( 195 | std::path::PathBuf::from(uri.path()) 196 | .clean() 197 | .to_str() 198 | .with_context(|| format!("failed to clean path: {:?}", uri.path()))?, 199 | ); 200 | 201 | let mut builder = http::uri::Builder::new(); 202 | if let Some(scheme) = uri.scheme() { 203 | builder = builder.scheme(scheme.as_str()); 204 | } 205 | if let Some(authority) = uri.authority() { 206 | builder = builder.authority(authority.as_str()); 207 | } 208 | if let Some(query) = uri.query() { 209 | builder = builder.path_and_query(format!("{clean_path}?{query}")); 210 | } else { 211 | builder = builder.path_and_query(clean_path); 212 | } 213 | builder.build()? 214 | }; 215 | 216 | if req.uri().path() == "/ws" { 217 | if !hyper_tungstenite::is_upgrade_request(&req) { 218 | log::info!("Not an upgrade request"); 219 | return Ok(http::response::Builder::new() 220 | .status(http::StatusCode::NOT_FOUND) 221 | .body(Body::empty()) 222 | .expect("unable to build response")); 223 | } 224 | 225 | let (response, websocket) = hyper_tungstenite::upgrade(req, None)?; 226 | 227 | tokio::spawn(async move { 228 | log::info!("Incoming TCP connection from: {}", remote_addr); 229 | 230 | let ws_stream = match websocket.await { 231 | Ok(ws_stream) => ws_stream, 232 | Err(e) => { 233 | log::error!("error in websocket upgrade: {:#}", e); 234 | return; 235 | } 236 | }; 237 | if let Err(e) = 238 | handle_connection(rfb_addr, ws_stream, &authentication, enable_audio).await 239 | { 240 | log::error!("error in websocket connection: {:#}", e); 241 | } 242 | log::info!("{} disconnected", remote_addr); 243 | }); 244 | return Ok(response); 245 | } 246 | 247 | Ok(hyper_staticfile::Static::new(std::path::Path::new("./")) 248 | .serve(req) 249 | .await?) 250 | } 251 | 252 | #[doc(hidden)] 253 | #[tokio::main] 254 | async fn main() -> Result<()> { 255 | env_logger::init(); 256 | 257 | let matches = clap::App::new("rfbproxy") 258 | .about("An RFB proxy that enables WebSockets and audio") 259 | .arg( 260 | clap::Arg::with_name("address") 261 | .long("address") 262 | .value_name("HOST:PORT") 263 | .default_value("0.0.0.0:5900") 264 | .help("The hostname and port in which the server will bind") 265 | .takes_value(true), 266 | ) 267 | .arg( 268 | clap::Arg::with_name("rfb-server") 269 | .long("rfb-server") 270 | .value_name("HOST:PORT") 271 | .default_value("127.0.0.1:5901") 272 | .help("The hostname and port where the original RFB server is listening") 273 | .takes_value(true), 274 | ) 275 | .arg( 276 | clap::Arg::with_name("http-server") 277 | .long("http-server") 278 | .help( 279 | "Whether a normal HTTP server will start to serve the current directory's contents", 280 | ), 281 | ) 282 | .arg( 283 | clap::Arg::with_name("enable-audio") 284 | .long("enable-audio") 285 | .help("Whether the muxer will support audio muxing or be a simple WebSocket proxy"), 286 | ) 287 | .arg( 288 | clap::Arg::with_name("replid") 289 | .long("replid") 290 | .takes_value(true) 291 | .help("The ID of the Repl. Used for authentication"), 292 | ) 293 | .arg( 294 | clap::Arg::with_name("pubkeys") 295 | .long("pubkeys") 296 | .takes_value(true) 297 | .help("A JSON-encoded mapping of key IDs to base64-encoded ed25519 public keys"), 298 | ) 299 | .get_matches(); 300 | 301 | if matches.value_of("replid").is_some() != matches.value_of("pubkeys").is_some() { 302 | bail!("--replid and --pubkeys must be passed together"); 303 | } 304 | 305 | // Create the event loop and TCP listener we'll accept connections on. 306 | let local_addr = matches 307 | .value_of("address") 308 | .context("missing --address arg")?; 309 | let rfb_addr: std::net::SocketAddr = matches 310 | .value_of("rfb-server") 311 | .context("missing --rfb-server arg")? 312 | .parse()?; 313 | let enable_audio = matches.is_present("enable-audio") 314 | || !std::env::var("VNC_ENABLE_EXPERIMENTAL_AUDIO") 315 | .unwrap_or_else(|_| String::new()) 316 | .is_empty(); 317 | let authentication = if matches.value_of("replid").is_some() { 318 | let mut pubkeys_base64: HashMap = 319 | serde_json::from_str(matches.value_of("pubkeys").unwrap().trim())?; 320 | let pubkeys = pubkeys_base64 321 | .drain() 322 | .map(|(keyid, pubkey)| { 323 | let pubkey = base64::decode(pubkey) 324 | .with_context(|| anyhow::anyhow!("failed to parse pubkey {}", &keyid))?; 325 | Ok((keyid, pubkey)) 326 | }) 327 | .collect::>, anyhow::Error>>()?; 328 | Arc::new(auth::RfbAuthentication::Replit { 329 | replid: matches.value_of("replid").unwrap().to_string(), 330 | pubkeys, 331 | }) 332 | } else if enable_audio { 333 | Arc::new(auth::RfbAuthentication::Passthrough) 334 | } else { 335 | // If both audio and the replit authentications are disabled, we can let the server and 336 | // client talk directly to each other without interfering since we don't need to parse any 337 | // of the messages. 338 | Arc::new(auth::RfbAuthentication::Null) 339 | }; 340 | 341 | if matches.is_present("http-server") { 342 | let server = Server::bind(&local_addr.parse()?).serve(hyper::service::make_service_fn( 343 | |conn: &hyper::server::conn::AddrStream| { 344 | let remote_addr = conn.remote_addr(); 345 | let authentication = authentication.clone(); 346 | async move { 347 | Ok::<_, hyper::Error>(hyper::service::service_fn(move |req: Request| { 348 | let authentication = authentication.clone(); 349 | async move { 350 | handle_request( 351 | rfb_addr, 352 | req, 353 | remote_addr, 354 | authentication.clone(), 355 | enable_audio, 356 | ) 357 | .await 358 | } 359 | })) 360 | } 361 | }, 362 | )); 363 | log::info!("Listening on: {}", local_addr); 364 | 365 | server.await?; 366 | } else { 367 | let listener = TcpListener::bind(&local_addr).await?; 368 | log::info!("Listening on: {}", local_addr); 369 | 370 | while let Ok((raw_stream, remote_addr)) = listener.accept().await { 371 | let ws_stream = match tokio_tungstenite::accept_hdr_async( 372 | raw_stream, 373 | |request: &tungstenite::handshake::server::Request, 374 | mut response: tungstenite::handshake::server::Response| { 375 | const PROTOCOL_HEADER: &str = "Sec-WebSocket-Protocol"; 376 | if let Some(val) = request.headers().get(PROTOCOL_HEADER) { 377 | response.headers_mut().insert(PROTOCOL_HEADER, val.clone()); 378 | } 379 | Ok(response) 380 | }, 381 | ) 382 | .await 383 | { 384 | Ok(ws_stream) => ws_stream, 385 | Err(e) => { 386 | log::error!("error in websocket upgrade: {:#}", e); 387 | continue; 388 | } 389 | }; 390 | let authentication = authentication.clone(); 391 | tokio::spawn(async move { 392 | log::info!("Incoming TCP connection from: {}", remote_addr); 393 | if let Err(e) = 394 | handle_connection(rfb_addr, ws_stream, &authentication, enable_audio).await 395 | { 396 | log::error!("error in websocket connection: {:#}", e); 397 | } 398 | log::info!("{} disconnected", remote_addr); 399 | }); 400 | } 401 | } 402 | 403 | Ok(()) 404 | } 405 | -------------------------------------------------------------------------------- /src/rfb.rs: -------------------------------------------------------------------------------- 1 | //! A Remote Framebuffer connection that can inject audio messages into the stream. 2 | //! 3 | //! This connection also translates from a TCP socket (server) to a Websocket (client) for use with 4 | //! noVNC. 5 | 6 | use crate::audio; 7 | use crate::messages; 8 | use crate::messages::{client, server}; 9 | 10 | use std::sync::atomic::{AtomicUsize, Ordering}; 11 | use std::sync::Arc; 12 | 13 | use anyhow::{bail, Context, Result}; 14 | 15 | use bytes::{Buf, BytesMut}; 16 | 17 | use futures::{SinkExt, StreamExt}; 18 | use tokio::io::{AsyncReadExt, AsyncWriteExt}; 19 | use tokio::sync::{mpsc, oneshot}; 20 | 21 | /// The shared state between the RFB client and server. 22 | #[derive(Debug)] 23 | pub struct RfbConnectionState { 24 | bytes_per_pixel: AtomicUsize, 25 | } 26 | 27 | /// A Remote Framebuffer connection between a TCP server and a Websocket client. 28 | #[derive(Debug)] 29 | #[allow(dead_code)] 30 | pub struct RfbConnection { 31 | stream: tokio::net::TcpStream, 32 | connection_state: Arc, 33 | server_tx: mpsc::Sender>, 34 | client_tx: mpsc::Sender>, 35 | } 36 | 37 | impl RfbConnection { 38 | pub async fn new( 39 | mut stream: tokio::net::TcpStream, 40 | ws_stream: &mut tokio_tungstenite::WebSocketStream, 41 | server_tx: mpsc::Sender>, 42 | client_tx: mpsc::Sender>, 43 | ) -> Result 44 | where 45 | Stream: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin, 46 | { 47 | let mut buf = [0u8; 1024]; 48 | 49 | // ClientInit 50 | match ws_stream.next().await { 51 | Some(msg) => match msg.context("bad ClientInit handshake")? { 52 | tokio_tungstenite::tungstenite::protocol::Message::Binary(payload) => { 53 | log::debug!("->: {:?}", &payload); 54 | stream.write_all(&payload).await?; 55 | } 56 | unexpected_msg => bail!("unexpected message {:?}", unexpected_msg), 57 | }, 58 | None => bail!("missing ClientInit handshake"), 59 | } 60 | 61 | // ServerInit handshake. 62 | let n = stream.read(&mut buf).await?; 63 | let pixel_format = messages::PixelFormat::new(&buf[4..20]); 64 | log::debug!("<-: {:?}", pixel_format); 65 | ws_stream 66 | .send(tokio_tungstenite::tungstenite::protocol::Message::Binary( 67 | buf[0..n].to_vec(), 68 | )) 69 | .await?; 70 | 71 | log::debug!("\n--: Handshake finished\n"); 72 | 73 | Ok(RfbConnection { 74 | stream, 75 | connection_state: Arc::new(RfbConnectionState { 76 | bytes_per_pixel: AtomicUsize::new((pixel_format.bits_per_pixel / 8) as usize), 77 | }), 78 | server_tx, 79 | client_tx, 80 | }) 81 | } 82 | 83 | pub fn split<'a>(&'a mut self) -> (ReadHalf<'a>, WriteHalf<'a>) { 84 | let (rs, ws) = self.stream.split(); 85 | ( 86 | ReadHalf { 87 | read_ws: rs, 88 | connection_state: Arc::clone(&self.connection_state), 89 | buf: BytesMut::with_capacity(4096), 90 | }, 91 | WriteHalf { 92 | write_ws: ws, 93 | connection_state: Arc::clone(&self.connection_state), 94 | buf: BytesMut::with_capacity(1024), 95 | client_tx: &mut self.client_tx, 96 | stop_chan: None, 97 | audio_stream: None, 98 | }, 99 | ) 100 | } 101 | } 102 | 103 | /// The readable half of [`RfbConnection::split`]. Represents the server-to-client half. 104 | #[derive(Debug)] 105 | pub struct ReadHalf<'a> { 106 | read_ws: tokio::net::tcp::ReadHalf<'a>, 107 | connection_state: Arc, 108 | buf: BytesMut, 109 | } 110 | 111 | impl ReadHalf<'_> { 112 | pub async fn read_server_message(&mut self) -> Result> { 113 | loop { 114 | // Attempt to parse a frame from the buffered data. If enough data 115 | // has been buffered, the frame is returned. 116 | let mut cur = std::io::Cursor::new(&self.buf[..]); 117 | if let Some(msg) = match server::Message::parse( 118 | &mut cur, 119 | self.connection_state.bytes_per_pixel.load(Ordering::SeqCst), 120 | ) { 121 | Ok(msg) => { 122 | self.buf.advance(msg.len()); 123 | Some(msg) 124 | } 125 | // There is not enough data present in the read buffer to parse a 126 | // single frame. We must wait for more data to be received from the 127 | // socket. Reading from the socket will be done in the statement 128 | // after this `match`. 129 | // 130 | // We do not want to return `Err` from here as this "error" is an 131 | // expected runtime condition. 132 | Err(messages::Error::Incomplete) => None, 133 | // An error was encountered while parsing the frame. The connection 134 | // is now in an invalid state. Returning `Err` from here will result 135 | // in the connection being closed. 136 | Err(e) => { 137 | log::error!("<-: !!welp: {:?}", e); 138 | return Err(e.into()); 139 | } 140 | } { 141 | return Ok(Some(msg)); 142 | } 143 | 144 | if !self.buf.is_empty() { 145 | log::debug!("<-: [incomplete]: {:?}", &self.buf); 146 | log::debug!( 147 | " gonna read a bit, the current buffer of size {} is not enough. brb.", 148 | self.buf.len() 149 | ); 150 | } 151 | 152 | // There is not enough buffered data to read a frame. Attempt to 153 | // read more data from the socket. 154 | // 155 | // On success, the number of bytes is returned. `0` indicates "end 156 | // of stream". 157 | let read_result = self.read_ws.read_buf(&mut self.buf).await; 158 | let n = match read_result { 159 | Ok(0) => { 160 | log::warn!("!!: could not read anything D:."); 161 | // The remote closed the connection. For this to be a clean 162 | // shutdown, there should be no data in the read buffer. If 163 | // there is, this means that the peer closed the socket while 164 | // sending a frame. 165 | if !self.buf.is_empty() { 166 | bail!("connection reset by peer"); 167 | } 168 | return Ok(None); 169 | } 170 | Ok(n) => n, 171 | Err(e) => return Err(e.into()), 172 | }; 173 | log::debug!("<-: okay, let's parse the {} bytes.", n); 174 | } 175 | } 176 | } 177 | 178 | /// The writable half of [`RfbConnection::split`]. Represents the client-to-server half. 179 | pub struct WriteHalf<'a> { 180 | write_ws: tokio::net::tcp::WriteHalf<'a>, 181 | connection_state: Arc, 182 | buf: BytesMut, 183 | client_tx: &'a mut mpsc::Sender>, 184 | stop_chan: Option>, 185 | audio_stream: Option, 186 | } 187 | 188 | impl WriteHalf<'_> { 189 | pub async fn write_all(&mut self, buf: &[u8]) -> Result<()> { 190 | let mut cur = std::io::Cursor::new(buf); 191 | loop { 192 | if cur.remaining() > 0 { 193 | if self.buf.capacity() - self.buf.len() < cur.remaining() { 194 | let old_capacity = self.buf.capacity(); 195 | let old_space_available = self.buf.capacity() - self.buf.len(); 196 | self.buf.reserve(std::cmp::max(cur.remaining(), 4096)); 197 | log::debug!("had to grow buffer to be able to add more stuff. old capacity: {}, old space availabile: {}, capacity: {}, space available: {}", 198 | old_capacity, old_space_available, self.buf.capacity(), self.buf.capacity() - self.buf.len()); 199 | } 200 | 201 | let written = std::cmp::min(self.buf.capacity() - self.buf.len(), cur.remaining()); 202 | if written == 0 { 203 | log::warn!( 204 | "could not write to the buffer! capacity: {}, space available: {}, remaining: {}", 205 | self.buf.capacity(), 206 | self.buf.capacity()- self.buf.len(), 207 | cur.remaining() 208 | ); 209 | } else { 210 | self.buf.extend_from_slice(&cur.get_ref()[0..written]); 211 | cur.advance(written); 212 | } 213 | } 214 | 215 | let result = self.read_client_message().await; 216 | match result { 217 | Err(messages::Error::Incomplete) => break, 218 | Ok(client::Message::ReplitClientAudioStartEncoder( 219 | _payload, 220 | enabled, 221 | channels, 222 | codec, 223 | kbps, 224 | )) => { 225 | if !enabled { 226 | // This drops the sending channel and cause the stream to stop. 227 | self.stop_chan.take(); 228 | self.audio_stream.take(); 229 | self.client_tx 230 | .send( 231 | server::Message::ReplitAudioServerMessage(vec![ 232 | 0xF5, // message-type 233 | 0x00, // submessage-type 234 | 0x00, 0x01, // message length 235 | 0x00, // enabled 236 | ]) 237 | .into_data(), 238 | ) 239 | .await?; 240 | continue; 241 | } 242 | 243 | self.audio_stream = match audio::Stream::new(channels, codec, kbps) { 244 | Ok(stream) => Some(stream), 245 | Err(e) => { 246 | log::error!("failed to create an audio stream: {:#}", e); 247 | self.client_tx 248 | .send( 249 | server::Message::ReplitAudioServerMessage(vec![ 250 | 0xF5, // message-type 251 | 0x00, // submessage-type 252 | 0x00, 0x01, // message length 253 | 0x00, // enabled 254 | ]) 255 | .into_data(), 256 | ) 257 | .await?; 258 | continue; 259 | } 260 | }; 261 | self.client_tx 262 | .send( 263 | server::Message::ReplitAudioServerMessage(vec![ 264 | 0xF5, // message-type 265 | 0x00, // submessage-type 266 | 0x00, 0x01, // message length 267 | 0x01, // enabled 268 | ]) 269 | .into_data(), 270 | ) 271 | .await?; 272 | continue; 273 | } 274 | Ok(client::Message::ReplitClientAudioStartContinuousUpdates(_payload)) => { 275 | let audio_stream = match self.audio_stream.take() { 276 | Some(stream) => stream, 277 | None => { 278 | self.client_tx 279 | .send( 280 | server::Message::ReplitAudioServerMessage(vec![ 281 | 0xF5, // message-type 282 | 0x02, // submessage-type 283 | 0x00, 0x01, // message length 284 | 0x00, // enabled 285 | ]) 286 | .into_data(), 287 | ) 288 | .await?; 289 | continue; 290 | } 291 | }; 292 | self.client_tx 293 | .send( 294 | server::Message::ReplitAudioServerMessage(vec![ 295 | 0xF5, // message-type 296 | 0x02, // submessage-type 297 | 0x00, 0x01, // message length 298 | 0x01, // enabled 299 | ]) 300 | .into_data(), 301 | ) 302 | .await?; 303 | let (stop_chan_tx, stop_chan_rx) = oneshot::channel::<()>(); 304 | self.stop_chan = Some(stop_chan_rx); 305 | let chan = self.client_tx.clone(); 306 | tokio::task::spawn_blocking(move || { 307 | audio_stream.run(stop_chan_tx, chan.clone()); 308 | futures::executor::block_on(async { 309 | if let Err(e) = chan 310 | .send( 311 | server::Message::ReplitAudioServerMessage(vec![ 312 | 0xF5, // message-type 313 | 0x02, // submessage-type 314 | 0x00, 0x01, // message length 315 | 0x00, // enabled 316 | ]) 317 | .into_data(), 318 | ) 319 | .await 320 | { 321 | log::error!("failed to notify client of audio closure: {:#}", e); 322 | } 323 | }); 324 | }); 325 | } 326 | Ok(client::Message::SetEncodings(payload)) => { 327 | let mut cur = std::io::Cursor::new(&payload[4..]); 328 | while cur.has_remaining() { 329 | if cur.get_i32() == 0x52706C41 { 330 | self.client_tx 331 | .send( 332 | server::Message::FramebufferUpdate(vec![ 333 | 0x00, // message-type 334 | 0x00, // padding 335 | 0x00, 0x01, // number-of-rectangles 336 | 0x00, 0x00, // x-position 337 | 0x00, 0x00, // y-position 338 | 0x00, 0x00, // width 339 | 0x00, 0x00, // height 340 | 0x52, 0x70, 0x6C, 341 | 0x41, // Replit Audio Pseudo-encoding 342 | 0x00, 0x00, // Version 343 | 0x00, 0x02, // Number of encodings 344 | 0x00, 0x00, // Opus codec, WebM container 345 | 0x00, 0x01, // MP3 codec, MPEG-1 container 346 | ]) 347 | .into_data(), 348 | ) 349 | .await?; 350 | } 351 | } 352 | 353 | self.write_ws.write_all(&payload).await?; 354 | } 355 | Ok(client::Message::SetPixelFormat(payload)) => { 356 | let pixel_format = messages::PixelFormat::new(&payload[4..20]); 357 | log::debug!("->: SetPixelFormat({:?})", pixel_format); 358 | self.connection_state 359 | .bytes_per_pixel 360 | .store(pixel_format.bits_per_pixel as usize / 8, Ordering::SeqCst); 361 | self.write_ws.write_all(&payload).await?; 362 | } 363 | Ok(msg) => { 364 | log::debug!("->: {}", &msg); 365 | self.write_ws.write_all(&msg.into_data()).await?; 366 | } 367 | Err(err) => { 368 | bail!("failed to read client-to-server message: {:?}", err); 369 | } 370 | } 371 | } 372 | Ok(()) 373 | } 374 | 375 | async fn read_client_message(&mut self) -> Result { 376 | let mut cur = std::io::Cursor::new(&self.buf[..]); 377 | match client::Message::parse(&mut cur) { 378 | Ok(msg) => { 379 | self.buf.advance(msg.len()); 380 | Ok(msg) 381 | } 382 | // There is not enough data present in the read buffer to parse a 383 | // single frame. We must wait for more data to be received from the 384 | // socket. Reading from the socket will be done in the statement 385 | // after this `match`. 386 | // 387 | // We do not want to return `Err` from here as this "error" is an 388 | // expected runtime condition. 389 | Err(messages::Error::Incomplete) => { 390 | if !self.buf.is_empty() { 391 | log::debug!("->: [incomplete]: {:?}", &self.buf); 392 | } 393 | Err(messages::Error::Incomplete) 394 | } 395 | // An error was encountered while parsing the frame. The connection 396 | // is now in an invalid state. Returning `Err` from here will result 397 | // in the connection being closed. 398 | Err(e) => { 399 | log::error!("->: !! welp: {:?}", e); 400 | Err(e) 401 | } 402 | } 403 | } 404 | 405 | pub async fn shutdown(&mut self) -> Result<()> { 406 | // Drop the audio thread. 407 | self.stop_chan.take(); 408 | self.write_ws.shutdown().await?; 409 | Ok(()) 410 | } 411 | } 412 | -------------------------------------------------------------------------------- /src/auth.rs: -------------------------------------------------------------------------------- 1 | //! A wrapper for performing Remote Framebuffer authentication. 2 | 3 | use std::collections::HashMap; 4 | 5 | use anyhow::{anyhow, bail, Context, Result}; 6 | use bytes::BytesMut; 7 | use cipher::{BlockEncrypt, KeyInit}; 8 | use futures::{SinkExt, StreamExt}; 9 | use tokio::io::{AsyncReadExt, AsyncWriteExt}; 10 | use tokio_tungstenite::tungstenite::protocol::Message as WebSocketMessage; 11 | 12 | /// What kind of authentication to use for an RFB connection. 13 | pub enum RfbAuthentication { 14 | /// A null authentication. It does not perform the initial handshake, so it relies on the rest 15 | /// of the connection to pass through the data as-is without parsing. 16 | Null, 17 | 18 | /// An authentication that parses the initial ProtocolVersion and Security handshakes, and 19 | /// passes them as-is to the peer, without acting on it. This leaves the stream in a state 20 | /// where the ClientInit handshake is expected to appear next, followed by a stream of normal 21 | /// RFB messages can appear. 22 | Passthrough, 23 | 24 | /// An authentication that uses a Plain authentication mechanism, where the 25 | /// username is a Replit token for the Repl in which this proxy is being run. It acts as if this 26 | /// were the real server and exposes the Plain authentication as the only valid authentication 27 | /// mechanism, which should be fine since all connections should go over TLS anyways. 28 | /// 29 | /// The password will be sent to the upstream RFB server if it claims to support VncAuth. 30 | /// Otherwise, the password won't be sent (or checked at all!) and the None authentication type 31 | /// will be used. 32 | Replit { 33 | replid: String, 34 | pubkeys: HashMap>, 35 | }, 36 | } 37 | 38 | /// A way of authenticating an RFB connection. 39 | pub async fn authenticate( 40 | authentication: &RfbAuthentication, 41 | stream: &mut SocketStream, 42 | ws_stream: &mut WebSocketStream, 43 | ) -> Result<()> 44 | where 45 | SocketStream: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send, 46 | WebSocketStream: futures::Sink 47 | + futures::Stream< 48 | Item = std::result::Result< 49 | WebSocketMessage, 50 | tokio_tungstenite::tungstenite::error::Error, 51 | >, 52 | > + Unpin 53 | + Send, 54 | { 55 | match authentication { 56 | RfbAuthentication::Null => Ok(()), 57 | RfbAuthentication::Passthrough => authenticate_passthrough(stream, ws_stream).await, 58 | RfbAuthentication::Replit { replid, pubkeys } => { 59 | authenticate_replit(stream, ws_stream, replid, pubkeys).await 60 | } 61 | } 62 | } 63 | 64 | async fn authenticate_passthrough( 65 | stream: &mut SocketStream, 66 | ws_stream: &mut WebSocketStream, 67 | ) -> Result<()> 68 | where 69 | SocketStream: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send, 70 | WebSocketStream: futures::Sink 71 | + futures::Stream< 72 | Item = std::result::Result< 73 | WebSocketMessage, 74 | tokio_tungstenite::tungstenite::error::Error, 75 | >, 76 | > + Unpin 77 | + Send, 78 | { 79 | let mut buf = [0u8; 1024]; 80 | 81 | // ProtocolVersion handshake. 82 | let n = stream.read(&mut buf[0..12]).await?; 83 | if n != 12 { 84 | bail!("unexpected server handshake: {:?}", &buf[..n]); 85 | } 86 | log::debug!("<-: {:?}", std::str::from_utf8(&buf[0..12])?); 87 | ws_stream 88 | .send(WebSocketMessage::Binary(buf[0..12].to_vec())) 89 | .await?; 90 | match ws_stream.next().await { 91 | Some(msg) => match msg.context("bad client ProtocolVersion handshake")? { 92 | WebSocketMessage::Binary(payload) => { 93 | log::debug!("->: {:?}", std::str::from_utf8(&payload)?); 94 | stream.write_all(&payload).await?; 95 | } 96 | unexpected_msg => bail!("unexpected message {:?}", unexpected_msg), 97 | }, 98 | None => bail!("missing client ProtocolVersion handshake"), 99 | } 100 | 101 | // Security handshake. 102 | let mut n = stream.read(&mut buf).await?; 103 | log::debug!("<-: {:?}", &buf[0..n]); 104 | ws_stream 105 | .send(WebSocketMessage::Binary(buf[0..n].to_vec())) 106 | .await?; 107 | let client_security_handshake = match ws_stream.next().await { 108 | Some(msg) => match msg.context("bad client security handshake")? { 109 | WebSocketMessage::Binary(payload) => { 110 | if payload.len() != 1 { 111 | bail!( 112 | "unexpected security-type length. got {}, expected 1", 113 | payload.len() 114 | ); 115 | } 116 | log::debug!("->: {:?}", &payload); 117 | stream.write_all(&payload).await?; 118 | payload[0] 119 | } 120 | unexpected_msg => bail!("unexpected message {:?}", unexpected_msg), 121 | }, 122 | None => bail!("missing client security handshake"), 123 | }; 124 | match client_security_handshake { 125 | 1 => { 126 | // None security type 127 | } 128 | 2 => { 129 | // VNC Authentication security type 130 | n = stream.read(&mut buf).await?; 131 | log::debug!("<-: {:?}", &buf[0..n]); 132 | ws_stream 133 | .send(WebSocketMessage::Binary(buf[0..n].to_vec())) 134 | .await?; 135 | match ws_stream.next().await { 136 | Some(msg) => match msg.context("bad client VNCAuth security handshake")? { 137 | WebSocketMessage::Binary(payload) => { 138 | log::debug!("->: {:?}", &payload); 139 | stream.write_all(&payload).await?; 140 | } 141 | unexpected_msg => bail!("unexpected message {:?}", unexpected_msg), 142 | }, 143 | None => bail!("missing client VNCAuth security handshake"), 144 | } 145 | } 146 | unsupported => bail!("unsupported security type {}", unsupported), 147 | } 148 | 149 | // SecurityResult handshake. 150 | n = stream.read(&mut buf).await?; 151 | log::debug!("<-: {:?}", &buf[0..n]); 152 | ws_stream 153 | .send(WebSocketMessage::Binary(buf[0..n].to_vec())) 154 | .await?; 155 | 156 | Ok(()) 157 | } 158 | 159 | async fn authenticate_replit( 160 | stream: &mut SocketStream, 161 | ws_stream: &mut WebSocketStream, 162 | replid: &str, 163 | pubkeys: &HashMap>, 164 | ) -> Result<()> 165 | where 166 | SocketStream: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send, 167 | WebSocketStream: futures::Sink 168 | + futures::Stream< 169 | Item = std::result::Result< 170 | WebSocketMessage, 171 | tokio_tungstenite::tungstenite::error::Error, 172 | >, 173 | > + Unpin 174 | + Send, 175 | { 176 | client_handshake_step(ws_stream, b"RFB 003.008\n", b"RFB 003.008\n") 177 | .await 178 | .context("client ProtocolVersion handshake")?; 179 | 180 | // Only support the VeNCrypt authentication. 181 | client_handshake_step(ws_stream, b"\x01\x13", b"\x13") 182 | .await 183 | .context("client security type")?; 184 | 185 | // Only support VeNCrypt 0.2. 186 | client_handshake_step(ws_stream, b"\x00\x02", b"\x00\x02") 187 | .await 188 | .context("client VeNCrypt version")?; 189 | 190 | // Only support Plain authentication. 191 | ws_stream 192 | .send(WebSocketMessage::Binary(b"\x00".to_vec())) 193 | .await?; 194 | client_handshake_step(ws_stream, b"\x01\x00\x00\x01\x00", b"\x00\x00\x01\x00") 195 | .await 196 | .context("client VeNCrypt subtype")?; 197 | 198 | // All the preamble is done, now receive username+password. 199 | let (username, password) = client_username_password(ws_stream).await?; 200 | 201 | // We can let users through if they use the "runner" username, _but_ that's only if a password 202 | // is provided _and_ the server requires a password. 203 | let use_token = username != "runner" || password.is_empty(); 204 | if use_token { 205 | if let Err(err) = validate_token(&username, replid, pubkeys).context("token validation") { 206 | // Let the client know that it messed up. 207 | ws_stream 208 | .send(WebSocketMessage::Binary((b"\x00\x00\x00\x01").to_vec())) 209 | .await?; 210 | return Err(err); 211 | } 212 | } 213 | 214 | // Now that the token itself was validated, we perform the handshake against the upstream RFB 215 | // server. 216 | server_handshake_step(stream, b"RFB 003.008\n", b"RFB 003.008\n") 217 | .await 218 | .context("server ProtocolVersion handshake")?; 219 | 220 | let mut buf = [0u8; 1024]; 221 | 222 | // SecurityType handshake. This consists of a byte indicating the number of SecurityTypes, 223 | // followed by that many bytes describing a SecurityType supported by the server. 224 | let mut n = stream.read(&mut buf[..]).await?; 225 | if n <= 1 || buf[0] as usize + 1 != n { 226 | bail!("invalid SecurityType payload {:?}", &buf[..n]); 227 | } 228 | log::debug!("<-: {:?}", &buf[..n]); 229 | 230 | // The server sends the supported SecurityTypes in an arbitrary order, so we cannot check any 231 | // specific byte. Also, skipping the first byte since that's the length of the list. 232 | if buf[1..n].contains(&2) { 233 | // VncAuth. Relay the user-provided password. 234 | stream.write_all(b"\x02").await?; 235 | log::debug!("->: {:?}", b"\x02"); 236 | n = stream.read(&mut buf[..16]).await?; 237 | if n != 16 { 238 | bail!("invalid VncAuth nonce: {:?}", &buf[..n]); 239 | } 240 | log::debug!("<-: {:?}", &buf[..n]); 241 | 242 | vnc_des_encrypt(&password, &mut buf[..n]); 243 | log::debug!("->: {:?}", &buf[..n]); 244 | stream.write_all(&buf[..n]).await?; 245 | } else if buf[1..n].contains(&1) { 246 | // None 247 | if !use_token { 248 | // Let the client know that it messed up, since "runner" cannot be used if the server 249 | // does not validate the password. 250 | ws_stream 251 | .send(WebSocketMessage::Binary((b"\x00\x00\x00\x01").to_vec())) 252 | .await?; 253 | bail!( 254 | "server does not have a password set up, cannot use basic password authentication." 255 | ); 256 | } 257 | stream.write_all(b"\x01").await?; 258 | } else { 259 | bail!("no supported SecurityTypes found: {:?}", &buf[1..n]); 260 | } 261 | 262 | // SecurityResult handshake. 263 | n = stream.read(&mut buf[..]).await?; 264 | log::debug!("<-: {:?}", &buf[..n]); 265 | ws_stream 266 | .send(WebSocketMessage::Binary(buf[..n].to_vec())) 267 | .await?; 268 | if &buf[..n] != b"\x00\x00\x00\x00" { 269 | // Bail after sending the reply so that the client can display the "authentication failed" 270 | // message. 271 | bail!("authentication failure: {:?}", &buf[..n]); 272 | } 273 | 274 | Ok(()) 275 | } 276 | 277 | /// Sends a message to the client, reads its response, and checks that the response matches what we 278 | /// expect. 279 | async fn client_handshake_step( 280 | ws_stream: &mut WebSocketStream, 281 | message: &[u8], 282 | expected: &[u8], 283 | ) -> Result<()> 284 | where 285 | WebSocketStream: futures::Sink 286 | + futures::Stream< 287 | Item = std::result::Result< 288 | WebSocketMessage, 289 | tokio_tungstenite::tungstenite::error::Error, 290 | >, 291 | > + Unpin 292 | + Send, 293 | { 294 | ws_stream 295 | .send(WebSocketMessage::Binary(message.to_vec())) 296 | .await?; 297 | log::debug!("->: {:?}", &message); 298 | match ws_stream.next().await { 299 | Some(msg) => match msg.context("bad client ProtocolVersion handshake")? { 300 | WebSocketMessage::Binary(payload) => { 301 | if payload != expected { 302 | bail!("mismatched payload {:?}", payload); 303 | } 304 | } 305 | unexpected_msg => bail!("unexpected message {:?}", unexpected_msg), 306 | }, 307 | None => bail!("missing client message"), 308 | } 309 | 310 | Ok(()) 311 | } 312 | 313 | /// Reads a message from the server, checks that it matches what we expect, and sends a message in 314 | /// response. 315 | async fn server_handshake_step( 316 | stream: &mut SocketStream, 317 | expected: &[u8], 318 | message: &[u8], 319 | ) -> Result<()> 320 | where 321 | SocketStream: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send, 322 | { 323 | let mut buf = [0u8; 1024]; 324 | 325 | // ProtocolVersion handshake. 326 | let n = stream.read(&mut buf[..]).await?; 327 | if &buf[..n] != expected { 328 | bail!("mismatched payload {:?}", &buf[..n]); 329 | } 330 | log::debug!("<-: {:?}", &buf[..n]); 331 | stream.write_all(message).await?; 332 | 333 | Ok(()) 334 | } 335 | 336 | /// Reads the username / password from the client. Some implementations (like noVNC) send each 337 | /// length and the strings in separate WebSocket messages, but we should not rely on implementation 338 | /// details. This buffers the WebSocketStream and provides a stream-like view for 339 | /// [`parse_client_username_password`], which does the reading, and signals the loop in this 340 | /// function to read more from the client before retrying. 341 | async fn client_username_password( 342 | ws_stream: &mut WebSocketStream, 343 | ) -> Result<(String, String)> 344 | where 345 | WebSocketStream: futures::Sink 346 | + futures::Stream< 347 | Item = std::result::Result< 348 | WebSocketMessage, 349 | tokio_tungstenite::tungstenite::error::Error, 350 | >, 351 | > + Unpin 352 | + Send, 353 | { 354 | let mut buf = BytesMut::with_capacity(1024); 355 | loop { 356 | match ws_stream.next().await { 357 | Some(msg) => match msg.context("bad WebSocket message")? { 358 | WebSocketMessage::Binary(payload) => { 359 | buf.extend_from_slice(&payload); 360 | } 361 | unexpected_msg => bail!("unexpected message {:?}", unexpected_msg), 362 | }, 363 | None => bail!("missing client message"), 364 | } 365 | 366 | let mut cur = std::io::Cursor::new(&buf[..]); 367 | match parse_client_username_password(&mut cur) { 368 | Ok((username, password)) => { 369 | return Ok((String::from_utf8(username)?, String::from_utf8(password)?)); 370 | } 371 | Err(crate::messages::Error::Incomplete) => {} 372 | Err(e) => { 373 | return Err(e.into()); 374 | } 375 | } 376 | } 377 | } 378 | 379 | /// Reads the username / password from the client, where `cur` is a stream view of the client's 380 | /// WebSocket connection. 381 | fn parse_client_username_password( 382 | cur: &mut std::io::Cursor<&[u8]>, 383 | ) -> Result<(Vec, Vec), crate::messages::Error> { 384 | let username_length = crate::messages::io::get_u32(cur)? as usize; 385 | let password_length = crate::messages::io::get_u32(cur)? as usize; 386 | 387 | let username = crate::messages::io::read(cur, username_length)?; 388 | let password = crate::messages::io::read(cur, password_length)?; 389 | 390 | Ok((username.to_vec(), password.to_vec())) 391 | } 392 | 393 | /// Validate a Goval Handshake v5 token. It should be: 394 | /// 395 | /// - Issued by one of the known public keys. 396 | /// - Be valid at this point in time. 397 | /// - Be issued for the repl where this is being run. 398 | fn validate_token(token: &str, replid: &str, pubkeys: &HashMap>) -> Result<()> { 399 | use prost::Message; 400 | 401 | let token_parts = token.split('.').collect::>(); 402 | if token_parts.len() != 4 { 403 | bail!("token has wrong number of parts: {}", token_parts.len()); 404 | } 405 | let raw_footer = base64::decode_config(token_parts[3], base64::URL_SAFE_NO_PAD) 406 | .context("failed to extract the PASETO footer")?; 407 | let footer = crate::api::GovalTokenMetadata::decode( 408 | &*base64::decode(&raw_footer).context("failed to base64-decode the PASETO footer")?, 409 | ) 410 | .context("failed to parse the PASETO footer")?; 411 | 412 | let repl_token = crate::api::ReplToken::decode( 413 | &*base64::decode( 414 | match paseto::v2::verify_paseto( 415 | token, 416 | Some(std::str::from_utf8(&raw_footer)?), 417 | pubkeys 418 | .get(&footer.key_id) 419 | .ok_or_else(|| anyhow!("could not find {} in pubkeys", &footer.key_id))?, 420 | ) { 421 | Ok(message) => message, 422 | Err(err) => bail!("failed to verify PASETO: {}", err), 423 | }, 424 | ) 425 | .context("failed to base64-decode the PASETO message")?, 426 | ) 427 | .context("failed to parse the PASETO message")?; 428 | 429 | // Validate issue / expiration timestamps. 430 | let iat = match repl_token.iat.as_ref() { 431 | Some(ts) => std::time::SystemTime::UNIX_EPOCH 432 | .checked_add(std::time::Duration::new(ts.seconds as u64, ts.nanos as u32)) 433 | .ok_or_else(|| anyhow!("overflow decoding iat: {:?}", repl_token.iat.as_ref()))?, 434 | None => std::time::SystemTime::UNIX_EPOCH, 435 | }; 436 | let exp = match repl_token.exp.as_ref() { 437 | Some(ts) => std::time::SystemTime::UNIX_EPOCH 438 | .checked_add(std::time::Duration::new(ts.seconds as u64, ts.nanos as u32)) 439 | .ok_or_else(|| anyhow!("overflow decoding exp: {:?}", repl_token.exp.as_ref()))?, 440 | None => iat 441 | .checked_add(std::time::Duration::from_secs(3600)) 442 | .ok_or_else(|| anyhow!("overflow providing fallback iat: {:?}", &iat))?, 443 | }; 444 | let now = std::time::SystemTime::now(); 445 | if now < iat { 446 | bail!( 447 | "token issued in the past: {}", 448 | chrono::DateTime::::from(iat).to_rfc3339() 449 | ); 450 | } 451 | if now > exp { 452 | bail!( 453 | "token expired: {}", 454 | chrono::DateTime::::from(exp).to_rfc3339() 455 | ); 456 | } 457 | 458 | // Validate ReplID. 459 | let token_replid = match &repl_token.metadata { 460 | Some(crate::api::repl_token::Metadata::Repl(repl)) => repl.id.clone(), 461 | Some(crate::api::repl_token::Metadata::Id(id)) => id.id.clone(), 462 | _ => bail!("token does not contain a replid: {:?}", &repl_token), 463 | }; 464 | let pruned_replid = replid.split_once(":").map(|(replid, _)| replid.to_string()); 465 | if token_replid != replid && Some(&token_replid) != pruned_replid.as_ref() { 466 | bail!( 467 | "token not issued for replid {:?}: {:?}", 468 | &token_replid, 469 | replid, 470 | ); 471 | } 472 | 473 | Ok(()) 474 | } 475 | 476 | /// Encrypts a buffer with the password acting as a DES key, compatible with VNC authentication. 477 | /// The key is generated by truncating the password to 8 characters (or extends it to 8 characters, 478 | /// filling with zeroes), and reversing the bits of each byte of the password. 479 | fn vnc_des_encrypt(password: &str, buf: &mut [u8]) { 480 | let mut key = String::from(password).into_bytes(); 481 | key.resize(8, 0); 482 | // Reverse the bits of each byte. 483 | for key_byte in &mut key { 484 | let mut x = *key_byte; 485 | *key_byte = 0; 486 | for _ in 0..8 { 487 | *key_byte = (*key_byte << 1) | (x & 1); 488 | x >>= 1; 489 | } 490 | } 491 | let des = des::Des::new(cipher::Key::::from_slice(&key)); 492 | for i in (0..buf.len()).step_by(8) { 493 | des.encrypt_block(cipher::Block::::from_mut_slice( 494 | &mut buf[i..i + 8], 495 | )); 496 | } 497 | } 498 | 499 | #[cfg(test)] 500 | #[allow(deprecated)] 501 | mod tests { 502 | use super::*; 503 | 504 | use prost::Message; 505 | use ring::signature::KeyPair; 506 | use tokio_test::io::Builder; 507 | 508 | fn init() { 509 | let _ = env_logger::builder().is_test(true).try_init(); 510 | } 511 | 512 | #[test] 513 | fn test_passthrough_none_security_type() { 514 | init(); 515 | 516 | let mut socket_mock = Builder::new() 517 | .read(b"RFB 003.008\n") 518 | .write(b"RFB 003.008\n") 519 | // Only the None(1) security type is supported. 520 | .read(b"\x01\x01") 521 | .write(b"\x01") 522 | // Success! 523 | .read(b"\x00\x00\x00\x00") 524 | .build(); 525 | // Acting as a server to avoid having to unmask the frames. 526 | let mut websocket_stream = 527 | tokio_test::block_on(tokio_tungstenite::WebSocketStream::from_raw_socket( 528 | Builder::new() 529 | .write(b"\x82\x0cRFB 003.008\n") 530 | .read(b"\x82\x0cRFB 003.008\n") 531 | .write(b"\x82\x02\x01\x01") 532 | .read(b"\x82\x01\x01") 533 | .write(b"\x82\x04\x00\x00\x00\x00") 534 | .build(), 535 | tokio_tungstenite::tungstenite::protocol::Role::Server, 536 | Some(tokio_tungstenite::tungstenite::protocol::WebSocketConfig { 537 | max_send_queue: None, 538 | write_buffer_size: 128 * 1024, 539 | max_write_buffer_size: usize::MAX, 540 | max_message_size: None, 541 | max_frame_size: None, 542 | accept_unmasked_frames: true, 543 | }), 544 | )); 545 | 546 | tokio_test::block_on(authenticate_passthrough( 547 | &mut socket_mock, 548 | &mut websocket_stream, 549 | )) 550 | .expect("could not authenticate"); 551 | } 552 | 553 | #[test] 554 | fn test_passthrough_vncauth_security_type() { 555 | init(); 556 | 557 | let mut socket_mock = Builder::new() 558 | .read(b"RFB 003.008\n") 559 | .write(b"RFB 003.008\n") 560 | // Only the VncAuth(2) security type is supported. 561 | .read(b"\x01\x02") 562 | .write(b"\x02") 563 | // Challenge + Response. The password is, unsurprisingly, "password". 564 | .read(b"\x9e\xdd\x1d\xc2\xee\x5a\x5e\x78\x7f\x55\x21\xf2\x67\x9f\x71\xd6") 565 | .write(b"\x15\x6d\x69\xd7\x0f\x22\x21\xb5\x6f\x46\xe2\x92\xa3\xe2\x68\x37") 566 | // Success! 567 | .read(b"\x00\x00\x00\x00") 568 | .build(); 569 | // Acting as a server to avoid having to unmask the frames. 570 | let mut websocket_stream = 571 | tokio_test::block_on(tokio_tungstenite::WebSocketStream::from_raw_socket( 572 | Builder::new() 573 | .write(b"\x82\x0cRFB 003.008\n") 574 | .read(b"\x82\x0cRFB 003.008\n") 575 | .write(b"\x82\x02\x01\x02") 576 | .read(b"\x82\x01\x02") 577 | .write( 578 | b"\x82\x10\x9e\xdd\x1d\xc2\xee\x5a\x5e\x78\x7f\x55\x21\xf2\x67\x9f\x71\xd6", 579 | ) 580 | .read( 581 | b"\x82\x10\x15\x6d\x69\xd7\x0f\x22\x21\xb5\x6f\x46\xe2\x92\xa3\xe2\x68\x37", 582 | ) 583 | .write(b"\x82\x04\x00\x00\x00\x00") 584 | .build(), 585 | tokio_tungstenite::tungstenite::protocol::Role::Server, 586 | Some(tokio_tungstenite::tungstenite::protocol::WebSocketConfig { 587 | max_send_queue: None, 588 | write_buffer_size: 128 * 1024, 589 | max_write_buffer_size: usize::MAX, 590 | max_message_size: None, 591 | max_frame_size: None, 592 | accept_unmasked_frames: true, 593 | }), 594 | )); 595 | 596 | tokio_test::block_on(authenticate_passthrough( 597 | &mut socket_mock, 598 | &mut websocket_stream, 599 | )) 600 | .expect("could not authenticate"); 601 | } 602 | 603 | #[test] 604 | fn test_replit_none_security_type() { 605 | init(); 606 | 607 | let replid = "repl"; 608 | let keyid = "keyid"; 609 | 610 | let sys_rand = ring::rand::SystemRandom::new(); 611 | let key_pkcs8 = ring::signature::Ed25519KeyPair::generate_pkcs8(&sys_rand) 612 | .expect("Failed to generate pkcs8 key!"); 613 | let keypair = ring::signature::Ed25519KeyPair::from_pkcs8(key_pkcs8.as_ref()) 614 | .expect("Failed to parse keypair"); 615 | let pubkey = keypair.public_key(); 616 | 617 | let token = mint_token( 618 | replid, 619 | keyid, 620 | None, 621 | Some(prost_types::Timestamp { 622 | seconds: 253402329599, 623 | nanos: 0, 624 | }), 625 | &keypair, 626 | ) 627 | .expect("Failed to generate PASETO"); 628 | log::debug!( 629 | "{} --replid={} --pubkeys={{\"{}\":\"{}\"}}\n", 630 | token, 631 | &replid, 632 | &keyid, 633 | base64::encode(&pubkey) 634 | ); 635 | 636 | // Acting as a server to avoid having to unmask the frames. 637 | let mut websocket_stream = 638 | tokio_test::block_on(tokio_tungstenite::WebSocketStream::from_raw_socket( 639 | Builder::new() 640 | .write(b"\x82\x0cRFB 003.008\n") 641 | .read(b"\x82\x0cRFB 003.008\n") 642 | // Only the VeNCrypt security type is supported. 643 | .write(b"\x82\x02\x01\x13") 644 | .read(b"\x82\x01\x13") 645 | // VeNCrypt version 0.2 646 | .write(b"\x82\x02\x00\x02") 647 | .read(b"\x82\x02\x00\x02") 648 | // Only the Plain subtype is supported. 649 | .write(b"\x82\x01\x00") 650 | .write(b"\x82\x05\x01\x00\x00\x01\x00") 651 | .read(b"\x82\x04\x00\x00\x01\x00") 652 | // Username length 653 | .read(&[ 654 | 0x82, 655 | 0x04, 656 | 0x00, 657 | 0x00, 658 | (token.len() >> 8 & 0xFF) as u8, 659 | (token.len() & 0xFF) as u8, 660 | ]) 661 | // Password length 662 | .read(b"\x82\x04\x00\x00\x00\x08") 663 | // Username 664 | .read( 665 | &[ 666 | &[ 667 | 0x82, 668 | 0x7E, 669 | (token.len() >> 8 & 0xFF) as u8, 670 | (token.len() & 0xFF) as u8, 671 | ], 672 | token.as_bytes(), 673 | ] 674 | .concat(), 675 | ) 676 | // Password 677 | .read(b"\x82\x08password") 678 | // Success! 679 | .write(b"\x82\x04\x00\x00\x00\x00") 680 | .build(), 681 | tokio_tungstenite::tungstenite::protocol::Role::Server, 682 | Some(tokio_tungstenite::tungstenite::protocol::WebSocketConfig { 683 | max_send_queue: None, 684 | write_buffer_size: 128 * 1024, 685 | max_write_buffer_size: usize::MAX, 686 | max_message_size: None, 687 | max_frame_size: None, 688 | accept_unmasked_frames: true, 689 | }), 690 | )); 691 | let mut socket_mock = Builder::new() 692 | .read(b"RFB 003.008\n") 693 | .write(b"RFB 003.008\n") 694 | // Only the None(2) security type is supported. 695 | .read(b"\x01\x01") 696 | .write(b"\x01") 697 | // Success! 698 | .read(b"\x00\x00\x00\x00") 699 | .build(); 700 | 701 | let mut pubkeys = HashMap::>::new(); 702 | pubkeys.insert(keyid.to_string(), pubkey.as_ref().to_vec()); 703 | 704 | tokio_test::block_on(authenticate_replit( 705 | &mut socket_mock, 706 | &mut websocket_stream, 707 | &replid.to_string(), 708 | &pubkeys, 709 | )) 710 | .expect("could not authenticate"); 711 | } 712 | 713 | #[test] 714 | fn test_replit_none_security_type_with_runner_username() { 715 | init(); 716 | 717 | let replid = "repl"; 718 | let keyid = "keyid"; 719 | 720 | let sys_rand = ring::rand::SystemRandom::new(); 721 | let key_pkcs8 = ring::signature::Ed25519KeyPair::generate_pkcs8(&sys_rand) 722 | .expect("Failed to generate pkcs8 key!"); 723 | let keypair = ring::signature::Ed25519KeyPair::from_pkcs8(key_pkcs8.as_ref()) 724 | .expect("Failed to parse keypair"); 725 | let pubkey = keypair.public_key(); 726 | 727 | let token = mint_token( 728 | replid, 729 | keyid, 730 | None, 731 | Some(prost_types::Timestamp { 732 | seconds: 253402329599, 733 | nanos: 0, 734 | }), 735 | &keypair, 736 | ) 737 | .expect("Failed to generate PASETO"); 738 | log::debug!( 739 | "{} --replid={} --pubkeys={{\"{}\":\"{}\"}}\n", 740 | token, 741 | &replid, 742 | &keyid, 743 | base64::encode(&pubkey) 744 | ); 745 | 746 | // Acting as a server to avoid having to unmask the frames. 747 | let mut websocket_stream = 748 | tokio_test::block_on(tokio_tungstenite::WebSocketStream::from_raw_socket( 749 | Builder::new() 750 | .write(b"\x82\x0cRFB 003.008\n") 751 | .read(b"\x82\x0cRFB 003.008\n") 752 | // Only the VeNCrypt security type is supported. 753 | .write(b"\x82\x02\x01\x13") 754 | .read(b"\x82\x01\x13") 755 | // VeNCrypt version 0.2 756 | .write(b"\x82\x02\x00\x02") 757 | .read(b"\x82\x02\x00\x02") 758 | // Only the Plain subtype is supported. 759 | .write(b"\x82\x01\x00") 760 | .write(b"\x82\x05\x01\x00\x00\x01\x00") 761 | .read(b"\x82\x04\x00\x00\x01\x00") 762 | // Username length 763 | .read(b"\x82\x04\x00\x00\x00\x06") 764 | // Password length 765 | .read(b"\x82\x04\x00\x00\x00\x08") 766 | // Username 767 | .read(b"\x82\x06runner") 768 | // Password 769 | .read(b"\x82\x08password") 770 | // Oh noes! 771 | .write(b"\x82\x04\x00\x00\x00\x01") 772 | .build(), 773 | tokio_tungstenite::tungstenite::protocol::Role::Server, 774 | Some(tokio_tungstenite::tungstenite::protocol::WebSocketConfig { 775 | max_send_queue: None, 776 | write_buffer_size: 128 * 1024, 777 | max_write_buffer_size: usize::MAX, 778 | max_message_size: None, 779 | max_frame_size: None, 780 | accept_unmasked_frames: true, 781 | }), 782 | )); 783 | let mut socket_mock = Builder::new() 784 | .read(b"RFB 003.008\n") 785 | .write(b"RFB 003.008\n") 786 | // Only the None(2) security type is supported. 787 | .read(b"\x01\x01") 788 | .build(); 789 | 790 | let mut pubkeys = HashMap::>::new(); 791 | pubkeys.insert(keyid.to_string(), pubkey.as_ref().to_vec()); 792 | 793 | let auth_err = tokio_test::block_on(authenticate_replit( 794 | &mut socket_mock, 795 | &mut websocket_stream, 796 | &replid.to_string(), 797 | &pubkeys, 798 | )) 799 | .expect_err("Should have rejected the runner username"); 800 | assert_eq!( 801 | auth_err.to_string(), 802 | "server does not have a password set up, cannot use basic password authentication." 803 | ); 804 | } 805 | 806 | #[test] 807 | fn test_replit_vncauth_security_type() { 808 | init(); 809 | 810 | let replid = "repl"; 811 | let keyid = "keyid"; 812 | 813 | let sys_rand = ring::rand::SystemRandom::new(); 814 | let key_pkcs8 = ring::signature::Ed25519KeyPair::generate_pkcs8(&sys_rand) 815 | .expect("Failed to generate pkcs8 key!"); 816 | let keypair = ring::signature::Ed25519KeyPair::from_pkcs8(key_pkcs8.as_ref()) 817 | .expect("Failed to parse keypair"); 818 | let pubkey = keypair.public_key(); 819 | 820 | let token = mint_token( 821 | replid, 822 | keyid, 823 | None, 824 | Some(prost_types::Timestamp { 825 | seconds: 253402329599, 826 | nanos: 0, 827 | }), 828 | &keypair, 829 | ) 830 | .expect("Failed to generate PASETO"); 831 | log::debug!( 832 | "{} --replid={} --pubkeys={{\"{}\":\"{}\"}}\n", 833 | token, 834 | &replid, 835 | &keyid, 836 | base64::encode(&pubkey) 837 | ); 838 | 839 | // Acting as a server to avoid having to unmask the frames. 840 | let mut websocket_stream = 841 | tokio_test::block_on(tokio_tungstenite::WebSocketStream::from_raw_socket( 842 | Builder::new() 843 | .write(b"\x82\x0cRFB 003.008\n") 844 | .read(b"\x82\x0cRFB 003.008\n") 845 | // Only the VeNCrypt security type is supported. 846 | .write(b"\x82\x02\x01\x13") 847 | .read(b"\x82\x01\x13") 848 | // VeNCrypt version 0.2 849 | .write(b"\x82\x02\x00\x02") 850 | .read(b"\x82\x02\x00\x02") 851 | // Only the Plain subtype is supported. 852 | .write(b"\x82\x01\x00") 853 | .write(b"\x82\x05\x01\x00\x00\x01\x00") 854 | .read(b"\x82\x04\x00\x00\x01\x00") 855 | // Username length 856 | .read(&[ 857 | 0x82, 858 | 0x04, 859 | 0x00, 860 | 0x00, 861 | (token.len() >> 8 & 0xFF) as u8, 862 | (token.len() & 0xFF) as u8, 863 | ]) 864 | // Password length 865 | .read(b"\x82\x04\x00\x00\x00\x08") 866 | // Username 867 | .read( 868 | &[ 869 | &[ 870 | 0x82, 871 | 0x7E, 872 | (token.len() >> 8 & 0xFF) as u8, 873 | (token.len() & 0xFF) as u8, 874 | ], 875 | token.as_bytes(), 876 | ] 877 | .concat(), 878 | ) 879 | // Password 880 | .read(b"\x82\x08password") 881 | // Success! 882 | .write(b"\x82\x04\x00\x00\x00\x00") 883 | .build(), 884 | tokio_tungstenite::tungstenite::protocol::Role::Server, 885 | Some(tokio_tungstenite::tungstenite::protocol::WebSocketConfig { 886 | max_send_queue: None, 887 | write_buffer_size: 128 * 1024, 888 | max_write_buffer_size: usize::MAX, 889 | max_message_size: None, 890 | max_frame_size: None, 891 | accept_unmasked_frames: true, 892 | }), 893 | )); 894 | let mut socket_mock = Builder::new() 895 | .read(b"RFB 003.008\n") 896 | .write(b"RFB 003.008\n") 897 | // Only the VncAuth(2) security type is supported. 898 | .read(b"\x01\x02") 899 | .write(b"\x02") 900 | // Challenge + Response. The password is, unsurprisingly, "password". 901 | .read(b"\x9e\xdd\x1d\xc2\xee\x5a\x5e\x78\x7f\x55\x21\xf2\x67\x9f\x71\xd6") 902 | .write(b"\x15\x6d\x69\xd7\x0f\x22\x21\xb5\x6f\x46\xe2\x92\xa3\xe2\x68\x37") 903 | // Success! 904 | .read(b"\x00\x00\x00\x00") 905 | .build(); 906 | 907 | let mut pubkeys = HashMap::>::new(); 908 | pubkeys.insert(keyid.to_string(), pubkey.as_ref().to_vec()); 909 | 910 | tokio_test::block_on(authenticate_replit( 911 | &mut socket_mock, 912 | &mut websocket_stream, 913 | &replid.to_string(), 914 | &pubkeys, 915 | )) 916 | .expect("could not authenticate"); 917 | } 918 | 919 | #[test] 920 | fn test_replit_vncauth_security_type_with_runner_username() { 921 | init(); 922 | 923 | let replid = "repl"; 924 | let keyid = "keyid"; 925 | 926 | let sys_rand = ring::rand::SystemRandom::new(); 927 | let key_pkcs8 = ring::signature::Ed25519KeyPair::generate_pkcs8(&sys_rand) 928 | .expect("Failed to generate pkcs8 key!"); 929 | let keypair = ring::signature::Ed25519KeyPair::from_pkcs8(key_pkcs8.as_ref()) 930 | .expect("Failed to parse keypair"); 931 | let pubkey = keypair.public_key(); 932 | 933 | let token = mint_token( 934 | replid, 935 | keyid, 936 | None, 937 | Some(prost_types::Timestamp { 938 | seconds: 253402329599, 939 | nanos: 0, 940 | }), 941 | &keypair, 942 | ) 943 | .expect("Failed to generate PASETO"); 944 | log::debug!( 945 | "{} --replid={} --pubkeys={{\"{}\":\"{}\"}}\n", 946 | token, 947 | &replid, 948 | &keyid, 949 | base64::encode(&pubkey) 950 | ); 951 | 952 | // Acting as a server to avoid having to unmask the frames. 953 | let mut websocket_stream = 954 | tokio_test::block_on(tokio_tungstenite::WebSocketStream::from_raw_socket( 955 | Builder::new() 956 | .write(b"\x82\x0cRFB 003.008\n") 957 | .read(b"\x82\x0cRFB 003.008\n") 958 | // Only the VeNCrypt security type is supported. 959 | .write(b"\x82\x02\x01\x13") 960 | .read(b"\x82\x01\x13") 961 | // VeNCrypt version 0.2 962 | .write(b"\x82\x02\x00\x02") 963 | .read(b"\x82\x02\x00\x02") 964 | // Only the Plain subtype is supported. 965 | .write(b"\x82\x01\x00") 966 | .write(b"\x82\x05\x01\x00\x00\x01\x00") 967 | .read(b"\x82\x04\x00\x00\x01\x00") 968 | // Username length 969 | .read(b"\x82\x04\x00\x00\x00\x06") 970 | // Password length 971 | .read(b"\x82\x04\x00\x00\x00\x08") 972 | // Username 973 | .read(b"\x82\x06runner") 974 | // Password 975 | .read(b"\x82\x08password") 976 | // Success! 977 | .write(b"\x82\x04\x00\x00\x00\x00") 978 | .build(), 979 | tokio_tungstenite::tungstenite::protocol::Role::Server, 980 | Some(tokio_tungstenite::tungstenite::protocol::WebSocketConfig { 981 | max_send_queue: None, 982 | write_buffer_size: 128 * 1024, 983 | max_write_buffer_size: usize::MAX, 984 | max_message_size: None, 985 | max_frame_size: None, 986 | accept_unmasked_frames: true, 987 | }), 988 | )); 989 | let mut socket_mock = Builder::new() 990 | .read(b"RFB 003.008\n") 991 | .write(b"RFB 003.008\n") 992 | // Only the VncAuth(2) security type is supported. 993 | .read(b"\x01\x02") 994 | .write(b"\x02") 995 | // Challenge + Response. The password is, unsurprisingly, "password". 996 | .read(b"\x9e\xdd\x1d\xc2\xee\x5a\x5e\x78\x7f\x55\x21\xf2\x67\x9f\x71\xd6") 997 | .write(b"\x15\x6d\x69\xd7\x0f\x22\x21\xb5\x6f\x46\xe2\x92\xa3\xe2\x68\x37") 998 | // Success! 999 | .read(b"\x00\x00\x00\x00") 1000 | .build(); 1001 | 1002 | let mut pubkeys = HashMap::>::new(); 1003 | pubkeys.insert(keyid.to_string(), pubkey.as_ref().to_vec()); 1004 | 1005 | tokio_test::block_on(authenticate_replit( 1006 | &mut socket_mock, 1007 | &mut websocket_stream, 1008 | &replid.to_string(), 1009 | &pubkeys, 1010 | )) 1011 | .expect("could not authenticate"); 1012 | } 1013 | 1014 | #[test] 1015 | fn test_validate_token() { 1016 | init(); 1017 | 1018 | let replid = "repl"; 1019 | 1020 | let sys_rand = ring::rand::SystemRandom::new(); 1021 | 1022 | let keyid = "keyid"; 1023 | let keypair = ring::signature::Ed25519KeyPair::from_pkcs8( 1024 | ring::signature::Ed25519KeyPair::generate_pkcs8(&sys_rand) 1025 | .expect("Failed to generate pkcs8 key!") 1026 | .as_ref(), 1027 | ) 1028 | .expect("Failed to parse keypair"); 1029 | let pubkey = keypair.public_key(); 1030 | let mut pubkeys = HashMap::>::new(); 1031 | pubkeys.insert(keyid.to_string(), pubkey.as_ref().to_vec()); 1032 | 1033 | let keyid_other = "keyid_other"; 1034 | let keypair_other = ring::signature::Ed25519KeyPair::from_pkcs8( 1035 | ring::signature::Ed25519KeyPair::generate_pkcs8(&sys_rand) 1036 | .expect("Failed to generate pkcs8 key!") 1037 | .as_ref(), 1038 | ) 1039 | .expect("Failed to parse keypair"); 1040 | let pubkey_other = keypair_other.public_key(); 1041 | let mut pubkeys_other = HashMap::>::new(); 1042 | pubkeys_other.insert(keyid_other.to_string(), pubkey_other.as_ref().to_vec()); 1043 | 1044 | let mut pubkeys_wrong_pubkey = HashMap::>::new(); 1045 | pubkeys_wrong_pubkey.insert(keyid.to_string(), pubkey_other.as_ref().to_vec()); 1046 | 1047 | let token = mint_token( 1048 | replid, 1049 | keyid, 1050 | None, 1051 | Some(prost_types::Timestamp { 1052 | seconds: 253402329599, 1053 | nanos: 0, 1054 | }), 1055 | &keypair, 1056 | ) 1057 | .expect("Failed to generate PASETO"); 1058 | 1059 | validate_token(&token, &replid.to_string(), &pubkeys).expect("Failed to validate token"); 1060 | validate_token(&token, &format!("{replid}:01").to_string(), &pubkeys) 1061 | .expect("Failed to validate token"); 1062 | validate_token( 1063 | &String::from("this is not a token"), 1064 | &replid.to_string(), 1065 | &pubkeys, 1066 | ) 1067 | .expect_err("Should have rejected an invalid token"); 1068 | validate_token(&token, &replid.to_string(), &pubkeys_wrong_pubkey) 1069 | .expect_err("Should have rejected a token signed with a mismatched key"); 1070 | validate_token(&token, &replid.to_string(), &pubkeys_other) 1071 | .expect_err("Should have rejected a token signed with an unknown key"); 1072 | validate_token(&token, &String::from("other repl"), &pubkeys) 1073 | .expect_err("Should have rejected a token signed for another repl"); 1074 | 1075 | validate_token( 1076 | &mint_token( 1077 | replid, 1078 | keyid, 1079 | None, 1080 | Some(prost_types::Timestamp { 1081 | seconds: 0, 1082 | nanos: 0, 1083 | }), 1084 | &keypair, 1085 | ) 1086 | .expect("Failed to generate PASETO"), 1087 | &replid.to_string(), 1088 | &pubkeys, 1089 | ) 1090 | .expect_err("Should have rejected an expired token"); 1091 | validate_token( 1092 | &mint_token(replid, keyid, None, None, &keypair).expect("Failed to generate PASETO"), 1093 | &replid.to_string(), 1094 | &pubkeys, 1095 | ) 1096 | .expect_err("Should have rejected an (implicitly) expired token"); 1097 | validate_token( 1098 | &mint_token( 1099 | replid, 1100 | keyid, 1101 | Some(prost_types::Timestamp { 1102 | seconds: 253402329599, 1103 | nanos: 0, 1104 | }), 1105 | Some(prost_types::Timestamp { 1106 | seconds: 253402329599, 1107 | nanos: 0, 1108 | }), 1109 | &keypair, 1110 | ) 1111 | .expect("Failed to generate PASETO"), 1112 | &replid.to_string(), 1113 | &pubkeys, 1114 | ) 1115 | .expect_err("Should have rejected a not-yet-issued token"); 1116 | } 1117 | 1118 | fn mint_token( 1119 | replid: &str, 1120 | keyid: &str, 1121 | iat: Option, 1122 | exp: Option, 1123 | keypair: &ring::signature::Ed25519KeyPair, 1124 | ) -> Result { 1125 | let mut buf = BytesMut::with_capacity(1024); 1126 | 1127 | let mut repl_token = crate::api::ReplToken::default(); 1128 | repl_token.iat = iat; 1129 | repl_token.exp = exp; 1130 | repl_token.cluster = String::from("development"); 1131 | repl_token.metadata = Some(crate::api::repl_token::Metadata::Id( 1132 | crate::api::repl_token::ReplId { 1133 | id: String::from(replid), 1134 | source_repl: String::from(""), 1135 | }, 1136 | )); 1137 | repl_token.encode(&mut buf).expect("could not encode token"); 1138 | let message = base64::encode(&buf); 1139 | 1140 | buf.clear(); 1141 | crate::api::GovalTokenMetadata { 1142 | key_id: String::from(keyid), 1143 | } 1144 | .encode(&mut buf) 1145 | .expect("could not encode footer"); 1146 | let footer = base64::encode(&buf); 1147 | 1148 | let token = match paseto::v2::public_paseto(&message, Some(&footer), keypair) { 1149 | Ok(token) => token, 1150 | Err(err) => bail!("failed to generate PASETO: {}", err), 1151 | }; 1152 | 1153 | Ok(token) 1154 | } 1155 | } 1156 | --------------------------------------------------------------------------------