From 9f214fbbff6324af3d834c2ed473862861428556 Mon Sep 17 00:00:00 2001
From: "Aode (lion)" <asonix@asonix.dog>
Date: Thu, 18 Nov 2021 10:18:06 -0600
Subject: [PATCH] Support arbitrary schemes

---
 examples/fetch.rs    |  9 ++++++++-
 examples/resolver.rs |  3 ++-
 src/lib.rs           | 30 +++++++++++++++++++++++++-----
 3 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/examples/fetch.rs b/examples/fetch.rs
index 626ad7a..1bfcf83 100644
--- a/examples/fetch.rs
+++ b/examples/fetch.rs
@@ -5,7 +5,14 @@ use std::error::Error;
 #[actix_rt::main]
 async fn main() -> Result<(), Box<dyn Error>> {
     let client = Client::default();
-    let wf = Webfinger::fetch(&client, "asonix@localhost:8000", "localhost:8000", false).await?;
+    let wf = Webfinger::fetch(
+        &client,
+        Some("acct:"),
+        "asonix@localhost:8000",
+        "localhost:8000",
+        false,
+    )
+    .await?;
 
     println!("asonix's webfinger:\n{:#?}", wf);
     Ok(())
diff --git a/examples/resolver.rs b/examples/resolver.rs
index f1e4542..f437335 100644
--- a/examples/resolver.rs
+++ b/examples/resolver.rs
@@ -14,11 +14,12 @@ impl Resolver for MyResolver {
     type Error = actix_web::error::JsonPayloadError;
 
     fn find(
+        scheme: Option<&str>,
         account: &str,
         domain: &str,
         state: Data<MyState>,
     ) -> Pin<Box<dyn Future<Output = Result<Option<Webfinger>, Self::Error>>>> {
-        let w = if domain == state.domain {
+        let w = if scheme == Some("acct:") && domain == state.domain {
             Some(Webfinger::new(&format!("{}@{}", account, domain)))
         } else {
             None
diff --git a/src/lib.rs b/src/lib.rs
index 1679b05..257cfad 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -29,7 +29,7 @@
 //! #[actix_rt::main]
 //! async fn main() -> Result<(), Box<dyn Error>> {
 //!     let client = Client::default();
-//!     let wf = Webfinger::fetch(&client, "asonix@asonix.dog", "localhost:8000", false).await?;
+//!     let wf = Webfinger::fetch(&client, None, "asonix@asonix.dog", "localhost:8000", false).await?;
 //!
 //!     println!("asonix's webfinger:\n{:#?}", wf);
 //!     Ok(())
@@ -249,6 +249,7 @@ pub struct InvalidResource(String);
 /// formatting before the request reaches the route handler.
 #[derive(Clone, Debug)]
 pub struct WebfingerResource {
+    pub scheme: Option<String>,
     pub account: String,
     pub domain: String,
 }
@@ -257,12 +258,24 @@ impl std::str::FromStr for WebfingerResource {
     type Err = InvalidResource;
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
-        let trimmed = s.trim_start_matches("acct:").trim_start_matches('@');
+        let (scheme, trimmed) = s
+            .find(':')
+            .map(|index| {
+                let (scheme, trimmed) = s.split_at(index);
+                (
+                    Some(scheme.to_owned() + ":"),
+                    trimmed.trim_start_matches(':'),
+                )
+            })
+            .unwrap_or((None, s));
+
+        let trimmed = trimmed.trim_start_matches('@');
 
         if let Some(index) = trimmed.find('@') {
             let (account, domain) = trimmed.split_at(index);
 
             Ok(WebfingerResource {
+                scheme,
                 account: account.to_owned(),
                 domain: domain.trim_start_matches('@').to_owned(),
             })
@@ -353,6 +366,7 @@ pub trait Resolver {
     type Error: ResponseError + 'static;
 
     fn find(
+        scheme: Option<&str>,
         account: &str,
         domain: &str,
         state: Self::State,
@@ -366,10 +380,14 @@ pub fn endpoint<R>(
 where
     R: Resolver,
 {
-    let WebfingerResource { account, domain } = query.into_inner().resource;
+    let WebfingerResource {
+        scheme,
+        account,
+        domain,
+    } = query.into_inner().resource;
 
     Box::pin(async move {
-        match R::find(&account, &domain, state).await? {
+        match R::find(scheme.as_deref(), &account, &domain, state).await? {
             Some(w) => Ok(w.respond()),
             None => Ok(HttpResponse::NotFound().finish()),
         }
@@ -611,14 +629,16 @@ impl Webfinger {
     /// rather this library generating it's own http clients.
     pub async fn fetch(
         client: &Client,
+        scheme: Option<&str>,
         user: &str,
         domain: &str,
         https: bool,
     ) -> Result<Self, FetchError> {
         let url = format!(
-            "{}://{}/.well-known/webfinger?resource=acct:{}",
+            "{}://{}/.well-known/webfinger?resource={}{}",
             if https { "https" } else { "http" },
             domain,
+            scheme.unwrap_or("acct:"),
             user
         );