Merge pull request #55 from bvanneerven/master

Add `token-from-header` example for `axum`

Author: sunli829

Cargo.toml

@@ -31,6 +31,7 @@ members = [
     "axum/starwars",
     "axum/subscription",
     "axum/upload",
+    "axum/token-from-header",
 
     "tide/starwars",
     "tide/dataloader",

axum/token-from-header/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "axum-token-from-header"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+async-graphql = { path = "../../.." }
+async-graphql-axum = { path = "../../../integrations/axum" }
+tokio = { version = "1.8", features = ["macros", "rt-multi-thread"] }
+token = { path = "../../models/token" }
+hyper = "0.14"
+axum = { version = "0.5.1", features = ["ws", "headers"] }

axum/token-from-header/src/main.rs

@@ -0,0 +1,75 @@
+use async_graphql::{
+    http::{playground_source, GraphQLPlaygroundConfig, ALL_WEBSOCKET_PROTOCOLS},
+    Data, EmptyMutation, Schema,
+};
+use async_graphql_axum::{GraphQLProtocol, GraphQLRequest, GraphQLResponse, GraphQLWebSocket};
+use axum::{
+    extract::{ws::WebSocketUpgrade, Extension},
+    http::header::HeaderMap,
+    response::{Html, IntoResponse, Response},
+    routing::get,
+    Router, Server,
+};
+use token::{on_connection_init, QueryRoot, SubscriptionRoot, Token, TokenSchema};
+
+async fn graphql_playground() -> impl IntoResponse {
+    Html(playground_source(
+        GraphQLPlaygroundConfig::new("/").subscription_endpoint("/ws"),
+    ))
+}
+
+fn get_token_from_headers(headers: &HeaderMap) -> Option<Token> {
+    headers
+        .get("Token")
+        .and_then(|value| value.to_str().map(|s| Token(s.to_string())).ok())
+}
+
+async fn graphql_handler(
+    req: GraphQLRequest,
+    Extension(schema): Extension<TokenSchema>,
+    headers: HeaderMap,
+) -> GraphQLResponse {
+    let mut req = req.into_inner();
+    if let Some(token) = get_token_from_headers(&headers) {
+        req = req.data(token);
+    }
+    schema.execute(req).await.into()
+}
+
+async fn graphql_ws_handler(
+    Extension(schema): Extension<TokenSchema>,
+    protocol: GraphQLProtocol,
+    websocket: WebSocketUpgrade,
+    headers: HeaderMap,
+) -> Response {
+    let mut data = Data::default();
+    if let Some(token) = get_token_from_headers(&headers) {
+        data.insert(token);
+    }
+
+    websocket
+        .protocols(ALL_WEBSOCKET_PROTOCOLS)
+        .on_upgrade(move |stream| {
+            GraphQLWebSocket::new(stream, schema.clone(), protocol)
+                .with_data(data)
+                .on_connection_init(on_connection_init)
+                .serve()
+        })
+}
+
+#[tokio::main]
+async fn main() {
+    let schema = Schema::new(QueryRoot, EmptyMutation, SubscriptionRoot);
+
+    let app = Router::new()
+        .route("/", get(graphql_playground).post(graphql_handler))
+        .route("/ws", get(graphql_ws_handler))
+        .layer(Extension(schema));
+
+    println!("Playground: http://localhost:8000");
+
+    Server::bind(&"0.0.0.0:8000".parse().unwrap())
+        .serve(app.into_make_service())
+        .await
+        .unwrap();
+}