From dd77179d74df9cf89eb02f57c19c8e63cf59c8d9 Mon Sep 17 00:00:00 2001
From: Marcin Kulik <m@ku1ik.com>
Date: Thu, 11 Apr 2024 21:56:19 +0200
Subject: [PATCH] Bail when trying to nest stream sessions using the same relay
 URL

---
 Cargo.lock        |  8 +++++---
 Cargo.toml        |  1 +
 src/cmd/mod.rs    |  6 +++++-
 src/cmd/rec.rs    |  2 +-
 src/cmd/stream.rs | 35 ++++++++++++++++++++++++++++++++++-
 src/util.rs       | 13 +++++++++++++
 6 files changed, 59 insertions(+), 6 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index a2459e2..5022106 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -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]]
diff --git a/Cargo.toml b/Cargo.toml
index 2149782..6ed48ce 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -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
diff --git a/src/cmd/mod.rs b/src/cmd/mod.rs
index dd17e04..6d97b53 100644
--- a/src/cmd/mod.rs
+++ b/src/cmd/mod.rs
@@ -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
 }
diff --git a/src/cmd/rec.rs b/src/cmd/rec.rs
index 5f7baa8..f6785cd 100644
--- a/src/cmd/rec.rs
+++ b/src/cmd/rec.rs
@@ -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);
 
diff --git a/src/cmd/stream.rs b/src/cmd/stream.rs
index fd898e8..52ab4be 100644
--- a/src/cmd/stream.rs
+++ b/src/cmd/stream.rs
@@ -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()
+}
diff --git a/src/util.rs b/src/util.rs
index f3e69a6..f37ab94 100644
--- a/src/util.rs
+++ b/src/util.rs
@@ -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;