Add support for secret_token in built-in webhooks

Former-commit-id: 8806cb9d78

Cargo.toml

@@ -14,7 +14,8 @@ exclude = ["media"]
 [features]
 default = ["native-tls", "ctrlc_handler", "teloxide-core/default",  "auto-send"]
 
-webhooks-axum = ["axum", "tower", "tower-http"]
+webhooks = ["rand"]
+webhooks-axum = ["webhooks", "axum", "tower", "tower-http"]
 
 sqlite-storage = ["sqlx"]
 redis-storage = ["redis"]
@@ -92,6 +93,7 @@ bincode = { version = "1.3", optional = true }
 axum = { version = "0.4.8", optional = true }
 tower = { version = "0.4.12", optional = true }
 tower-http = { version = "0.2.5", features = ["trace"], optional = true }
+rand = { version = "0.8.5", optional = true }
 
 [dev-dependencies]
 rand = "0.8.3"

src/dispatching/update_listeners.rs

@@ -27,7 +27,7 @@
 
 /// Implementations of webhook update listeners - an alternative (to
 /// [`fn@polling`]) way of receiving updates from telegram.
-#[cfg(any(feature = "webhooks-axum"))]
+#[cfg(feature = "webhooks")]
 pub mod webhooks;
 
 use futures::Stream;

src/dispatching/update_listeners/webhooks.rs

@@ -42,13 +42,28 @@ pub struct Options {
     ///
     /// Default - false.
     pub drop_pending_updates: bool,
+
+    /// A secret token to be sent in a header “X-Telegram-Bot-Api-Secret-Token”
+    /// in every webhook request, 1-256 characters. Only characters `A-Z`,
+    /// `a-z`, `0-9`, `_` and `-` are allowed. The header is useful to ensure
+    /// that the request comes from a webhook set by you.
+    ///
+    /// Default - teloxide will generate a random token.
+    pub secret_token: Option<String>,
 }
 
 impl Options {
     /// Construct a new webhook options, see [`Options::address`] and
     /// [`Options::url`] for details.
     pub fn new(address: SocketAddr, url: url::Url) -> Self {
-        Self { address, url, certificate: None, max_connections: None, drop_pending_updates: false }
+        Self {
+            address,
+            url,
+            certificate: None,
+            max_connections: None,
+            drop_pending_updates: false,
+            secret_token: None,
+        }
     }
 
     /// Upload your public key certificate so that the root certificate in use
@@ -71,6 +86,32 @@ impl Options {
     pub fn drop_pending_updates(self) -> Self {
         Self { drop_pending_updates: true, ..self }
     }
+
+    /// A secret token to be sent in a header “X-Telegram-Bot-Api-Secret-Token”
+    /// in every webhook request, 1-256 characters. Only characters `A-Z`,
+    /// `a-z`, `0-9`, `_` and `-` are allowed. The header is useful to ensure
+    /// that the request comes from a webhook set by you.
+    ///
+    /// ## Panics
+    ///
+    /// If the token is invalid.
+    #[track_caller]
+    pub fn secret_token(self, token: String) -> Self {
+        check_secret(token.as_bytes()).expect("Invalid secret token");
+
+        Self { secret_token: Some(token), ..self }
+    }
+
+    /// Returns `self.secret_token`, generating a new one if it's `None`.
+    ///
+    /// After a call to this function `self.secret_token` is always `Some(_)`.
+    ///
+    /// **Note**: if you leave webhook setup to teloxide, it will automatically
+    /// generate a secret token. Call this function only if you need to know the
+    /// secret (for example because you are calling `set_webhook` by yourself).
+    pub fn get_or_gen_secret_token(&mut self) -> &str {
+        self.secret_token.get_or_insert_with(gen_secret_token)
+    }
 }
 
 #[cfg(feature = "webhooks-axum")]
@@ -91,6 +132,7 @@ where
     use crate::requests::Request;
     use teloxide_core::requests::HasPayload;
 
+    let secret = options.get_or_gen_secret_token().to_owned();
     let &mut Options {
         ref url, ref mut certificate, max_connections, drop_pending_updates, ..
     } = options;
@@ -99,12 +141,49 @@ where
     req.payload_mut().certificate = certificate.take();
     req.payload_mut().max_connections = max_connections;
     req.payload_mut().drop_pending_updates = Some(drop_pending_updates);
+    req.payload_mut().secret_token = Some(secret);
 
     req.send().await?;
 
     Ok(())
 }
 
+/// Generates a random string consisting of 32 characters (`a-z`, `A-Z`, `0-9`,
+/// `_` and `-`).
+fn gen_secret_token() -> String {
+    use rand::{distributions::Uniform, Rng};
+    const CHARSET: &[u8] = b"abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-";
+    const SECRET_LENGTH: usize = 32;
+
+    let random = rand::thread_rng()
+        .sample_iter(Uniform::new(0, CHARSET.len()))
+        .map(|idx| CHARSET[idx] as char)
+        .take(SECRET_LENGTH);
+
+    let mut secret = String::with_capacity(SECRET_LENGTH);
+    secret.extend(random);
+
+    secret
+}
+
+fn check_secret(bytes: &[u8]) -> Result<&[u8], &'static str> {
+    let len = bytes.len();
+
+    // Check that length is in bounds
+    if !(1 <= len && len <= 256) {
+        return Err("secret token length must be in range 1..=256");
+    }
+
+    // Check that all characters of the secret are supported by telegram
+    let is_not_supported =
+        |c: &_| !matches!(c, b'a'..=b'z' | b'A'..=b'Z' | b'0'..=b'9' | b'_' | b'-');
+    if bytes.iter().any(is_not_supported) {
+        return Err("secret token must only contain of `a-z`, `A-Z`, `0-9`, `_` and `-` characters");
+    }
+
+    Ok(bytes)
+}
+
 /// Returns first (`.0`) field from a tuple as a `&mut` reference.
 ///
 /// This hack is needed because there isn't currently a way to easily force a

src/dispatching/update_listeners/webhooks/axum.rs

@@ -1,12 +1,14 @@
-use std::convert::Infallible;
+use std::{convert::Infallible, future::Future, pin::Pin};
+
+use axum::{
+    extract::{FromRequest, RequestParts},
+    http::status::StatusCode,
+};
 
 use crate::{
     dispatching::{
         stop_token::{AsyncStopFlag, StopToken},
-        update_listeners::{
-            webhooks::{setup_webhook, tuple_first_mut, Options},
-            UpdateListener,
-        },
+        update_listeners::{webhooks::Options, UpdateListener},
     },
     requests::Requester,
 };
@@ -105,15 +107,12 @@ where
 pub async fn axum_to_router<R>(
     bot: R,
     mut options: Options,
-) -> Result<
-    (impl UpdateListener<Infallible>, impl std::future::Future<Output = ()> + Send, axum::Router),
-    R::Err,
->
+) -> Result<(impl UpdateListener<Infallible>, impl Future<Output = ()> + Send, axum::Router), R::Err>
 where
     R: Requester + Send,
     <R as Requester>::DeleteWebhook: Send,
 {
-    use crate::requests::Request;
+    use crate::{dispatching::update_listeners::webhooks::setup_webhook, requests::Request};
     use futures::FutureExt;
 
     setup_webhook(&bot, &mut options).await?;
@@ -149,12 +148,15 @@ where
 /// function.
 pub fn axum_no_setup(
     options: Options,
-) -> (impl UpdateListener<Infallible>, impl std::future::Future<Output = ()>, axum::Router) {
+) -> (impl UpdateListener<Infallible>, impl Future<Output = ()>, axum::Router) {
     use crate::{
-        dispatching::{stop_token::AsyncStopToken, update_listeners},
+        dispatching::{
+            stop_token::AsyncStopToken,
+            update_listeners::{self, webhooks::tuple_first_mut},
+        },
         types::Update,
     };
-    use axum::{extract::Extension, http::StatusCode, response::IntoResponse, routing::post};
+    use axum::{extract::Extension, response::IntoResponse, routing::post};
     use tokio::sync::mpsc;
     use tokio_stream::wrappers::UnboundedReceiverStream;
     use tower::ServiceBuilder;
@@ -167,9 +169,16 @@ pub fn axum_no_setup(
 
     async fn telegram_request(
         input: String,
+        secret_header: XTelegramBotApiSecretToken,
+        secret: Extension<Option<String>>,
         tx: Extension<CSender>,
         flag: Extension<AsyncStopFlag>,
     ) -> impl IntoResponse {
+        // FIXME: use constant time comparison here
+        if secret_header.0.as_deref() != secret.as_deref().map(str::as_bytes) {
+            return StatusCode::UNAUTHORIZED;
+        }
+
         let tx = match tx.get() {
             None => return StatusCode::SERVICE_UNAVAILABLE,
             // Do not process updates after `.stop()` is called even if the server is still
@@ -206,6 +215,7 @@ pub fn axum_no_setup(
             .layer(TraceLayer::new_for_http())
             .layer(Extension(ClosableSender::new(tx)))
             .layer(Extension(stop_flag.clone()))
+            .layer(Extension(options.secret_token))
             .into_inner(),
     );
 
@@ -245,3 +255,32 @@ impl<T> ClosableSender<T> {
         self.origin.write().unwrap().take();
     }
 }
+
+struct XTelegramBotApiSecretToken(Option<Vec<u8>>);
+
+impl<B> FromRequest<B> for XTelegramBotApiSecretToken {
+    type Rejection = StatusCode;
+
+    fn from_request<'l0, 'at>(
+        req: &'l0 mut RequestParts<B>,
+    ) -> Pin<Box<dyn Future<Output = Result<Self, Self::Rejection>> + Send + 'at>>
+    where
+        'l0: 'at,
+        Self: 'at,
+    {
+        use crate::dispatching::update_listeners::webhooks::check_secret;
+
+        let res = req
+            .headers_mut()
+            .and_then(|map| map.remove("x-telegram-bot-api-secret-token"))
+            .map(|header| {
+                check_secret(header.as_bytes())
+                    .map(<_>::to_owned)
+                    .map_err(|_| StatusCode::BAD_REQUEST)
+            })
+            .transpose()
+            .map(Self);
+
+        Box::pin(async { res }) as _
+    }
+}