Moved document validation into a sepparate func

This commit is contained in:
StNicolay 2023-07-20 23:28:26 +03:00
parent ff6dcf6dfe
commit 63acd976e7
Signed by: StNicolay
GPG Key ID: 9693D04DCD962B0D

View File

@ -2,15 +2,36 @@ use crate::prelude::*;
use futures::future::try_join; use futures::future::try_join;
use itertools::Itertools; use itertools::Itertools;
use rustc_hash::FxHashSet; use rustc_hash::FxHashSet;
use std::fmt::Write; use std::{fmt::Write, path::Path};
use teloxide::{net::Download, types::Document}; use teloxide::{
net::Download,
types::{Document, FileMeta},
};
use tokio::task::spawn_blocking; use tokio::task::spawn_blocking;
use trim_in_place::TrimInPlace; use trim_in_place::TrimInPlace;
#[inline] #[inline]
async fn download_file(bot: &Throttle<Bot>, document: &Document) -> crate::Result<Vec<u8>> { fn validate_document(document: Option<&Document>) -> Result<&Document, &'static str> {
let path = bot.get_file(document.file.id.as_str()).await?.path; let document = match document {
let mut data = Vec::with_capacity(document.file.size as usize); Some(document) => document,
None => return Err("You didn't send a file. Try again"),
};
let name = match document.file_name.as_deref() {
Some(name) => Path::new(name.trim_end()),
None => return Err("Couldn't get the name of the file. Try sending it again"),
};
match name.extension() {
Some(ext) if ext.eq_ignore_ascii_case("json") => Ok(document),
_ => Err("Invalid file name. You need to send a json file. Try again"),
}
}
#[inline]
async fn download_file(bot: &Throttle<Bot>, file: &FileMeta) -> crate::Result<Vec<u8>> {
let path = bot.get_file(file.id.as_str()).await?.path;
let mut data = Vec::with_capacity(file.size as usize);
bot.download_file_stream(&path) bot.download_file_stream(&path)
.try_for_each(|bytes| { .try_for_each(|bytes| {
data.extend(bytes); data.extend(bytes);
@ -131,31 +152,15 @@ pub async fn get_user(
return Ok(()); return Ok(());
} }
let document = match msg.document() { let file = match validate_document(msg.document()) {
Some(document) => document, Ok(document) => &document.file,
None => { Err(text) => {
let msg = bot let msg = bot.send_message(msg.chat.id, text).await?;
.send_message(msg.chat.id, "You didn't send a file. Try again")
.await?;
handler.previous = MessageIds::from(&msg); handler.previous = MessageIds::from(&msg);
return Ok(()); return Ok(());
} }
}; };
match document.file_name.as_deref() {
Some(name) if name.trim_end().ends_with(".json") => (),
_ => {
let msg = bot
.send_message(
msg.chat.id,
"Invalid file name. You need to send a json file. Try again",
)
.await?;
handler.previous = MessageIds::from(&msg);
return Ok(());
}
}
let existing_names = async { let existing_names = async {
Account::get_names(user_id, &db) Account::get_names(user_id, &db)
.await? .await?
@ -164,7 +169,7 @@ pub async fn get_user(
.map_err(Into::into) .map_err(Into::into)
}; };
let (data, existing_names) = try_join(download_file(&bot, document), existing_names).await?; let (data, existing_names) = try_join(download_file(&bot, file), existing_names).await?;
let user = match spawn_blocking(move || user_from_vec(data, existing_names)).await? { let user = match spawn_blocking(move || user_from_vec(data, existing_names)).await? {
Ok(Ok(user)) => user, Ok(Ok(user)) => user,