flibusta-apps/telegram_files_server
1
0
Fork 0
mirror of https://github.com/flibusta-apps/telegram_files_server.git synced 2025-12-06 20:45:37 +01:00
Code Issues Actions Packages Projects Releases Wiki Activity

Rewrite to rust

Browse Source 
Co-authored-by: Kurbanov Bulat <kurbanovbul@gmail.com>
This commit is contained in:
Шатунов Антон
2024-05-06 23:04:24 +03:00
parent c52950e937
commit 26de50da3a
28 changed files with 3044 additions and 1681 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
.github/workflows/codeql-analysis.yml vendored
View File

@@ -1,35 +0,0 @@
name: "CodeQL"
on:
push:
branches: [ main ]
pull_request:
branches: [ main ]
schedule:
- cron: '0 12 * * *'
jobs:
analyze:
name: Analyze
runs-on: ubuntu-latest
permissions:
actions: read
contents: read
security-events: write
strategy:
fail-fast: false
matrix:
language: [ 'python' ]
steps:
- name: Checkout repository
uses: actions/checkout@v4
- name: Initialize CodeQL
uses: github/codeql-action/init@v3
with:
languages: ${{ matrix.language }}
- name: Perform CodeQL Analysis
uses: github/codeql-action/analyze@v3

35
.github/workflows/linters.yaml vendored
View File

@@ -1,35 +0,0 @@
name: Linters
on:
push:
branches:
- main
pull_request:
types: [opened, synchronize, reopened]
jobs:
Run-Pre-Commit:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
with:
fetch-depth: 32
- uses: actions/setup-python@v5
with:
python-version: 3.11
- name: Install pre-commit
run: pip3 install pre-commit
- name: Pre-commit (Push)
env:
SETUPTOOLS_USE_DISTUTILS: stdlib
if: ${{ github.event_name == 'push' }}
run: pre-commit run --source ${{ github.event.before }} --origin ${{ github.event.after }} --show-diff-on-failure
- name: Pre-commit (Pull-Request)
env:
SETUPTOOLS_USE_DISTUTILS: stdlib
if: ${{ github.event_name == 'pull_request' }}
run: pre-commit run --source ${{ github.event.pull_request.base.sha }} --origin ${{ github.event.pull_request.head.sha }} --show-diff-on-failure

5
.gitignore vendored
View File

@@ -9,3 +9,8 @@ __pycache__
*.session-journal *.session-journal
venv venv
# Added by cargo
/target

21
.pre-commit-config.yaml
View File

@@ -1,18 +1,7 @@
exclude: 'docs|node_modules|migrations|.git|.tox'
repos: repos:
- repo: https://github.com/ambv/black - repo: https://github.com/doublify/pre-commit-rust
rev: 23.3.0 rev: v1.0
hooks: hooks:
- id: black - id: fmt
language_version: python3.11 - id: cargo-check
- id: clippy
- repo: https://github.com/charliermarsh/ruff-pre-commit
rev: 'v0.0.267'
hooks:
- id: ruff
- repo: https://github.com/crate-ci/typos
rev: typos-dict-v0.9.26
hooks:
- id: typos

2668
Cargo.lock generated Normal file
View File

File diff suppressed because it is too large Load Diff

34
Cargo.toml Normal file
View File

@@ -0,0 +1,34 @@
[package]
name = "telegram_files_server"
version = "0.1.0"
edition = "2021"
[dependencies]
serde = "1.0.200"
serde_json = "1.0.116"
axum = { version = "0.7.5", features = ["multipart"] }
axum_typed_multipart = "0.11.1"
tracing = "0.1.40"
tracing-subscriber = { version = "0.3.18", features = ["env-filter"]}
tower-http = { version = "0.5.2", features = ["trace"] }
tokio = "1.37.0"
tokio-util = { version = "0.7.11", features = [ "full" ] }
axum-prometheus = "0.6.1"
futures = "0.3.30"
once_cell = "1.19.0"
teloxide = "0.12.2"
sentry = "0.32.3"
dotenv = "0.15.0"
reqwest = { version = "0.11.10", features = [
"json",
"stream",
"multipart",
], default-features = false }

40
docker/production.dockerfile
View File

@@ -1,32 +1,24 @@
FROM ghcr.io/flibusta-apps/base_docker_images:3.12-poetry-buildtime AS build-image FROM rust:bullseye AS builder
WORKDIR /root/poetry
COPY pyproject.toml poetry.lock /root/poetry/
ENV VENV_PATH=/opt/venv
RUN poetry export --without-hashes > requirements.txt \
&& . /opt/venv/bin/activate \
&& pip install -r requirements.txt --no-cache-dir
FROM ghcr.io/flibusta-apps/base_docker_images:3.12-postgres-runtime AS runtime-image
RUN apt-get update \
&& apt-get install -y curl jq \
&& rm -rf /var/lib/apt/lists/*
WORKDIR /app WORKDIR /app
ENV VENV_PATH=/opt/venv COPY . .
ENV PATH="$VENV_PATH/bin:$PATH"
COPY --from=build-image $VENV_PATH $VENV_PATH RUN cargo build --release --bin telegram_files_server
COPY ./fastapi_file_server/ /app/
COPY ./scripts/* /
FROM debian:bullseye-slim
RUN apt-get update \
&& apt-get install -y openssl ca-certificates curl jq \
&& rm -rf /var/lib/apt/lists/*
RUN update-ca-certificates
COPY ./scripts/*.sh /
RUN chmod +x /*.sh RUN chmod +x /*.sh
EXPOSE 8080 WORKDIR /app
CMD ["/start.sh"] COPY --from=builder /app/target/release/telegram_files_server /usr/local/bin
ENTRYPOINT ["/start.sh"]

11
fastapi_file_server/app/depends.py
View File

@@ -1,11 +0,0 @@
from fastapi import HTTPException, Security, status
from core.auth import default_security
from core.config import env_config
async def check_token(api_key: str = Security(default_security)):
if api_key != env_config.API_KEY:
raise HTTPException(
status_code=status.HTTP_403_FORBIDDEN, detail="Wrong api key!"
)

5
fastapi_file_server/app/on_start.py
View File

@@ -1,5 +0,0 @@
from app.services.storages import StoragesContainer
async def on_start():
await StoragesContainer.prepare()

19
fastapi_file_server/app/serializers.py
View File

@@ -1,19 +0,0 @@
import enum
from pydantic import BaseModel
from typing_extensions import TypedDict
class UploadBackend(enum.StrEnum):
bot = "bot"
user = "user"
class Data(TypedDict):
chat_id: int
message_id: int
class UploadedFile(BaseModel):
backend: UploadBackend
data: Data

58
fastapi_file_server/app/services/file_downloader.py
View File

@@ -1,58 +0,0 @@
from app.serializers import UploadBackend
from app.services.storages import BotStorage, StoragesContainer, UserStorage
class FileDownloader:
_bot_storage_index = 0
_user_storage_index = 0
@classmethod
@property
def bot_storages(cls) -> list[BotStorage]:
return StoragesContainer.BOT_STORAGES
@classmethod
@property
def user_storages(cls) -> list[UserStorage]:
return StoragesContainer.USER_STORAGES
@classmethod
def get_bot_storage(cls) -> BotStorage:
if not cls.bot_storages:
raise ValueError("Aiogram storage not exist!")
bot_storages: list[BotStorage] = cls.bot_storages # type: ignore
cls._bot_storage_index = (cls._bot_storage_index + 1) % len(bot_storages)
return bot_storages[cls._bot_storage_index]
@classmethod
def get_user_storage(cls) -> UserStorage:
if not cls.user_storages:
raise ValueError("Telethon storage not exists!")
user_storages: list[UserStorage] = cls.user_storages # type: ignore
cls._user_storage_index = (cls._user_storage_index + 1) % len(user_storages)
return user_storages[cls._user_storage_index]
@classmethod
async def _download_via(cls, message_id: int, storage_type: UploadBackend):
if storage_type == UploadBackend.bot:
storage = cls.get_bot_storage()
else:
storage = cls.get_user_storage()
return await storage.download(message_id)
@classmethod
async def download_by_message_id(cls, message_id: int):
if not cls.bot_storages and not cls.user_storages:
raise ValueError("Files storage not exist!")
if (data := await cls._download_via(message_id, UploadBackend.bot)) is not None:
return data
return await cls._download_via(message_id, UploadBackend.user)

117
fastapi_file_server/app/services/file_uploader.py
View File

@@ -1,117 +0,0 @@
from typing import Optional
from fastapi import UploadFile
from app.serializers import Data, UploadBackend, UploadedFile
from app.services.storages import BotStorage, StoragesContainer, UserStorage
def seekable(*args, **kwargs):
return True
class FileUploader:
_bot_storage_index = 0
_user_storage_index = 0
@classmethod
@property
def bot_storages(cls) -> list[BotStorage]:
return StoragesContainer.BOT_STORAGES
@classmethod
@property
def user_storages(cls) -> list[UserStorage]:
return StoragesContainer.USER_STORAGES
def __init__(
self,
file: UploadFile,
filename: str,
file_size: int,
caption: Optional[str] = None,
) -> None:
self.file = file
self.filename = filename
self.file_size = file_size
self.caption = caption
self.upload_data: Optional[Data] = None
self.upload_backend: Optional[UploadBackend] = None
async def _upload(self) -> bool:
if not self.bot_storages and not self.user_storages:
raise ValueError("Files storage not exist!")
if await self._upload_via(UploadBackend.bot):
return True
return await self._upload_via(UploadBackend.user)
async def _upload_via(self, storage_type: UploadBackend) -> bool:
if storage_type == UploadBackend.bot:
storage = self.get_bot_storage()
else:
storage = self.get_user_storage()
file = self.file
setattr(file, "seekable", seekable) # noqa: B010
data = await storage.upload(
file, # type: ignore
file_size=self.file_size,
filename=self.filename,
caption=self.caption,
)
if not data:
return False
self.upload_data = {"chat_id": data[0], "message_id": data[1]}
self.upload_backend = storage_type
return True
def get_result(self) -> UploadedFile:
assert self.upload_backend is not None
assert self.upload_data is not None
return UploadedFile(backend=self.upload_backend, data=self.upload_data)
@classmethod
def get_bot_storage(cls) -> BotStorage:
if not cls.bot_storages:
raise ValueError("Aiogram storage not exist!")
bot_storages: list[BotStorage] = cls.bot_storages # type: ignore
cls._bot_storage_index = (cls._bot_storage_index + 1) % len(bot_storages)
return bot_storages[cls._bot_storage_index]
@classmethod
def get_user_storage(cls) -> UserStorage:
if not cls.user_storages:
raise ValueError("Telethon storage not exists!")
user_storages: list[UserStorage] = cls.user_storages # type: ignore
cls._user_storage_index = (cls._user_storage_index + 1) % len(user_storages)
return user_storages[cls._user_storage_index]
@classmethod
async def upload(
cls,
file: UploadFile,
filename: str,
file_size: int,
caption: Optional[str] = None,
) -> Optional[UploadedFile]:
uploader = cls(file, filename, file_size, caption)
upload_result = await uploader._upload()
if not upload_result:
return None
return uploader.get_result()

146
fastapi_file_server/app/services/storages.py
View File

@@ -1,146 +0,0 @@
from typing import AsyncIterator, Optional
import telethon.client
import telethon.errors
import telethon.hints
import telethon.tl.types
from core.config import env_config
class BaseStorage:
def __init__(self, channel_id: int, app_id: int, api_hash: str, session: str):
self.channel_id = channel_id
self.client = telethon.client.TelegramClient(session, app_id, api_hash)
self.ready = False
async def prepare(self):
...
async def upload(
self,
file: telethon.hints.FileLike,
filename: str,
file_size: int,
caption: Optional[str] = None,
) -> Optional[tuple[int, int]]:
try:
uploaded_file = await self.client.upload_file(
file, file_size=file_size, file_name=filename
)
if caption:
message = await self.client.send_file(
entity=self.channel_id,
file=uploaded_file,
caption=caption,
force_document=True,
)
else:
message = await self.client.send_file(
entity=self.channel_id, file=uploaded_file, force_document=True
)
except telethon.errors.FilePartInvalidError:
return None
except telethon.errors.FilePartsInvalidError:
return None
except telethon.errors.PhotoInvalidError:
return None
if not message.media:
return None
return self.channel_id, message.id
async def download(self, message_id: int) -> Optional[AsyncIterator[bytes]]:
messages = await self.client.get_messages(self.channel_id, ids=[message_id])
if not messages:
return None
message: Optional[telethon.tl.types.Message] = messages[0]
if message is None:
return None
if message.media is None:
return None
return self.client.iter_download(message.media)
class UserStorage(BaseStorage):
async def prepare(self):
if self.ready:
return
await self.client.start() # type: ignore
if not await self.client.is_user_authorized():
await self.client.sign_in()
try:
await self.client.sign_in(code=input("Enter code: "))
except telethon.errors.SessionPasswordNeededError:
await self.client.sign_in(password=input("Enter password: "))
self.ready = True
class BotStorage(BaseStorage):
def __init__(
self,
channel_id: int,
app_id: int,
api_hash: str,
session: str,
token: str,
) -> None:
super().__init__(channel_id, app_id, api_hash, session)
self.token = token
async def prepare(self):
if self.ready:
return
await self.client.start(bot_token=self.token) # type: ignore
self.ready = True
class StoragesContainer:
BOT_STORAGES: list[BotStorage] = []
USER_STORAGES: list[UserStorage] = []
@classmethod
async def prepare(cls):
if not env_config.TELETHON_APP_CONFIG:
return
if env_config.BOT_TOKENS:
cls.BOT_STORAGES: list[BotStorage] = [
BotStorage(
env_config.TELEGRAM_CHAT_ID,
env_config.TELETHON_APP_CONFIG.APP_ID,
env_config.TELETHON_APP_CONFIG.API_HASH,
token.split(":")[0],
token,
)
for token in env_config.BOT_TOKENS
]
if env_config.TELETHON_SESSIONS:
cls.USER_STORAGES: list[UserStorage] = [
UserStorage(
env_config.TELEGRAM_CHAT_ID,
env_config.TELETHON_APP_CONFIG.APP_ID,
env_config.TELETHON_APP_CONFIG.API_HASH,
session,
)
for session in env_config.TELETHON_SESSIONS
]
for storage in [*cls.BOT_STORAGES, *cls.USER_STORAGES]:
await storage.prepare()

34
fastapi_file_server/app/views.py
View File

@@ -1,34 +0,0 @@
from typing import Annotated, Optional
from fastapi import APIRouter, Depends, Form, HTTPException, UploadFile, status
from fastapi.responses import StreamingResponse
from app.depends import check_token
from app.serializers import UploadedFile
from app.services.file_downloader import FileDownloader
from app.services.file_uploader import FileUploader
router = APIRouter(
prefix="/api/v1/files", dependencies=[Depends(check_token)], tags=["files"]
)
@router.post("/upload/", response_model=UploadedFile)
async def upload_file(
file: UploadFile,
file_size: Annotated[int, Form()],
filename: Annotated[str, Form()],
caption: Annotated[Optional[str], Form()],
):
return await FileUploader.upload(file, filename, file_size, caption=caption)
@router.get("/download_by_message/{chat_id}/{message_id}")
async def download_by_message(chat_id: str, message_id: int):
data = await FileDownloader.download_by_message_id(message_id)
if data is None:
raise HTTPException(status.HTTP_400_BAD_REQUEST)
return StreamingResponse(data)

31
fastapi_file_server/core/app.py
View File

@@ -1,31 +0,0 @@
from fastapi import FastAPI
from fastapi.responses import ORJSONResponse
from prometheus_fastapi_instrumentator import Instrumentator
import sentry_sdk
from app.on_start import on_start
from app.views import router
from core.config import env_config
sentry_sdk.init(
env_config.SENTRY_DSN,
)
def start_app() -> FastAPI:
app = FastAPI(default_response_class=ORJSONResponse)
app.include_router(router)
@app.on_event("startup")
async def startup() -> None:
await on_start()
Instrumentator(
should_ignore_untemplated=True,
excluded_handlers=["/docs", "/metrics", "/healthcheck"],
).instrument(app).expose(app, include_in_schema=True)
return app

4
fastapi_file_server/core/auth.py
View File

@@ -1,4 +0,0 @@
from fastapi.security import APIKeyHeader
default_security = APIKeyHeader(name="Authorization")

34
fastapi_file_server/core/config.py
View File

@@ -1,34 +0,0 @@
from typing import Optional
from pydantic import BaseModel
from pydantic_settings import (
BaseSettings,
SettingsConfigDict,
)
BotToken = str
TelethonSessionName = str
class TelethonConfig(BaseModel):
APP_ID: int
API_HASH: str
class EnvConfig(BaseSettings):
model_config = SettingsConfigDict(env_file='.env', env_file_encoding='utf-8')
API_KEY: str
TELEGRAM_CHAT_ID: int
BOT_TOKENS: Optional[list[BotToken]] = None
TELETHON_APP_CONFIG: Optional[TelethonConfig] = None
TELETHON_SESSIONS: Optional[list[TelethonSessionName]] = None
SENTRY_DSN: str
env_config = EnvConfig()

4
fastapi_file_server/main.py
View File

@@ -1,4 +0,0 @@
from core.app import start_app
app = start_app()

1025
poetry.lock generated
View File

File diff suppressed because it is too large Load Diff

77
pyproject.toml
View File

@@ -1,77 +0,0 @@
[tool.poetry]
name = "fastapi_file_server"
version = "0.1.0"
description = ""
authors = ["Kurbanov Bulat <kurbanovbul@gmail.com>"]
license = "MIT"
[tool.poetry.dependencies]
python = "^3.12"
fastapi = "^0.110.2"
pydantic = "^2.6.4"
python-multipart = "^0.0.9"
telethon = "^1.35.0"
prometheus-fastapi-instrumentator = "^7.0.0"
orjson = "^3.10.1"
sentry-sdk = "^2.0.0"
pydantic-settings = "^2.2.1"
typing-extensions = "^4.11.0"
pyyaml = "^6.0.1"
certifi = "^2024.2.2"
granian = "^1.2.3"
[tool.poetry.group.dev.dependencies]
pytest = "^7.2.0"
mypy = "^0.991"
pre-commit = "^2.21.0"
[build-system]
requires = ["poetry-core>=1.0.0"]
build-backend = "poetry.core.masonry.api"
[tool.black]
include = '\.pyi?$'
exclude = '''
/(
\.git
| \.vscode
| \venv
| alembic
)/
'''
[tool.ruff]
fix = true
target-version = "py311"
src = ["fastapi_file_server"]
line-length=88
ignore = []
select = ["B", "C", "E", "F", "W", "B9", "I001"]
exclude = [
# No need to traverse our git directory
".git",
# There's no value in checking cache directories
"__pycache__",
# The conf file is mostly autogenerated, ignore it
"fastapi_file_server/app/alembic",
]
[tool.ruff.flake8-bugbear]
extend-immutable-calls = ["fastapi.File", "fastapi.Form", "fastapi.Security"]
[tool.ruff.mccabe]
max-complexity = 15
[tool.ruff.isort]
known-first-party = ["core", "app"]
force-sort-within-sections = true
force-wrap-aliases = true
section-order = ["future", "standard-library", "base_framework", "framework_ext", "third-party", "first-party", "local-folder"]
lines-after-imports = 2
[tool.ruff.isort.sections]
base_framework = ["fastapi",]
framework_ext = ["starlette"]
[tool.ruff.pyupgrade]
keep-runtime-typing = true

5
scripts/start.sh
View File

@@ -1,8 +1,5 @@
#! /usr/bin/env sh #! /usr/bin/env sh
cd /app
mkdir -p prometheus
/env.sh > ./.env /env.sh > ./.env
granian --interface asgi --host 0.0.0.0 --port 8080 --loop uvloop main:app exec /usr/local/bin/telegram_files_server

35
src/config.rs Normal file
View File

@@ -0,0 +1,35 @@
use once_cell::sync::Lazy;
fn get_env(env: &'static str) -> String {
std::env::var(env).unwrap_or_else(|_| panic!("Cannot get the {} env variable", env))
}
pub struct Config {
pub api_key: String,
pub telegram_api_url: String,
pub telegram_chat_id: i64,
pub telegram_temp_chat_id: i64,
pub bot_tokens: Vec<String>,
pub sentry_dsn: String,
}
impl Config {
pub fn load() -> Config {
Config {
api_key: get_env("API_KEY"),
telegram_api_url: get_env("API_URL"),
telegram_chat_id: get_env("TELEGRAM_CHAT_ID").parse().unwrap(),
telegram_temp_chat_id: get_env("TELEGRAM_TEMP_CHAT_ID").parse().unwrap(),
bot_tokens: serde_json::from_str(&get_env("BOT_TOKENS")).unwrap(),
sentry_dsn: get_env("SENTRY_DSN"),
}
}
}
pub static CONFIG: Lazy<Config> = Lazy::new(Config::load);

32
src/core/bot.rs Normal file
View File

@@ -0,0 +1,32 @@
use std::sync::{
atomic::{AtomicUsize, Ordering},
Arc,
};
use once_cell::sync::Lazy;
use teloxide::Bot;
use crate::config::{self, CONFIG};
pub struct RoundRobinBot {
bot_tokens: Arc<Vec<String>>,
current_index: AtomicUsize,
}
impl RoundRobinBot {
pub fn new(bot_tokens: Vec<String>) -> Self {
RoundRobinBot {
bot_tokens: Arc::new(bot_tokens),
current_index: AtomicUsize::new(0),
}
}
pub fn get_bot(&self) -> Bot {
let index = self.current_index.fetch_add(1, Ordering::Relaxed) % self.bot_tokens.len();
Bot::new(self.bot_tokens[index].clone())
.set_api_url(reqwest::Url::parse(CONFIG.telegram_api_url.as_str()).unwrap())
}
}
pub static ROUND_ROBIN_BOT: Lazy<RoundRobinBot> =
Lazy::new(|| RoundRobinBot::new(config::CONFIG.bot_tokens.clone()));

104
src/core/file_utils.rs Normal file
View File

@@ -0,0 +1,104 @@
use std::pin::Pin;
use axum::body::Bytes;
use futures::TryStreamExt;
use serde::Serialize;
use teloxide::{
net::Download,
requests::Requester,
types::{ChatId, InputFile, MessageId},
Bot,
};
use tokio::io::AsyncRead;
use tokio_util::compat::FuturesAsyncReadCompatExt;
use crate::config::CONFIG;
use super::bot::ROUND_ROBIN_BOT;
#[derive(Serialize)]
pub struct UploadedFile {
pub backend: String,
pub data: MessageInfo,
}
#[derive(Serialize)]
pub struct MessageInfo {
pub chat_id: i64,
pub message_id: i32,
}
pub async fn upload_file(
file: Bytes,
filename: String,
caption: Option<String>,
) -> Result<UploadedFile, String> {
let bot = ROUND_ROBIN_BOT.get_bot();
let document = InputFile::memory(file).file_name(filename);
let mut request = bot.send_document(ChatId(CONFIG.telegram_chat_id), document);
request.caption = caption;
let result = request.await;
match result {
Ok(message) => Ok(UploadedFile {
backend: "bot".to_string(),
data: MessageInfo {
chat_id: message.chat.id.0,
message_id: message.id.0,
},
}),
Err(err) => Err(err.to_string()),
}
}
pub async fn download_file(chat_id: i64, message_id: i32) -> Option<BotDownloader> {
let bot = ROUND_ROBIN_BOT.get_bot();
let result = bot
.forward_message(
ChatId(CONFIG.telegram_temp_chat_id),
ChatId(chat_id),
MessageId(message_id),
)
.await;
match result {
Ok(message) => {
if message.document() == None {
return Option::None;
}
let file_id = message.document().unwrap().file.id.clone();
let path = bot.get_file(file_id.clone()).await.unwrap().path;
return Some(BotDownloader::new(bot, path));
}
Err(_) => None,
}
}
pub struct BotDownloader {
bot: Bot,
file_path: String,
}
impl BotDownloader {
pub fn new(bot: Bot, file_path: String) -> Self {
Self { bot, file_path }
}
pub fn get_async_read(self) -> Pin<Box<dyn AsyncRead + Send>> {
let stream = self.bot.download_file_stream(&self.file_path);
Box::pin(
stream
.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e))
.into_async_read()
.compat()
)
}
}

3
src/core/mod.rs Normal file
View File

@@ -0,0 +1,3 @@
pub mod bot;
pub mod views;
pub mod file_utils;

103
src/core/views.rs Normal file
View File

@@ -0,0 +1,103 @@
use axum::{
body::Bytes,
extract::{DefaultBodyLimit, Path},
http::{self, Request, StatusCode},
middleware::{self, Next},
response::{IntoResponse, Response},
routing::{get, post},
Router,
};
use axum_prometheus::PrometheusMetricLayer;
use axum_typed_multipart::{TryFromMultipart, TypedMultipart};
use tokio_util::io::ReaderStream;
use tower_http::trace::{self, TraceLayer};
use tracing::Level;
use crate::config::CONFIG;
use super::file_utils::{download_file, upload_file};
const BODY_LIMIT: usize = 4 * (2 << 30); // bytes: 4GB
async fn auth(req: Request<axum::body::Body>, next: Next) -> Result<Response, StatusCode> {
let auth_header = req
.headers()
.get(http::header::AUTHORIZATION)
.and_then(|header| header.to_str().ok());
let auth_header = if let Some(auth_header) = auth_header {
auth_header
} else {
return Err(StatusCode::UNAUTHORIZED);
};
if auth_header != CONFIG.api_key {
return Err(StatusCode::UNAUTHORIZED);
}
Ok(next.run(req).await)
}
pub async fn get_router() -> Router {
let (prometheus_layer, metric_handle) = PrometheusMetricLayer::pair();
let app_router = Router::new()
.route("/upload", post(upload))
.route("/download_by_message/:chat_id/:message_id", get(download))
.layer(DefaultBodyLimit::max(BODY_LIMIT))
.layer(middleware::from_fn(auth))
.layer(prometheus_layer);
let metric_router =
Router::new().route("/metrics", get(|| async move { metric_handle.render() }));
Router::new()
.nest("/api/v1/files", app_router)
.nest("/", metric_router)
.layer(
TraceLayer::new_for_http()
.make_span_with(trace::DefaultMakeSpan::new().level(Level::INFO))
.on_response(trace::DefaultOnResponse::new().level(Level::INFO)),
)
}
#[derive(TryFromMultipart)]
pub struct UploadFileRequest {
#[form_data(limit = "unlimited")]
file: Bytes,
filename: String,
caption: Option<String>,
}
async fn upload(data: TypedMultipart<UploadFileRequest>) -> impl IntoResponse {
let result = match upload_file(
data.file.clone(),
data.filename.to_string(),
data.caption.clone(),
)
.await
{
Ok(file) => serde_json::to_string(&file),
Err(err) => Ok(err),
};
result.unwrap()
}
async fn download(Path(chat_id): Path<i64>, Path(message_id): Path<i32>) -> impl IntoResponse {
let downloader = download_file(chat_id, message_id).await;
let data = match downloader {
Some(v) => v.get_async_read(),
None => return StatusCode::NOT_FOUND.into_response()
};
let reader = ReaderStream::new(data);
axum::body::Body::from_stream(reader).into_response()
}

36
src/main.rs Normal file
View File

@@ -0,0 +1,36 @@
mod config;
mod core;
use std::{net::SocketAddr, str::FromStr};
use sentry::{integrations::debug_images::DebugImagesIntegration, types::Dsn, ClientOptions};
use crate::core::views::get_router;
#[tokio::main]
async fn main() {
dotenv::dotenv().ok();
let options = ClientOptions {
dsn: Some(Dsn::from_str(&config::CONFIG.sentry_dsn).unwrap()),
default_integrations: false,
..Default::default()
}
.add_integration(DebugImagesIntegration::new());
let _guard = sentry::init(options);
tracing_subscriber::fmt()
.with_target(false)
.compact()
.init();
let addr = SocketAddr::from(([0, 0, 0, 0], 8080));
let app = get_router().await;
println!("Start webserver...");
let listener = tokio::net::TcpListener::bind(&addr).await.unwrap();
axum::serve(listener, app).await.unwrap();
println!("Webserver shutdown...");
}