relay_actions/cache/

agent.rs

use anyhow::{Context, Result, bail};
use serde::{Deserialize, Serialize};
use std::{
    collections::HashMap,
    path::Path,
    time::{SystemTime, UNIX_EPOCH},
};

use super::Cache;
use relay_lib::prelude::AgentId;

const FILE: &str = "agents.ron";
const SRV_PREFIX: &str = "_relay._tcp";

#[derive(Debug, Default, Serialize, Deserialize)]
pub struct AgentsCache {
    pub agents: HashMap<AgentId, CachedAgent>,
}

impl Cache for AgentsCache {
    fn load(base: &Path) -> Result<Self> {
        let content = std::fs::read_to_string(base.join(FILE)).unwrap_or_default();
        let agents = ron::from_str(&content).unwrap_or_default();
        Ok(Self { agents })
    }

    fn save(&self, base: &Path) -> Result<()> {
        let content = ron::to_string(&self.agents)?.into_bytes();
        std::fs::create_dir_all(base)?;
        std::fs::write(base.join(FILE), content)?;
        Ok(())
    }

    fn cleanup(&mut self) -> Result<()> {
        self.agents.retain(|_, agent| agent.is_valid());
        Ok(())
    }
}

impl AgentsCache {
    pub fn fetch(&mut self, agent_id: &AgentId) -> Result<CachedAgent> {
        if let Some(agent) = self.agents.get(agent_id)
            && agent.is_valid()
        {
            return Ok(agent.clone());
        }

        let rt = tokio::runtime::Runtime::new()?;
        let agent = rt.block_on(async {
            match Self::fetch_via_srv(agent_id).await {
                Ok(agent) => Ok(agent),
                Err(srv_err) => Self::fetch_via_well_known(agent_id).await.with_context(|| {
                    format!(
                        "SRV failed: {} and well-known fallback also failed",
                        srv_err
                    )
                }),
            }
        })?;

        self.agents.insert(agent_id.clone(), agent.clone());
        Ok(agent)
    }

    pub fn url(&mut self, agent_id: &AgentId) -> Result<reqwest::Url> {
        let agent = self.fetch(agent_id)?;

        if !agent.is_valid() {
            bail!("Cached agent information is expired");
        }

        reqwest::Url::parse(&format!("{}://{}:{}", protocol(), agent.host, agent.port))
            .context("Failed to construct agent URL")
    }

    async fn fetch_via_srv(agent_id: &AgentId) -> Result<CachedAgent> {
        let resolver = hickory_resolver::Resolver::builder_tokio().unwrap().build();

        let response = resolver
            .srv_lookup(format!("{}.{}", SRV_PREFIX, agent_id))
            .await
            .context("SRV lookup failed")?;

        let mut records: Vec<_> = response.iter().collect();

        if records.is_empty() {
            bail!("No SRV records found");
        }

        records.sort_by_key(|r| r.priority());

        for record in records {
            let host = record
                .target()
                .to_string()
                .trim_end_matches('.')
                .to_string();

            let port = record.port();

            if host.is_empty() || port == 0 {
                continue;
            }

            if resolver.lookup_ip(host.clone()).await.is_err() {
                continue;
            }

            return Ok(CachedAgent {
                host,
                port,
                fetched_at: now_secs(),
                ttl_secs: 24 * 3600,
            });
        }

        bail!("No valid SRV records")
    }

    async fn fetch_via_well_known(agent_id: &AgentId) -> Result<CachedAgent> {
        let host = agent_id.to_string();
        let ports = if protocol() == "https" {
            [433, 8433, 2525]
        } else {
            [80, 8080, 2525]
        };

        for port in ports {
            let url = format!("{}://{}:{}/.well-known/relay", protocol(), host, port);

            match Self::try_fetch(&url).await {
                Ok(agent) => return Ok(agent),
                Err(_) => {
                    continue;
                }
            }
        }

        bail!("well-known failed on all known ports")
    }

    async fn try_fetch(url: &str) -> Result<CachedAgent> {
        let resp = reqwest::get(url).await.context("request failed")?;

        if !resp.status().is_success() {
            bail!("status {}", resp.status());
        }

        #[derive(Deserialize)]
        struct WellKnown {
            agent: AgentInfo,
        }

        #[derive(Deserialize)]
        struct AgentInfo {
            url: String,
        }

        let data: WellKnown = resp.json().await?;

        let url = reqwest::Url::parse(&data.agent.url)?;

        let host = url
            .host_str()
            .ok_or_else(|| anyhow::anyhow!("Missing host"))?;
        let port = url
            .port_or_known_default()
            .ok_or_else(|| anyhow::anyhow!("Missing port"))?;

        Ok(CachedAgent {
            host: host.to_string(),
            port,
            fetched_at: now_secs(),
            ttl_secs: 24 * 3600,
        })
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CachedAgent {
    pub host: String,
    pub port: u16,
    pub fetched_at: u64,
    pub ttl_secs: u64,
}

impl CachedAgent {
    pub fn is_valid(&self) -> bool {
        let now = now_secs();
        now < self.fetched_at + self.ttl_secs
    }
}

fn now_secs() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_secs()
}

/// Useful fofr testing local agents without SSL. In debug mode, the protocol can be overridden to
/// "http" by setting the environment variable RELAY_NO_SSL=1.
#[inline]
fn protocol() -> &'static str {
    if cfg!(debug_assertions) {
        match std::env::var("RELAY_NO_SSL") {
            Ok(val) if val == "1" => "http",
            Ok(_) => "https",
            Err(_) => "https",
        }
    } else {
        "https"
    }
}