forked from bot/tmp2
113 lines
3.3 KiB
Rust
113 lines
3.3 KiB
Rust
/*
|
|
* mCaptcha - A proof of work based DoS protection system
|
|
* Copyright © 2021 Aravinth Manivannan <realravinth@batsense.net>
|
|
*
|
|
* 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.
|
|
*
|
|
* 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/>.
|
|
*/
|
|
use std::cell::RefCell;
|
|
use std::rc::Rc;
|
|
|
|
use redis::cluster::ClusterClient;
|
|
use redis::FromRedisValue;
|
|
//use redis::cluster::ClusterConnection;
|
|
use redis::Client;
|
|
//use redis::Connection;
|
|
use redis::{aio::Connection, cluster::ClusterConnection};
|
|
|
|
use crate::errors::*;
|
|
|
|
#[derive(Clone)]
|
|
pub enum RedisConfig {
|
|
Single(String),
|
|
Cluster(Vec<String>),
|
|
}
|
|
|
|
impl RedisConfig {
|
|
pub fn connect(&self) -> RedisClient {
|
|
match self {
|
|
Self::Single(url) => {
|
|
let client = Client::open(url.as_str()).unwrap();
|
|
RedisClient::Single(client)
|
|
}
|
|
Self::Cluster(nodes) => {
|
|
let cluster_client = ClusterClient::open(nodes.to_owned()).unwrap();
|
|
RedisClient::Cluster(cluster_client)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub enum RedisConnection {
|
|
Single(Rc<RefCell<Connection>>),
|
|
Cluster(Rc<RefCell<ClusterConnection>>),
|
|
}
|
|
|
|
impl RedisConnection {
|
|
#[inline]
|
|
pub fn get_client(&self) -> Self {
|
|
match self {
|
|
Self::Single(con) => Self::Single(Rc::clone(&con)),
|
|
Self::Cluster(con) => Self::Cluster(Rc::clone(&con)),
|
|
}
|
|
}
|
|
#[inline]
|
|
pub async fn exec<T: FromRedisValue>(&self, cmd: &mut redis::Cmd) -> redis::RedisResult<T> {
|
|
match self {
|
|
RedisConnection::Single(con) => cmd.query_async(&mut *con.borrow_mut()).await,
|
|
RedisConnection::Cluster(con) => cmd.query(&mut *con.borrow_mut()),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[derive(Clone)]
|
|
pub enum RedisClient {
|
|
Single(Client),
|
|
Cluster(ClusterClient),
|
|
}
|
|
|
|
pub struct Redis {
|
|
_client: RedisClient,
|
|
connection: RedisConnection,
|
|
}
|
|
|
|
impl Redis {
|
|
pub async fn new(redis: RedisConfig) -> CaptchaResult<Self> {
|
|
let (_client, connection) = Self::connect(redis).await;
|
|
let master = Self {
|
|
_client,
|
|
connection,
|
|
};
|
|
Ok(master)
|
|
}
|
|
|
|
pub fn get_client(&self) -> RedisConnection {
|
|
self.connection.get_client()
|
|
}
|
|
|
|
async fn connect(redis: RedisConfig) -> (RedisClient, RedisConnection) {
|
|
let redis = redis.connect();
|
|
let client = match &redis {
|
|
RedisClient::Single(c) => {
|
|
let con = c.get_async_connection().await.unwrap();
|
|
RedisConnection::Single(Rc::new(RefCell::new(con)))
|
|
}
|
|
RedisClient::Cluster(c) => {
|
|
let con = c.get_connection().unwrap();
|
|
RedisConnection::Cluster(Rc::new(RefCell::new(con)))
|
|
}
|
|
};
|
|
(redis, client)
|
|
}
|
|
}
|