[puter] add wisp v1 fallback support

Author: r58Playz

server/src/config.rs

@@ -25,8 +25,7 @@ use wisp_mux::{
 };
 
 use crate::{
-	handle::wisp::utils::get_certificates_from_paths, puter::PuterPasswordProtocolExtensionBuilder,
-	CLI, CONFIG, RESOLVER,
+	CLI, CONFIG, RESOLVER, handle::wisp::utils::get_certificates_from_paths, puter::{PuterPasswordProtocolExtension, PuterPasswordProtocolExtensionBuilder}
 };
 
 pub const VERSION_STRING: &str = concat!(
@@ -290,6 +289,8 @@ struct ConfigCache {
 	pub blocked_udp_hosts: RegexSet,
 
 	pub socks5_server: Option<SocketAddr>,
+
+	pub puter_auth_server: Option<Url>,
 }
 
 lazy_static! {
@@ -319,6 +320,8 @@ lazy_static! {
 			blocked_udp_hosts: RegexSet::new(&CONFIG.stream.block_udp_hosts).unwrap(),
 
 			socks5_server: CONFIG.stream.socks5_server.as_ref().map(|x| x.to_socket_addrs().expect("failed to resolve socks5 server").next().expect("failed to resolve socks5 server")),
+
+			puter_auth_server: CONFIG.wisp.puter_auth_server.as_ref().map(|x| Url::parse(x).expect("failed to parse puter auth server")),
 		}
 	};
 }
@@ -406,6 +409,11 @@ impl WispConfig {
 		self.extensions.contains(&ProtocolExtension::Wispnet)
 	}
 
+	#[doc(hidden)]
+	pub fn puter_auth_server(&self) -> Option<&Url> {
+		CONFIG_CACHE.puter_auth_server.as_ref()
+	}
+
 	#[doc(hidden)]
 	pub async fn to_opts(&self) -> anyhow::Result<(Option<WispV2Handshake>, u32)> {
 		if self.wisp_v2 {
@@ -432,7 +440,8 @@ impl WispConfig {
 							.context("failed to parse puter auth server")?,
 						true,
 					),
-				))
+				));
+				required_extensions.push(PuterPasswordProtocolExtension::ID);
 			}
 
 			match self.auth_extension {

server/src/puter.rs

@@ -14,6 +14,36 @@ use wisp_mux::{
 
 use crate::REQWEST_CLIENT;
 
+#[derive(Serialize)]
+struct AuthRequest {
+	token: String,
+}
+#[derive(Deserialize)]
+struct AuthResponse {
+	#[serde(default)]
+	allow: bool,
+}
+
+pub async fn verify_relay_token(endpoint: &Url, token: &str) -> anyhow::Result<bool> {
+	let origin = endpoint.origin().ascii_serialization();
+
+	let res = REQWEST_CLIENT
+		.request(Method::POST, endpoint.clone())
+		.header("Content-Type", "application/json")
+		.header("Origin", origin)
+		.json(&AuthRequest {
+			token: token.to_string(),
+		})
+		.send()
+		.await
+		.context("failed to ask auth server for auth")?
+		.json::<AuthResponse>()
+		.await
+		.context("auth server gave invalid response")?;
+
+	Ok(res.allow)
+}
+
 /// ID of Puter password protocol extension.
 pub const PUTER_PASSWORD_PROTOCOL_EXTENSION_ID: u8 = 0x02;
 
@@ -55,15 +85,6 @@ impl PuterPasswordProtocolExtension {
 	pub const ID: u8 = PUTER_PASSWORD_PROTOCOL_EXTENSION_ID;
 }
 
-#[derive(Serialize)]
-struct AuthRequest {
-	token: String,
-}
-#[derive(Deserialize)]
-struct AuthResponse {
-	allow: bool,
-}
-
 #[async_trait]
 impl ProtocolExtension for PuterPasswordProtocolExtension {
 	fn get_id(&self) -> u8 {
@@ -81,26 +102,12 @@ impl ProtocolExtension for PuterPasswordProtocolExtension {
 				chosen_password,
 				..
 			} => {
-				let origin = endpoint.origin().ascii_serialization();
-
-				let res = REQWEST_CLIENT
-					.request(Method::POST, endpoint.clone())
-					.header("Content-Type", "application/json")
-					.header("Origin", origin)
-					.json(&AuthRequest {
-						token: chosen_password.clone(),
-					})
-					.send()
+				if verify_relay_token(endpoint, &chosen_password)
 					.await
-					.context("failed to ask auth server for auth")
+					.context("failed to verify relay token")
 					.map_err(|x| WispError::ExtensionImplError(x.into()))?
-					.json::<AuthResponse>()
-					.await
-					.context("auth server gave invalid response")
-					.map_err(|x| WispError::ExtensionImplError(x.into()))?;
-
-				if res.allow {
-                    Ok(None)
+				{
+					Ok(None)
 				} else {
 					Ok(Some((
 						CloseReason::ExtensionsPasswordAuthFailed,

server/src/route.rs

@@ -20,6 +20,7 @@ use crate::{
 	config::SocketTransport,
 	generate_stats,
 	listener::{ServerStream, ServerStreamExt, ServerStreamRead, ServerStreamWrite},
+	puter::verify_relay_token,
 	stream::WebSocketStreamWrapper,
 	upgrade::{is_upgrade_request, upgrade},
 	util_chain::{chain, Chain},
@@ -140,7 +141,31 @@ where
 	};
 
 	let ws_protocol = headers.get(SEC_WEBSOCKET_PROTOCOL);
-	let req_path = req.uri().path().to_string();
+	let mut req_path = req.uri().path().to_string();
+
+	if let Some(server) = CONFIG.wisp.puter_auth_server() {
+		let endpoint = CONFIG.wisp.prefix.clone() + "/";
+		if req_path != endpoint {
+			let trimmed = req_path.strip_prefix("/").unwrap_or(&req_path);
+
+			let Some(token_loc) = trimmed.find('/') else {
+				debug!("sent non_ws_response to http client [no token found]");
+				return non_ws_resp();
+			};
+
+			if !verify_relay_token(server, &trimmed[..token_loc])
+				.await
+				.context("failed to verify relay token")?
+			{
+				debug!("sent non_ws_response to http client [token invalid]");
+				return Ok(Response::builder()
+					.status(StatusCode::UNAUTHORIZED)
+					.body(Body::new(CONFIG.server.non_ws_response.as_bytes().into()))?);
+			}
+
+			req_path = trimmed[token_loc..].to_string();
+		}
+	}
 
 	if req_path.ends_with(&(CONFIG.wisp.prefix.clone() + "/")) {
 		let has_ws_protocol = ws_protocol.is_some();