parakeet / did-resolver / src / lib.rs

use error::Error;
use hickory_resolver::proto::serialize::binary::BinEncodable;
use hickory_resolver::TokioAsyncResolver;
use reqwest::header::CONTENT_TYPE;
use reqwest::{Client, StatusCode};

pub mod error;
pub mod types;

const DEFAULT_DUR: std::time::Duration = std::time::Duration::from_secs(3);
const DEFAULT_PLC: &str = "https://plc.directory";

#[derive(Default)]
pub struct ResolverOpts {
    pub plc_directory: Option<String>,
    pub timeout: Option<std::time::Duration>,
    pub user_agent: Option<String>,
}

pub struct Resolver {
    client: Client,
    plc: String,
    dns: TokioAsyncResolver,
}

impl Resolver {
    pub fn new(opts: ResolverOpts) -> Result<Self, Error> {
        let dns = hickory_resolver::AsyncResolver::tokio_from_system_conf()?;

        let mut client = Client::builder().timeout(opts.timeout.unwrap_or(DEFAULT_DUR));

        if let Some(user_agent) = opts.user_agent {
            client = client.user_agent(user_agent);
        }

        Ok(Resolver {
            client: client.build()?,
            plc: opts.plc_directory.unwrap_or(DEFAULT_PLC.to_string()),
            dns,
        })
    }

    pub async fn resolve_did(&self, did: &str) -> Result<Option<types::DidDocument>, Error> {
        let did_parts = did.splitn(3, ':').collect::<Vec<&str>>();

        if did_parts.len() != 3 {
            return Err(Error::BadDidFormat);
        }

        if did_parts[0] != "did" {
            return Err(Error::BadDidFormat);
        }

        match did_parts[1] {
            "plc" => self.resolve_did_plc(did).await,
            "web" => self.resolve_did_web(did_parts[2]).await,
            method => Err(Error::UnsupportedDidMethod(method.to_string())),
        }
    }

    async fn resolve_did_plc(&self, did: &str) -> Result<Option<types::DidDocument>, Error> {
        let res = self
            .client
            .get(format!("{}/{did}", self.plc))
            .send()
            .await?;

        let status = res.status();

        if status.is_server_error() {
            return Err(Error::ServerError);
        }

        if status == StatusCode::NOT_FOUND || status == StatusCode::GONE {
            return Ok(None);
        }

        let did_doc = res.json().await?;
        Ok(Some(did_doc))
    }

    async fn resolve_did_web(&self, id: &str) -> Result<Option<types::DidDocument>, Error> {
        let res = match self
            .client
            .get(format!("https://{id}/.well-known/did.json"))
            .send()
            .await
        {
            Ok(res) => res,
            Err(err) => {
                if err.is_timeout() {
                    return Err(Error::Timeout);
                } else if err.is_redirect() {
                    return Err(Error::TooManyRedirects);
                }
                return Err(Error::Http(err));
            }
        };

        let status = res.status();

        if status.is_server_error() {
            return Err(Error::ServerError);
        }

        if status == StatusCode::NOT_FOUND {
            return Ok(None);
        }

        let did_doc = res.json().await?;
        Ok(Some(did_doc))
    }

    pub async fn resolve_handle(&self, handle: &str) -> Result<Option<String>, Error> {
        // we want one of these to succeed
        let http = self.resolve_handle_http(handle);
        let dns = self.resolve_handle_dns(handle);

        match tokio::join!(http, dns) {
            (Ok(http), Ok(dns)) => Ok(dns.or(http)),
            (http, dns) => http.or(dns),
        }
    }

    async fn resolve_handle_http(&self, handle: &str) -> Result<Option<String>, Error> {
        let res = self
            .client
            .get(format!("https://{handle}/.well-known/atproto-did"))
            .send()
            .await?
            .error_for_status()?;

        if let Some(ct) = res
            .headers()
            .get(CONTENT_TYPE)
            .and_then(|ct| ct.to_str().ok())
        {
            if !ct.starts_with("text/plain") {
                return Ok(None);
            }
        }

        let did = res.text().await?;
        if !did.starts_with("did:") {
            return Ok(None);
        }

        Ok(Some(did))
    }

    async fn resolve_handle_dns(&self, handle: &str) -> Result<Option<String>, Error> {
        let res = match self.dns.txt_lookup(format!("_atproto.{handle}.")).await {
            Ok(txt) => txt,
            Err(err) => {
                return match err.kind() {
                    hickory_resolver::error::ResolveErrorKind::NoRecordsFound { .. } => Ok(None),
                    _ => Err(err.into()),
                }
            }
        };

        let Some(first) = res.as_lookup().records().first() else {
            return Ok(None);
        };

        let Some(Ok(data)) = first.data().and_then(|v| v.as_txt()).map(|v| v.to_bytes()) else {
            return Ok(None);
        };

        let Some(string_data) = String::from_utf8_lossy(&data)
            .strip_prefix("$did=")
            .map(|v| v.to_string())
        else {
            return Ok(None);
        };

        Ok(Some(string_data))
    }
}