stalwart/crates/imap/src/op/authenticate.rs

/*
 * Copyright (c) 2020-2022, Stalwart Labs Ltd.
 *
 * This file is part of the Stalwart IMAP Server.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of
 * the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
 * GNU Affero General Public License for more details.
 * in the LICENSE file at the top-level directory of this distribution.
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 * You can be released from the requirements of the AGPLv3 license by
 * purchasing a commercial license. Please contact licensing@stalw.art
 * for more details.
*/

use std::sync::Arc;

use imap_proto::{
    protocol::{authenticate::Mechanism, capability::Capability},
    receiver::{self, Request},
    Command, ResponseCode, StatusResponse,
};
use mail_parser::decoders::base64::base64_decode;
use mail_send::Credentials;
use tokio::io::AsyncRead;

use crate::core::{Session, SessionData, State};

impl<T: AsyncRead> Session<T> {
    pub async fn handle_authenticate(&mut self, request: Request<Command>) -> crate::OpResult {
        match request.parse_authenticate() {
            Ok(mut args) => match args.mechanism {
                Mechanism::Plain | Mechanism::OAuthBearer => {
                    if !args.params.is_empty() {
                        match base64_decode(args.params.pop().unwrap().as_bytes()) {
                            Some(challenge) => {
                                let result = if args.mechanism == Mechanism::Plain {
                                    decode_challenge_plain(&challenge)
                                } else {
                                    decode_challenge_oauth(&challenge)
                                };

                                match result {
                                    Ok(credentials) => {
                                        self.authenticate(credentials, args.tag).await
                                    }
                                    Err(err) => {
                                        self.write_bytes(
                                            StatusResponse::no(err).with_tag(args.tag).into_bytes(),
                                        )
                                        .await
                                    }
                                }
                            }
                            None => {
                                self.write_bytes(
                                    StatusResponse::no("Failed to decode challenge.")
                                        .with_tag(args.tag)
                                        .with_code(ResponseCode::Parse)
                                        .into_bytes(),
                                )
                                .await
                            }
                        }
                    } else {
                        self.receiver.request = receiver::Request {
                            tag: args.tag,
                            command: Command::Authenticate,
                            tokens: vec![receiver::Token::Argument(args.mechanism.into_bytes())],
                        };
                        self.receiver.state = receiver::State::Argument { last_ch: b' ' };
                        self.write_bytes(b"+ \"\"\r\n".to_vec()).await
                    }
                }
                _ => {
                    self.write_bytes(
                        StatusResponse::no("Authentication mechanism not supported.")
                            .with_tag(args.tag)
                            .with_code(ResponseCode::Cannot)
                            .into_bytes(),
                    )
                    .await
                }
            },
            Err(response) => self.write_bytes(response.into_bytes()).await,
        }
    }

    pub async fn authenticate(
        &mut self,
        credentials: Credentials<String>,
        tag: String,
    ) -> crate::Result<()> {
        // Throttle authentication requests
        if self.jmap.is_auth_allowed(self.remote_addr.clone()).is_err() {
            self.write_bytes(
                StatusResponse::bye("Too many authentication requests from this IP address.")
                    .into_bytes(),
            )
            .await?;
            tracing::debug!(parent: &self.span,
                event = "disconnect",
                "Too many authentication attempts, disconnecting.",
            );
            return Err(());
        }

        // Authenticate
        let access_token = match credentials {
            Credentials::Plain { username, secret } | Credentials::XOauth2 { username, secret } => {
                self.jmap.authenticate_plain(&username, &secret).await
            }
            Credentials::OAuthBearer { token } => {
                match self
                    .jmap
                    .validate_access_token("access_token", &token)
                    .await
                {
                    Ok((account_id, _, _)) => self.jmap.get_access_token(account_id).await,
                    Err(err) => {
                        tracing::debug!(
                            parent: &self.span,
                            context = "authenticate",
                            err = err,
                            "Failed to validate access token."
                        );
                        None
                    }
                }
            }
        };

        if let Some(access_token) = access_token {
            // Enforce concurrency limits
            let in_flight = self
                .imap
                .get_authenticated_limiter(access_token.primary_id())
                .lock()
                .concurrent_requests
                .is_allowed();
            if let Some(in_flight) = in_flight {
                // Cache access token
                let access_token = Arc::new(access_token);
                self.jmap.cache_access_token(access_token.clone());

                // Create session
                self.state = State::Authenticated {
                    data: Arc::new(SessionData::new(self, &access_token, in_flight).await?),
                };
                self.write_bytes(
                    StatusResponse::ok("Authentication successful")
                        .with_code(ResponseCode::Capability {
                            capabilities: Capability::all_capabilities(true, self.is_tls),
                        })
                        .with_tag(tag)
                        .into_bytes(),
                )
                .await?;
                Ok(())
            } else {
                self.write_bytes(
                    StatusResponse::bye("Too many concurrent IMAP connections.").into_bytes(),
                )
                .await?;
                tracing::debug!(parent: &self.span,
                    event = "disconnect",
                    "Too many concurrent connections, disconnecting.",
                );
                Err(())
            }
        } else {
            self.write_bytes(
                StatusResponse::no("Authentication failed")
                    .with_tag(tag)
                    .with_code(ResponseCode::AuthenticationFailed)
                    .into_bytes(),
            )
            .await?;

            let auth_failures = self.state.auth_failures();
            if auth_failures < self.imap.max_auth_failures {
                self.state = State::NotAuthenticated {
                    auth_failures: auth_failures + 1,
                };
                Ok(())
            } else {
                self.write_bytes(
                    StatusResponse::bye("Too many authentication failures").into_bytes(),
                )
                .await?;
                tracing::debug!(
                    parent: &self.span,
                    event = "disconnect",
                    "Too many authentication failures, disconnecting.",
                );
                Err(())
            }
        }
    }

    pub async fn handle_unauthenticate(&mut self, request: Request<Command>) -> crate::OpResult {
        self.state = State::NotAuthenticated { auth_failures: 0 };

        self.write_bytes(
            StatusResponse::completed(Command::Unauthenticate)
                .with_tag(request.tag)
                .into_bytes(),
        )
        .await
    }
}

pub fn decode_challenge_plain(challenge: &[u8]) -> Result<Credentials<String>, &'static str> {
    let mut username = Vec::new();
    let mut secret = Vec::new();
    let mut arg_num = 0;
    for &ch in challenge {
        if ch != 0 {
            if arg_num == 1 {
                username.push(ch);
            } else if arg_num == 2 {
                secret.push(ch);
            }
        } else {
            arg_num += 1;
        }
    }

    match (String::from_utf8(username), String::from_utf8(secret)) {
        (Ok(username), Ok(secret)) if !username.is_empty() && !secret.is_empty() => {
            Ok((username, secret).into())
        }
        _ => Err("Invalid AUTH=PLAIN challenge."),
    }
}

pub fn decode_challenge_oauth(challenge: &[u8]) -> Result<Credentials<String>, &'static str> {
    let mut saw_marker = true;
    for (pos, &ch) in challenge.iter().enumerate() {
        if saw_marker {
            if challenge
                .get(pos..)
                .map_or(false, |b| b.starts_with(b"auth=Bearer "))
            {
                let pos = pos + 12;
                return Ok(Credentials::OAuthBearer {
                    token: String::from_utf8(
                        challenge
                            .get(
                                pos..pos
                                    + challenge
                                        .get(pos..)
                                        .and_then(|c| c.iter().position(|&ch| ch == 0x01))
                                        .unwrap_or(challenge.len()),
                            )
                            .ok_or("Failed to find end of bearer token")?
                            .to_vec(),
                    )
                    .map_err(|_| "Bearer token is not a valid UTF-8 string.")?,
                });
            } else {
                saw_marker = false;
            }
        } else if ch == 0x01 {
            saw_marker = true;
        }
    }

    Err("Failed to find 'auth=Bearer' in challenge.")
}