Skip to content

Commit

Permalink
don't break free tunnels
Browse files Browse the repository at this point in the history
  • Loading branch information
@charlottea98
charlottea98 committed May 29, 2024
1 parent f844482 commit 38b852c
Show file tree
Hide file tree
Showing 6 changed files with 120 additions and 71 deletions.
16 changes: 8 additions & 8 deletions linkup-cli/src/background_booting.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,13 +20,13 @@ use crate::{CliError, LINKUP_LOCALDNS_INSTALL};

#[cfg_attr(test, mockall::automock)]
pub trait BackgroundServices {
fn boot_background_services(&self, state: LocalState) -> Result<(), CliError>;
fn boot_background_services(&self, state: LocalState) -> Result<LocalState, CliError>;
}

pub struct RealBackgroundServices;

impl BackgroundServices for RealBackgroundServices {
fn boot_background_services(&self, mut state: LocalState) -> Result<(), CliError> {
fn boot_background_services(&self, mut state: LocalState) -> Result<LocalState, CliError> {
let local_url = Url::parse(&format!("http://localhost:{}", LINKUP_LOCALSERVER_PORT))
.expect("linkup url invalid");

Expand Down Expand Up @@ -75,17 +75,17 @@ impl BackgroundServices for RealBackgroundServices {
boot_local_dns(state.domain_strings(), state.linkup.session_name.clone())?;
}

// if let Some(tunnel) = &state.linkup.tunnel {
// println!("Waiting for tunnel DNS to propogate at {}...", tunnel);
if let Some(tunnel) = &state.linkup.tunnel {
println!("Waiting for tunnel DNS to propagate at {}...", tunnel);

// wait_for_dns_ok(tunnel.clone())?;
wait_for_dns_ok(tunnel.clone())?;

// println!();
// }
println!();
}

print_session_names(&state);

Ok(())
Ok(state)
}
}

Expand Down
5 changes: 5 additions & 0 deletions linkup-cli/src/file_system.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ pub trait FileSystem {
fn file_exists(&self, file_path: &Path) -> bool;
fn create_dir_all(&self, path: &Path) -> Result<(), CliError>;
fn get_home(&self) -> Result<String, CliError>;
fn get_env(&self, key: &str) -> Result<String, CliError>;
}

pub struct RealFileSystem;
Expand Down Expand Up @@ -41,4 +42,8 @@ impl FileSystem for RealFileSystem {
fn get_home(&self) -> Result<String, CliError> {
Ok(env::var("HOME").expect("HOME is not set"))
}

fn get_env(&self, key: &str) -> Result<String, CliError> {
Ok(env::var(key).unwrap_or_else(|_| panic!("{} is not set", key)))
}
}
3 changes: 3 additions & 0 deletions linkup-cli/src/local_config.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,9 @@ impl LocalState {
}

pub fn save(&mut self) -> Result<(), CliError> {
if cfg!(test) {
return Ok(());
}
let yaml_string = match serde_yaml::to_string(self) {
Ok(yaml) => yaml,
Err(_) => {
Expand Down
51 changes: 32 additions & 19 deletions linkup-cli/src/paid_tunnel.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ use std::{env, fs, path::Path};
use reqwest::header::{HeaderMap, HeaderValue, AUTHORIZATION, CONTENT_TYPE};

use crate::file_system::{FileLike, FileSystem, RealFileSystem};
use crate::CliError;
use crate::{CliError, LINKUP_LOCALSERVER_PORT};
use serde::{Deserialize, Serialize};

use base64::prelude::*;
Expand Down Expand Up @@ -189,7 +189,8 @@ impl PaidTunnelManager for RealPaidTunnelManager {

println!("{}", body);

let _parsed : CreateDNSRecordResponse = send_request(&client, &url, headers, Some(body), "POST")?;
let _parsed: CreateDNSRecordResponse =
send_request(&client, &url, headers, Some(body), "POST")?;
Ok(())
}
}
Expand All @@ -205,7 +206,8 @@ fn save_tunnel_credentials(
tunnel_id: &str,
tunnel_secret: &str,
) -> Result<(), CliError> {
let account_id = env::var("LINKUP_CLOUDFLARE_ACCOUNT_ID")
let account_id = filesys
.get_env("LINKUP_CLOUDFLARE_ACCOUNT_ID")
.map_err(|err| CliError::BadConfig(err.to_string()))?;
let data = serde_json::json!({
"AccountTag": account_id,
Expand All @@ -214,7 +216,9 @@ fn save_tunnel_credentials(
});

// Determine the directory path
let home_dir = env::var("HOME").map_err(|err| CliError::BadConfig(err.to_string()))?;
let home_dir = filesys
.get_env("HOME")
.map_err(|err| CliError::BadConfig(err.to_string()))?;
let dir_path = Path::new(&home_dir).join(".cloudflared");

// Create the directory if it does not exist
Expand All @@ -239,7 +243,9 @@ fn save_tunnel_credentials(

fn create_config_yml(filesys: &dyn FileSystem, tunnel_id: &str) -> Result<(), CliError> {
// Determine the directory path
let home_dir = env::var("HOME").map_err(|err| CliError::BadConfig(err.to_string()))?;
let home_dir = filesys
.get_env("HOME")
.map_err(|err| CliError::BadConfig(err.to_string()))?;
let dir_path = Path::new(&home_dir).join(".cloudflared");

// Create the directory if it does not exist
Expand All @@ -255,7 +261,7 @@ fn create_config_yml(filesys: &dyn FileSystem, tunnel_id: &str) -> Result<(), Cl
let file_path_str = file_path.to_string_lossy().to_string();

let config = Config {
url: "http://localhost:8000".to_string(),
url: format!("http://localhost:{}", LINKUP_LOCALSERVER_PORT),
tunnel: tunnel_id.to_string(),
credentials_file: file_path_str,
};
Expand Down Expand Up @@ -325,11 +331,14 @@ mod tests {

#[test]
fn create_config_yml_when_no_config_dir() {
env::set_var("HOME", "/tmp/home");
let content = "url: http://localhost:8000\ntunnel: TUNNEL_ID\ncredentials-file: /tmp/home/.cloudflared/TUNNEL_ID.json\n";
let content = "url: http://localhost:9066\ntunnel: TUNNEL_ID\ncredentials-file: /tmp/home/.cloudflared/TUNNEL_ID.json\n";

let mut mock_fs = MockFileSystem::new();
// If .cloudflared directory does not exist:
mock_fs
.expect_get_env()
.with(predicate::eq("HOME"))
.returning(|_| Ok("/tmp/home".to_string()));
mock_fs
.expect_file_exists()
.withf(|path| path.ends_with(".cloudflared"))
Expand All @@ -352,16 +361,17 @@ mod tests {

let result = create_config_yml(&mock_fs, "TUNNEL_ID");
assert!(result.is_ok());

env::remove_var("HOME")
}

#[test]
fn create_config_yml_config_dir_exists() {
env::set_var("HOME", "/tmp/home");
let content = "url: http://localhost:8000\ntunnel: TUNNEL_ID\ncredentials-file: /tmp/home/.cloudflared/TUNNEL_ID.json\n";
let content = "url: http://localhost:9066\ntunnel: TUNNEL_ID\ncredentials-file: /tmp/home/.cloudflared/TUNNEL_ID.json\n";

let mut mock_fs = MockFileSystem::new();
mock_fs
.expect_get_env()
.with(predicate::eq("HOME"))
.returning(|_| Ok("/tmp/home".to_string()));
// If .cloudflared directory exists:
mock_fs
.expect_file_exists()
Expand All @@ -382,17 +392,23 @@ mod tests {

let result = create_config_yml(&mock_fs, "TUNNEL_ID");
assert!(result.is_ok());

env::remove_var("HOME")
}

#[test]
fn test_save_tunnel_credentials() {
env::set_var("HOME", "/tmp/home");
env::set_var("LINKUP_CLOUDFLARE_ACCOUNT_ID", "ACCOUNT_ID");
let content = "{\"AccountTag\":\"ACCOUNT_ID\",\"TunnelID\":\"TUNNEL_ID\",\"TunnelSecret\":\"AQIDBAUGBwgBAgMEBQYHCAECAwQFBgcIAQIDBAUGBwg=\"}";

let mut mock_fs = MockFileSystem::new();
mock_fs
.expect_get_env()
.with(predicate::eq("HOME"))
.returning(|_| Ok("/tmp/home".to_string()));

mock_fs
.expect_get_env()
.with(predicate::eq("LINKUP_CLOUDFLARE_ACCOUNT_ID"))
.returning(|_| Ok("ACCOUNT_ID".to_string()));

mock_fs
.expect_create_file()
.withf(|path| path.ends_with("TUNNEL_ID.json"))
Expand All @@ -408,9 +424,6 @@ mod tests {
"AQIDBAUGBwgBAgMEBQYHCAECAwQFBgcIAQIDBAUGBwg=",
);
assert!(result.is_ok());

env::remove_var("HOME");
env::remove_var("LINKUP_CLOUDFLARE_ACCOUNT_ID");
}

#[test]
Expand Down
23 changes: 17 additions & 6 deletions linkup-cli/src/services/tunnel.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ impl TunnelManager for RealTunnelManager {
let mut attempt = 0;
loop {
attempt += 1;
match try_run_tunnel(&state) {
match try_run_tunnel(state) {
Ok(url) => return Ok(url),
Err(CliError::StopErr(e)) => {
return Err(CliError::StopErr(format!(
Expand Down Expand Up @@ -80,7 +80,7 @@ fn try_run_tunnel(state: &LocalState) -> Result<Url, CliError> {

match daemonize.execute() {
Outcome::Child(child_result) => match child_result {
Ok(_) => daemonized_tunnel_child(&state),
Ok(_) => daemonized_tunnel_child(state),
Err(e) => {
return Err(CliError::StartLocalTunnel(format!(
"Failed to start local tunnel: {}",
Expand All @@ -99,8 +99,11 @@ fn try_run_tunnel(state: &LocalState) -> Result<Url, CliError> {
},
}

let is_paid = state.is_paid;
let session_name = state.linkup.session_name.clone();

let tunnel_url_re =
Regex::new(r"Starting tunnel tunnelID=.*").expect("Failed to compile regex");
Regex::new(r"https://[a-zA-Z0-9-]+\.trycloudflare\.com").expect("Failed to compile regex");
let tunnel_started_re =
Regex::new(r"Registered tunnel connection").expect("Failed to compile regex");

Expand All @@ -123,7 +126,15 @@ fn try_run_tunnel(state: &LocalState) -> Result<Url, CliError> {

for line in buf_reader.lines() {
let line = line.unwrap_or_default();
if let Some(url_match) = tunnel_url_re.find(&line) {
if is_paid {
url = Some(
Url::parse(
format!("https://tunnel-{}.mentimeter.dev", session_name)
.as_str(),
)
.expect("Failed to parse tunnel URL"),
);
} else if let Some(url_match) = tunnel_url_re.find(&line) {
let found_url =
Url::parse(url_match.as_str()).expect("Failed to parse tunnel URL");
url = Some(found_url);
Expand All @@ -149,7 +160,6 @@ fn try_run_tunnel(state: &LocalState) -> Result<Url, CliError> {
thread::sleep(Duration::from_millis(100));
}
});

match rx.recv_timeout(Duration::from_secs(TUNNEL_START_WAIT)) {
Ok(result) => result,
Err(e) => {
Expand All @@ -166,8 +176,9 @@ fn daemonized_tunnel_child(state: &LocalState) {
let url = format!("http://localhost:{}", LINKUP_LOCALSERVER_PORT);
let cmd_args: Vec<&str> = match state.is_paid {
true => vec!["tunnel", "run", state.linkup.session_name.as_str()],
false => vec!["tunnel", "run", "--url", url.as_str()],
false => vec!["tunnel", "--url", url.as_str()],
};
println!("Starting cloudflared tunnel with args: {:?}", cmd_args);
let mut child_cmd = Command::new("cloudflared")
.args(cmd_args)
.stdout(Stdio::inherit())
Expand Down
Loading

0 comments on commit 38b852c

Please sign in to comment.