355 lines
12 KiB
Rust
355 lines
12 KiB
Rust
|
//! Implementation of authentication using iroh hooks
|
||
|
|
//!
|
||
|
|
//! This implements an auth protocol that works with iroh hooks.
|
||
|
|
//! It allows to put authentication in front of iroh protocols. The protocols don't need any special support.
|
||
|
|
//! Authentication is handled prior to establishing the connections, over a separate connection.
|
||
|
|
|
||
|
|
use iroh::{Endpoint, EndpointAddr, endpoint::presets, protocol::Router};
|
||
|
|
use n0_error::{Result, StdResultExt};
|
||
|
|
|
||
|
|
use crate::echo::Echo;
|
||
|
|
|
||
|
|
#[tokio::main]
|
||
|
|
async fn main() -> Result<()> {
|
||
|
|
tracing_subscriber::fmt::init();
|
||
|
|
let server_router = accept_side(b"secret!!").await?;
|
||
|
|
server_router.endpoint().online().await;
|
||
|
|
let server_addr = server_router.endpoint().addr();
|
||
|
|
|
||
|
|
println!("-- no --");
|
||
|
|
let res = connect_side_no_auth(server_addr.clone()).await;
|
||
|
|
println!("echo without auth: {:#}", res.unwrap_err());
|
||
|
|
|
||
|
|
println!("-- wrong --");
|
||
|
|
let res = connect_side(server_addr.clone(), b"dunno").await;
|
||
|
|
println!("echo with wrong auth: {:#}", res.unwrap_err());
|
||
|
|
|
||
|
|
println!("-- correct --");
|
||
|
|
let res = connect_side(server_addr.clone(), b"secret!!").await;
|
||
|
|
println!("echo with correct auth: {res:?}");
|
||
|
|
|
||
|
|
server_router.shutdown().await.anyerr()?;
|
||
|
|
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn connect_side(remote_addr: EndpointAddr, token: &[u8]) -> Result<()> {
|
||
|
|
let (auth_hook, auth_task) = auth::outgoing(token.to_vec());
|
||
|
|
let endpoint = Endpoint::builder(presets::N0)
|
||
|
|
.hooks(auth_hook)
|
||
|
|
.bind()
|
||
|
|
.await?;
|
||
|
|
let _guard = auth_task.spawn(endpoint.clone());
|
||
|
|
Echo::connect(&endpoint, remote_addr, b"hello there!").await
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn connect_side_no_auth(remote_addr: EndpointAddr) -> Result<()> {
|
||
|
|
let endpoint = Endpoint::bind(presets::N0).await?;
|
||
|
|
Echo::connect(&endpoint, remote_addr, b"hello there!").await
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn accept_side(token: &[u8]) -> Result<Router> {
|
||
|
|
let (auth_hook, auth_protocol) = auth::incoming(token.to_vec());
|
||
|
|
let endpoint = Endpoint::builder(presets::N0)
|
||
|
|
.hooks(auth_hook)
|
||
|
|
.bind()
|
||
|
|
.await?;
|
||
|
|
|
||
|
|
let router = Router::builder(endpoint)
|
||
|
|
.accept(auth::ALPN, auth_protocol)
|
||
|
|
.accept(echo::ALPN, Echo)
|
||
|
|
.spawn();
|
||
|
|
|
||
|
|
Ok(router)
|
||
|
|
}
|
||
|
|
|
||
|
|
mod echo {
|
||
|
|
//! A bare-bones protocol with no knowledge of auth whatsoever.
|
||
|
|
|
||
|
|
use iroh::{
|
||
|
|
Endpoint, EndpointAddr,
|
||
|
|
endpoint::Connection,
|
||
|
|
protocol::{AcceptError, ProtocolHandler},
|
||
|
|
};
|
||
|
|
use n0_error::{Result, StdResultExt, anyerr};
|
||
|
|
|
||
|
|
#[derive(Debug, Clone)]
|
||
|
|
pub struct Echo;
|
||
|
|
|
||
|
|
pub const ALPN: &[u8] = b"iroh-example/echo/0";
|
||
|
|
|
||
|
|
impl Echo {
|
||
|
|
pub async fn connect(
|
||
|
|
endpoint: &Endpoint,
|
||
|
|
remote: impl Into<EndpointAddr>,
|
||
|
|
message: &[u8],
|
||
|
|
) -> Result<()> {
|
||
|
|
let conn = endpoint.connect(remote, ALPN).await?;
|
||
|
|
let (mut send, mut recv) = conn.open_bi().await.anyerr()?;
|
||
|
|
send.write_all(message).await.anyerr()?;
|
||
|
|
send.finish().anyerr()?;
|
||
|
|
let response = recv.read_to_end(1000).await.anyerr()?;
|
||
|
|
conn.close(0u32.into(), b"bye!");
|
||
|
|
if response == message {
|
||
|
|
Ok(())
|
||
|
|
} else {
|
||
|
|
Err(anyerr!("Received invalid response"))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
impl ProtocolHandler for Echo {
|
||
|
|
async fn accept(&self, connection: Connection) -> Result<(), AcceptError> {
|
||
|
|
let (mut send, mut recv) = connection.accept_bi().await?;
|
||
|
|
tokio::io::copy(&mut recv, &mut send).await?;
|
||
|
|
send.finish()?;
|
||
|
|
connection.closed().await;
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
mod auth {
|
||
|
|
//! Authentication hook
|
||
|
|
|
||
|
|
use std::{
|
||
|
|
collections::{HashMap, HashSet, hash_map},
|
||
|
|
sync::{Arc, Mutex},
|
||
|
|
};
|
||
|
|
|
||
|
|
use iroh::{
|
||
|
|
Endpoint, EndpointAddr, EndpointId,
|
||
|
|
endpoint::{
|
||
|
|
AfterHandshakeOutcome, BeforeConnectOutcome, Connection, ConnectionError, EndpointHooks,
|
||
|
|
},
|
||
|
|
protocol::{AcceptError, ProtocolHandler},
|
||
|
|
};
|
||
|
|
use n0_error::{AnyError, Result, StackResultExt, StdResultExt, anyerr};
|
||
|
|
use n0_future::task::AbortOnDropHandle;
|
||
|
|
use tokio::{
|
||
|
|
sync::{mpsc, oneshot},
|
||
|
|
task::JoinSet,
|
||
|
|
};
|
||
|
|
use tracing::debug;
|
||
|
|
|
||
|
|
pub const ALPN: &[u8] = b"iroh-example/auth/0";
|
||
|
|
|
||
|
|
const CLOSE_ACCEPTED: u32 = 1;
|
||
|
|
const CLOSE_DENIED: u32 = 403;
|
||
|
|
|
||
|
|
/// Outgoing side: Use this if you want to pre-auth outgoing connections.
|
||
|
|
pub fn outgoing(token: Vec<u8>) -> (OutgoingAuthHook, OutgoingAuthTask) {
|
||
|
|
let (tx, rx) = mpsc::channel(16);
|
||
|
|
let hook = OutgoingAuthHook { tx };
|
||
|
|
let connector = OutgoingAuthTask {
|
||
|
|
token,
|
||
|
|
rx,
|
||
|
|
allowed_remotes: Default::default(),
|
||
|
|
pending_remotes: Default::default(),
|
||
|
|
tasks: JoinSet::new(),
|
||
|
|
};
|
||
|
|
(hook, connector)
|
||
|
|
}
|
||
|
|
|
||
|
|
type AuthResult = Result<(), Arc<AnyError>>;
|
||
|
|
|
||
|
|
/// Hook to mount on the endpoint builder.
|
||
|
|
#[derive(Debug)]
|
||
|
|
pub struct OutgoingAuthHook {
|
||
|
|
tx: mpsc::Sender<(EndpointId, oneshot::Sender<AuthResult>)>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl OutgoingAuthHook {
|
||
|
|
async fn authenticate(&self, remote_id: EndpointId) -> Result<()> {
|
||
|
|
let (tx, rx) = oneshot::channel();
|
||
|
|
self.tx
|
||
|
|
.send((remote_id, tx))
|
||
|
|
.await
|
||
|
|
.std_context("authenticator stopped")?;
|
||
|
|
rx.await
|
||
|
|
.std_context("authenticator stopped")?
|
||
|
|
.context("failed to authenticate")
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
impl EndpointHooks for OutgoingAuthHook {
|
||
|
|
async fn before_connect<'a>(
|
||
|
|
&'a self,
|
||
|
|
remote_addr: &'a EndpointAddr,
|
||
|
|
alpn: &'a [u8],
|
||
|
|
) -> BeforeConnectOutcome {
|
||
|
|
// Don't intercept auth request themsevles
|
||
|
|
if alpn == ALPN {
|
||
|
|
BeforeConnectOutcome::Accept
|
||
|
|
} else {
|
||
|
|
match self.authenticate(remote_addr.id).await {
|
||
|
|
Ok(()) => BeforeConnectOutcome::Accept,
|
||
|
|
Err(err) => {
|
||
|
|
debug!("authentication denied: {err:#}");
|
||
|
|
BeforeConnectOutcome::Reject
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Connector task that initiates pre-auth request. Call [`Self::spawn`] once the endpoint is built.
|
||
|
|
pub struct OutgoingAuthTask {
|
||
|
|
token: Vec<u8>,
|
||
|
|
rx: mpsc::Receiver<(EndpointId, oneshot::Sender<AuthResult>)>,
|
||
|
|
allowed_remotes: HashSet<EndpointId>,
|
||
|
|
pending_remotes: HashMap<EndpointId, Vec<oneshot::Sender<AuthResult>>>,
|
||
|
|
tasks: JoinSet<(EndpointId, Result<()>)>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl OutgoingAuthTask {
|
||
|
|
pub fn spawn(self, endpoint: Endpoint) -> AbortOnDropHandle<()> {
|
||
|
|
AbortOnDropHandle::new(tokio::spawn(self.run(endpoint)))
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn run(mut self, endpoint: Endpoint) {
|
||
|
|
loop {
|
||
|
|
tokio::select! {
|
||
|
|
msg = self.rx.recv() => {
|
||
|
|
let Some((remote_id, tx)) = msg else {
|
||
|
|
break;
|
||
|
|
};
|
||
|
|
self.handle_msg(&endpoint, remote_id, tx);
|
||
|
|
}
|
||
|
|
Some(res) = self.tasks.join_next(), if !self.tasks.is_empty() => {
|
||
|
|
let (remote_id, res) = res.expect("connect task panicked");
|
||
|
|
let res = res.map_err(Arc::new);
|
||
|
|
self.handle_task(remote_id, res);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn handle_msg(
|
||
|
|
&mut self,
|
||
|
|
endpoint: &Endpoint,
|
||
|
|
remote_id: EndpointId,
|
||
|
|
tx: oneshot::Sender<Result<(), Arc<AnyError>>>,
|
||
|
|
) {
|
||
|
|
if self.allowed_remotes.contains(&remote_id) {
|
||
|
|
tx.send(Ok(())).ok();
|
||
|
|
} else {
|
||
|
|
match self.pending_remotes.entry(remote_id) {
|
||
|
|
hash_map::Entry::Occupied(mut entry) => {
|
||
|
|
entry.get_mut().push(tx);
|
||
|
|
}
|
||
|
|
hash_map::Entry::Vacant(entry) => {
|
||
|
|
let endpoint = endpoint.clone();
|
||
|
|
let token = self.token.clone();
|
||
|
|
self.tasks.spawn(async move {
|
||
|
|
let res = Self::connect(endpoint, remote_id, token).await;
|
||
|
|
(remote_id, res)
|
||
|
|
});
|
||
|
|
entry.insert(vec![tx]);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
fn handle_task(&mut self, remote_id: EndpointId, res: Result<(), Arc<AnyError>>) {
|
||
|
|
if res.is_ok() {
|
||
|
|
self.allowed_remotes.insert(remote_id);
|
||
|
|
}
|
||
|
|
let senders = self.pending_remotes.remove(&remote_id);
|
||
|
|
for tx in senders.into_iter().flatten() {
|
||
|
|
tx.send(res.clone()).ok();
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn connect(endpoint: Endpoint, remote_id: EndpointId, token: Vec<u8>) -> Result<()> {
|
||
|
|
let conn = endpoint.connect(remote_id, ALPN).await?;
|
||
|
|
let mut stream = conn.open_uni().await.anyerr()?;
|
||
|
|
stream.write_all(&token).await.anyerr()?;
|
||
|
|
stream.finish().anyerr()?;
|
||
|
|
let reason = conn.closed().await;
|
||
|
|
if let ConnectionError::ApplicationClosed(code) = &reason
|
||
|
|
&& code.error_code.into_inner() as u32 == CLOSE_ACCEPTED
|
||
|
|
{
|
||
|
|
Ok(())
|
||
|
|
} else if let ConnectionError::ApplicationClosed(code) = &reason
|
||
|
|
&& code.error_code.into_inner() as u32 == CLOSE_DENIED
|
||
|
|
{
|
||
|
|
Err(anyerr!("authentication denied by remote"))
|
||
|
|
} else {
|
||
|
|
Err(AnyError::from_std(reason))
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Incoming side: Use this if you want to only accept connections from peers with successful pre-auth requests.
|
||
|
|
pub fn incoming(token: Vec<u8>) -> (IncomingAuthHook, AuthProtocol) {
|
||
|
|
let allowed_remotes: Arc<Mutex<HashSet<EndpointId>>> = Default::default();
|
||
|
|
let hook = IncomingAuthHook {
|
||
|
|
allowed_remotes: allowed_remotes.clone(),
|
||
|
|
};
|
||
|
|
let protocol = AuthProtocol {
|
||
|
|
allowed_remotes,
|
||
|
|
token,
|
||
|
|
};
|
||
|
|
(hook, protocol)
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Accept-side auth hook: Mount this onto the endpoint.
|
||
|
|
///
|
||
|
|
/// This will reject incoming connections if the remote did not successfully authenticate before.
|
||
|
|
#[derive(Debug)]
|
||
|
|
pub struct IncomingAuthHook {
|
||
|
|
allowed_remotes: Arc<Mutex<HashSet<EndpointId>>>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl EndpointHooks for IncomingAuthHook {
|
||
|
|
async fn after_handshake<'a>(
|
||
|
|
&'a self,
|
||
|
|
conn: &'a iroh::endpoint::Connection,
|
||
|
|
) -> AfterHandshakeOutcome {
|
||
|
|
if conn.alpn() == ALPN
|
||
|
|
|| self
|
||
|
|
.allowed_remotes
|
||
|
|
.lock()
|
||
|
|
.expect("poisoned")
|
||
|
|
.contains(&conn.remote_id())
|
||
|
|
{
|
||
|
|
AfterHandshakeOutcome::Accept
|
||
|
|
} else {
|
||
|
|
AfterHandshakeOutcome::Reject {
|
||
|
|
error_code: 403u32.into(),
|
||
|
|
reason: b"not authenticated".to_vec(),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Accept-side auth protocol. Mount this on the router to accept authentication requests.
|
||
|
|
#[derive(Debug, Clone)]
|
||
|
|
pub struct AuthProtocol {
|
||
|
|
token: Vec<u8>,
|
||
|
|
allowed_remotes: Arc<Mutex<HashSet<EndpointId>>>,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl ProtocolHandler for AuthProtocol {
|
||
|
|
/// The `accept` method is called for each incoming connection for our ALPN.
|
||
|
|
///
|
||
|
|
/// The returned future runs on a newly spawned tokio task, so it can run as long as
|
||
|
|
/// the connection lasts.
|
||
|
|
async fn accept(&self, connection: Connection) -> Result<(), AcceptError> {
|
||
|
|
let mut stream = connection.accept_uni().await?;
|
||
|
|
let token = stream.read_to_end(256).await.anyerr()?;
|
||
|
|
let remote_id = connection.remote_id();
|
||
|
|
if token == self.token {
|
||
|
|
self.allowed_remotes
|
||
|
|
.lock()
|
||
|
|
.expect("poisoned")
|
||
|
|
.insert(remote_id);
|
||
|
|
connection.close(CLOSE_ACCEPTED.into(), b"accepted");
|
||
|
|
} else {
|
||
|
|
connection.close(CLOSE_DENIED.into(), b"rejected");
|
||
|
|
}
|
||
|
|
Ok(())
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|