use anyhow::Result;
use iroh::endpoint::Connection;
use iroh::{Endpoint, EndpointAddr, PublicKey, SecretKey};
use std::path::PathBuf;
use tokio::fs;
const DEFAULT_ALPN: &[u8] = b"xoq/p2p/0";
pub struct IrohServerBuilder {
key_path: Option<PathBuf>,
secret_key: Option<SecretKey>,
alpn: Vec<u8>,
}
impl IrohServerBuilder {
pub fn new() -> Self {
Self {
key_path: None,
secret_key: None,
alpn: DEFAULT_ALPN.to_vec(),
}
}
pub fn identity_path(mut self, path: impl Into<PathBuf>) -> Self {
self.key_path = Some(path.into());
self
}
pub fn secret_key(mut self, key: SecretKey) -> Self {
self.secret_key = Some(key);
self
}
pub fn alpn(mut self, alpn: &[u8]) -> Self {
self.alpn = alpn.to_vec();
self
}
pub async fn bind(self) -> Result<IrohServer> {
let secret_key = match (self.secret_key, self.key_path) {
(Some(key), _) => key,
(None, Some(path)) => load_or_generate_key(&path).await?,
(None, None) => SecretKey::generate(&mut rand::rng()),
};
let endpoint = Endpoint::builder()
.alpns(vec![self.alpn])
.secret_key(secret_key)
.bind()
.await?;
endpoint.online().await;
Ok(IrohServer { endpoint })
}
}
impl Default for IrohServerBuilder {
fn default() -> Self {
Self::new()
}
}
pub struct IrohClientBuilder {
alpn: Vec<u8>,
}
impl IrohClientBuilder {
pub fn new() -> Self {
Self {
alpn: DEFAULT_ALPN.to_vec(),
}
}
pub fn alpn(mut self, alpn: &[u8]) -> Self {
self.alpn = alpn.to_vec();
self
}
pub async fn connect(self, server_id: PublicKey) -> Result<IrohConnection> {
let endpoint = Endpoint::builder().bind().await?;
let addr = EndpointAddr::from(server_id);
let conn = endpoint.connect(addr, &self.alpn).await?;
Ok(IrohConnection {
conn,
_endpoint: endpoint,
})
}
pub async fn connect_str(self, server_id: &str) -> Result<IrohConnection> {
let id: PublicKey = server_id.parse()?;
self.connect(id).await
}
}
impl Default for IrohClientBuilder {
fn default() -> Self {
Self::new()
}
}
pub struct IrohServer {
endpoint: Endpoint,
}
impl IrohServer {
pub fn id(&self) -> PublicKey {
self.endpoint.id()
}
pub fn addr(&self) -> EndpointAddr {
self.endpoint.addr()
}
pub async fn accept(&self) -> Result<Option<IrohConnection>> {
if let Some(incoming) = self.endpoint.accept().await {
let conn = incoming.await?;
return Ok(Some(IrohConnection {
conn,
_endpoint: self.endpoint.clone(),
}));
}
Ok(None)
}
pub fn endpoint(&self) -> &Endpoint {
&self.endpoint
}
}
pub struct IrohConnection {
conn: Connection,
_endpoint: Endpoint,
}
impl IrohConnection {
pub fn remote_id(&self) -> PublicKey {
self.conn.remote_id()
}
pub async fn open_stream(&self) -> Result<IrohStream> {
let (send, recv) = self.conn.open_bi().await?;
Ok(IrohStream { send, recv })
}
pub async fn accept_stream(&self) -> Result<IrohStream> {
let (send, recv) = self.conn.accept_bi().await?;
Ok(IrohStream { send, recv })
}
pub fn connection(&self) -> &Connection {
&self.conn
}
}
pub struct IrohStream {
send: iroh::endpoint::SendStream,
recv: iroh::endpoint::RecvStream,
}
impl IrohStream {
pub async fn write(&mut self, data: &[u8]) -> Result<()> {
self.send.write_all(data).await?;
Ok(())
}
pub async fn write_str(&mut self, data: &str) -> Result<()> {
self.write(data.as_bytes()).await
}
pub async fn read(&mut self, buf: &mut [u8]) -> Result<Option<usize>> {
Ok(self.recv.read(buf).await?)
}
pub async fn read_string(&mut self) -> Result<Option<String>> {
let mut buf = vec![0u8; 4096];
if let Some(n) = self.read(&mut buf).await? {
return Ok(Some(String::from_utf8_lossy(&buf[..n]).to_string()));
}
Ok(None)
}
pub fn split(self) -> (iroh::endpoint::SendStream, iroh::endpoint::RecvStream) {
(self.send, self.recv)
}
}
async fn load_or_generate_key(path: &PathBuf) -> Result<SecretKey> {
if path.exists() {
let bytes = fs::read(path).await?;
let key_bytes: [u8; 32] = bytes
.try_into()
.map_err(|_| anyhow::anyhow!("Invalid key file"))?;
Ok(SecretKey::from_bytes(&key_bytes))
} else {
let key = SecretKey::generate(&mut rand::rng());
fs::write(path, key.to_bytes()).await?;
Ok(key)
}
}
pub fn generate_key() -> SecretKey {
SecretKey::generate(&mut rand::rng())
}
pub async fn save_key(key: &SecretKey, path: impl Into<PathBuf>) -> Result<()> {
fs::write(path.into(), key.to_bytes()).await?;
Ok(())
}
pub async fn load_key(path: impl Into<PathBuf>) -> Result<SecretKey> {
let bytes = fs::read(path.into()).await?;
let key_bytes: [u8; 32] = bytes
.try_into()
.map_err(|_| anyhow::anyhow!("Invalid key file"))?;
Ok(SecretKey::from_bytes(&key_bytes))
}