radio/iroh/examples/dht_address_lookup.rs

//! An example chat application using the iroh endpoint and
//! pkarr address lookup.
//!
//! Starting the example without args creates a server that publishes its
//! address to the DHT. Starting the example with an endpoint id as argument
//! looks up the address of the endpoint id in the DHT and connects to it.
//!
//! You can look at the published pkarr DNS record using <https://app.pkarr.org/>.
//!
//! To see what is going on, run with `RUST_LOG=iroh_pkarr_address_lookup=debug`.
use std::str::FromStr;

use clap::Parser;
use iroh::{Endpoint, EndpointId};
use n0_error::{Result, StdResultExt};
use tracing::warn;
use url::Url;

const CHAT_ALPN: &[u8] = b"pkarr-address-lookup-demo-chat";

#[derive(Parser)]
struct Args {
    /// The endpoint id to connect to. If not set, the program will start a server.
    endpoint_id: Option<EndpointId>,
    /// Disable using the mainline DHT for Address Lookup and publishing.
    #[clap(long)]
    disable_dht: bool,
    /// Pkarr relay to use.
    #[clap(long, default_value = "iroh")]
    pkarr_relay: PkarrRelay,
}

#[derive(Debug, Clone)]
enum PkarrRelay {
    /// Disable pkarr relay.
    Disabled,
    /// Use the iroh pkarr relay.
    Iroh,
    /// Use a custom pkarr relay.
    Custom(Url),
}

impl FromStr for PkarrRelay {
    type Err = url::ParseError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        match s {
            "disabled" => Ok(Self::Disabled),
            "iroh" => Ok(Self::Iroh),
            s => Ok(Self::Custom(Url::parse(s)?)),
        }
    }
}

fn build_address_lookup(args: Args) -> iroh::address_lookup::pkarr::dht::Builder {
    let builder = iroh::address_lookup::DhtAddressLookup::builder().dht(!args.disable_dht);
    match args.pkarr_relay {
        PkarrRelay::Disabled => builder,
        PkarrRelay::Iroh => builder.n0_dns_pkarr_relay(),
        PkarrRelay::Custom(url) => builder.pkarr_relay(url),
    }
}

async fn chat_server(args: Args) -> Result<()> {
    let secret_key = iroh::SecretKey::generate(&mut rand::rng());
    let endpoint_id = secret_key.public();
    let address_lookup = build_address_lookup(args);
    let endpoint = Endpoint::builder()
        .alpns(vec![CHAT_ALPN.to_vec()])
        .secret_key(secret_key)
        .address_lookup(address_lookup)
        .bind()
        .await?;
    let zid = pkarr::PublicKey::try_from(endpoint_id.as_bytes())
        .anyerr()?
        .to_z32();
    println!("Listening on {endpoint_id}");
    println!("pkarr z32: {zid}");
    println!("see https://app.pkarr.org/?pk={zid}");
    while let Some(incoming) = endpoint.accept().await {
        let accepting = match incoming.accept() {
            Ok(accepting) => accepting,
            Err(err) => {
                warn!("incoming connection failed: {err:#}");
                // we can carry on in these cases:
                // this can be caused by retransmitted datagrams
                continue;
            }
        };
        tokio::spawn(async move {
            let connection = accepting.await?;
            let remote_endpoint_id = connection.remote_id();
            println!("got connection from {remote_endpoint_id}");
            // just leave the tasks hanging. this is just an example.
            let (mut writer, mut reader) = connection.accept_bi().await.anyerr()?;
            let _copy_to_stdout = tokio::spawn(async move {
                tokio::io::copy(&mut reader, &mut tokio::io::stdout()).await
            });
            let _copy_from_stdin =
                tokio::spawn(
                    async move { tokio::io::copy(&mut tokio::io::stdin(), &mut writer).await },
                );
            n0_error::Ok(())
        });
    }
    Ok(())
}

async fn chat_client(args: Args) -> Result<()> {
    let remote_endpoint_id = args.endpoint_id.unwrap();
    let secret_key = iroh::SecretKey::generate(&mut rand::rng());
    let endpoint_id = secret_key.public();
    // note: we don't pass a secret key here, because we don't need to publish our address, don't spam the DHT
    let address_lookup = build_address_lookup(args).no_publish();
    // we do not need to specify the alpn here, because we are not going to accept connections
    let endpoint = Endpoint::builder()
        .secret_key(secret_key)
        .address_lookup(address_lookup)
        .bind()
        .await?;
    println!("We are {endpoint_id} and connecting to {remote_endpoint_id}");
    let connection = endpoint.connect(remote_endpoint_id, CHAT_ALPN).await?;
    println!("connected to {remote_endpoint_id}");
    let (mut writer, mut reader) = connection.open_bi().await.anyerr()?;
    let _copy_to_stdout =
        tokio::spawn(async move { tokio::io::copy(&mut reader, &mut tokio::io::stdout()).await });
    let _copy_from_stdin =
        tokio::spawn(async move { tokio::io::copy(&mut tokio::io::stdin(), &mut writer).await });
    _copy_to_stdout.await.anyerr()?.anyerr()?;
    _copy_from_stdin.await.anyerr()?.anyerr()?;
    Ok(())
}

#[tokio::main]
async fn main() -> Result<()> {
    tracing_subscriber::fmt::init();
    let args = Args::parse();
    if args.endpoint_id.is_some() {
        chat_client(args).await?;
    } else {
        chat_server(args).await?;
    }
    Ok(())
}