Bail when trying to nest stream sessions using the same relay URL

This commit is contained in:
Marcin Kulik
2024-04-11 21:56:19 +02:00
parent 83d08a071a
commit dd77179d74
6 changed files with 59 additions and 6 deletions

8
Cargo.lock generated
View File

@@ -99,6 +99,7 @@ dependencies = [
"scraper",
"serde",
"serde_json",
"sha2",
"signal-hook",
"tempfile",
"termion",
@@ -1627,16 +1628,17 @@ dependencies = [
[[package]]
name = "ring"
version = "0.17.7"
version = "0.17.8"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "688c63d65483050968b2a8937f7995f443e27041a0f7700aa59b0822aedebb74"
checksum = "c17fa4cb658e3583423e915b9f3acc01cceaee1860e33d59ebae66adc3a2dc0d"
dependencies = [
"cc",
"cfg-if",
"getrandom 0.2.12",
"libc",
"spin",
"untrusted",
"windows-sys 0.48.0",
"windows-sys 0.52.0",
]
[[package]]

View File

@@ -38,6 +38,7 @@ tracing-subscriber = { version = "0.3.18", features = ["env-filter"] }
rgb = "0.8.37"
url = "2.5.0"
tokio-tungstenite = { version = "0.21.0", features = ["rustls-tls-webpki-roots"] }
sha2 = "0.10.8"
[profile.release]
strip = true

View File

@@ -26,10 +26,14 @@ fn build_exec_command(command: Option<String>) -> Vec<String> {
vec!["/bin/sh".to_owned(), "-c".to_owned(), command]
}
fn build_exec_extra_env() -> HashMap<String, String> {
fn build_exec_extra_env(vars: &[(String, String)]) -> HashMap<String, String> {
let mut env = HashMap::new();
env.insert("ASCIINEMA_REC".to_owned(), "1".to_owned());
for (k, v) in vars {
env.insert(k.clone(), v.clone());
}
env
}

View File

@@ -83,7 +83,7 @@ impl Cli {
let notifier = super::get_notifier(config);
let record_input = self.input || config.cmd_rec_input();
let exec_command = super::build_exec_command(command.as_ref().cloned());
let exec_extra_env = super::build_exec_extra_env();
let exec_extra_env = super::build_exec_extra_env(&[]);
logger::info!("Recording session started, writing to {}", self.filename);

View File

@@ -4,11 +4,16 @@ use crate::logger;
use crate::pty;
use crate::streamer::{self, KeyBindings};
use crate::tty;
use crate::util;
use anyhow::bail;
use anyhow::{anyhow, Context, Result};
use clap::Args;
use reqwest::blocking::Client;
use reqwest::header;
use serde::Deserialize;
use std::collections::HashMap;
use std::env;
use std::fmt::Debug;
use std::fs;
use std::net::SocketAddr;
use std::net::TcpListener;
@@ -96,9 +101,20 @@ impl Cli {
let notifier = super::get_notifier(config);
let record_input = self.input || config.cmd_stream_input();
let exec_command = super::build_exec_command(command.as_ref().cloned());
let exec_extra_env = super::build_exec_extra_env();
let listener = self.get_listener()?;
let relay = self.get_relay(config)?;
let relay_id = relay.as_ref().map(|r| r.id());
let exec_extra_env = build_exec_extra_env(relay_id.as_ref());
if let (Some(id), Some(parent_id)) = (relay_id, parent_session_relay_id()) {
if id == parent_id {
if let Some(Relay { url: Some(url), .. }) = relay {
bail!("This shell is already being streamed at {url}");
} else {
bail!("This shell is already being streamed");
}
}
}
logger::info!("Streaming session started");
@@ -212,6 +228,12 @@ impl Cli {
}
}
impl Relay {
fn id(&self) -> String {
util::sha2_digest(self.ws_producer_url.as_ref())
}
}
fn get_server_stream(stream_id: String, config: &Config) -> Result<GetStreamResponse> {
let response = Client::new()
.get(stream_api_url(&config.get_server_url()?, stream_id))
@@ -251,3 +273,14 @@ fn get_key_bindings(config: &Config) -> Result<KeyBindings> {
Ok(keys)
}
fn build_exec_extra_env(relay_id: Option<&String>) -> HashMap<String, String> {
match relay_id {
Some(id) => super::build_exec_extra_env(&[("ASCIINEMA_RELAY_ID".to_string(), id.clone())]),
None => super::build_exec_extra_env(&[]),
}
}
fn parent_session_relay_id() -> Option<String> {
env::var("ASCIINEMA_RELAY_ID").ok()
}

View File

@@ -1,5 +1,6 @@
use anyhow::{anyhow, bail, Result};
use reqwest::Url;
use sha2::Digest;
use std::path::{Path, PathBuf};
use std::{io, thread};
use tempfile::NamedTempFile;
@@ -115,6 +116,18 @@ impl Utf8Decoder {
}
}
pub fn sha2_digest(s: &str) -> String {
let mut hasher = sha2::Sha224::new();
hasher.update(s.as_bytes());
hasher
.finalize()
.as_slice()
.iter()
.map(|byte| format!("{:02x}", byte))
.collect()
}
#[cfg(test)]
mod tests {
use super::Utf8Decoder;