mirror of
https://github.com/andwati/telecom-platform.git
synced 2026-10-18 11:02:54 +00:00
feat: Add charging engine with Redis-based credit management and usage tracking
- Add ChargingEngine struct with Redis client, rating plans, and sync interval - Add SubscriberAccount model with balance, data/voice/SMS limits and usage tracking - Add UsageEvent model for billing events with cost calculation - Add RatingPlan model with data/voice/SMS rates and limits - Add AccountStatus enum (Active, Suspended, Terminated, Blocked) - Add UsageType enum (Data, Voice, SMS) - Add sync_usage_data for
This commit is contained in:
nutcas3
@@ -0,0 +1,127 @@
|
||||
use axum::{
|
||||
extract::{Path, State},
|
||||
response::IntoResponse,
|
||||
routing::{get, post},
|
||||
Json, Router,
|
||||
};
|
||||
use tower_http::cors::{Any, CorsLayer};
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::errors::{ChargingError, ChargingResult, validate_ip, validate_bytes, ErrorContext};
|
||||
use crate::models::*;
|
||||
|
||||
pub struct AppState {
|
||||
pub charging_engine: std::sync::Arc<crate::charging::ChargingEngine>,
|
||||
}
|
||||
|
||||
pub fn create_router(state: AppState) -> Router {
|
||||
let cors = CorsLayer::new()
|
||||
.allow_origin(Any)
|
||||
.allow_methods(Any)
|
||||
.allow_headers(Any);
|
||||
|
||||
Router::new()
|
||||
.route("/v1/credit/:ip/check", post(check_credit))
|
||||
.route("/v1/credit/:ip/deduct", post(deduct_credit))
|
||||
.route("/v1/credit/:ip/add", post(add_credit))
|
||||
.route("/v1/credit/:ip/balance", get(get_balance))
|
||||
.route("/health", get(health_check))
|
||||
.layer(cors)
|
||||
.with_state(state)
|
||||
}
|
||||
|
||||
/// POST /v1/credit/:ip/check
|
||||
/// Check if user has enough credit for data usage
|
||||
pub async fn check_credit(
|
||||
Path(ip): Path<String>,
|
||||
State(state): State<AppState>,
|
||||
Json(req): Json<CreditCheckRequest>,
|
||||
) -> ChargingResult<Json<CreditCheckResponse>> {
|
||||
// Validate input
|
||||
validate_ip(&ip)?;
|
||||
validate_bytes(req.bytes_requested)?;
|
||||
|
||||
let allowed = state.charging_engine.check_credit(&ip, req.bytes_requested).await
|
||||
.with_context("Failed to check credit")?;
|
||||
|
||||
let remaining = state.charging_engine.get_balance(&ip).await
|
||||
.with_context("Failed to get balance")?;
|
||||
|
||||
info!(
|
||||
"Credit check for {}: {} bytes requested, {} bytes available, allowed: {}",
|
||||
ip, req.bytes_requested, remaining, allowed
|
||||
);
|
||||
|
||||
Ok(Json(CreditCheckResponse {
|
||||
allowed,
|
||||
remaining_bytes: remaining,
|
||||
}))
|
||||
}
|
||||
|
||||
/// POST /v1/credit/:ip/deduct
|
||||
/// Deduct bytes from user's credit balance
|
||||
pub async fn deduct_credit(
|
||||
Path(ip): Path<String>,
|
||||
State(state): State<AppState>,
|
||||
Json(req): Json<DeductRequest>,
|
||||
) -> ChargingResult<()> {
|
||||
// Validate input
|
||||
validate_ip(&ip)?;
|
||||
validate_bytes(req.bytes_used)?;
|
||||
|
||||
let new_balance = state.charging_engine.deduct_credit(&ip, req.bytes_used).await
|
||||
.with_context("Failed to deduct credit")?;
|
||||
|
||||
info!(
|
||||
"User {} deducted {} bytes, remaining: {}",
|
||||
ip, req.bytes_used, new_balance
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// POST /v1/credit/:ip/add
|
||||
/// Add bytes to user's credit balance
|
||||
pub async fn add_credit(
|
||||
Path(ip): Path<String>,
|
||||
State(state): State<AppState>,
|
||||
Json(req): Json<AddCreditRequest>,
|
||||
) -> ChargingResult<()> {
|
||||
// Validate input
|
||||
validate_ip(&ip)?;
|
||||
validate_bytes(req.bytes_to_add)?;
|
||||
|
||||
let new_balance = state.charging_engine.add_credit(&ip, req.bytes_to_add).await
|
||||
.with_context("Failed to add credit")?;
|
||||
|
||||
info!(
|
||||
"User {} added {} bytes, new balance: {}",
|
||||
ip, req.bytes_to_add, new_balance
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// GET /v1/credit/:ip/balance
|
||||
/// Get current credit balance
|
||||
pub async fn get_balance(
|
||||
Path(ip): Path<String>,
|
||||
State(state): State<AppState>,
|
||||
) -> ChargingResult<Json<BalanceResponse>> {
|
||||
// Validate input
|
||||
validate_ip(&ip)?;
|
||||
|
||||
let balance = state.charging_engine.get_balance(&ip).await
|
||||
.with_context("Failed to get balance")?;
|
||||
|
||||
Ok(Json(BalanceResponse {
|
||||
ip,
|
||||
balance_bytes: balance,
|
||||
}))
|
||||
}
|
||||
|
||||
/// GET /health
|
||||
/// Health check endpoint
|
||||
pub async fn health_check() -> ChargingResult<Json<HealthResponse>> {
|
||||
Ok(Json(HealthResponse::default()))
|
||||
}
|
||||
@@ -0,0 +1,478 @@
|
||||
use chrono::{DateTime, Utc};
|
||||
use redis::AsyncCommands;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::collections::HashMap;
|
||||
use std::time::Duration;
|
||||
use tokio::time::{interval, sleep};
|
||||
use tracing::{info, warn, error, debug};
|
||||
|
||||
use crate::errors::{ChargingError, ChargingResult, ErrorContext};
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct SubscriberAccount {
|
||||
pub imsi: String,
|
||||
pub balance: i64, // in smallest currency unit (e.g., cents)
|
||||
pub data_limit: u64, // bytes
|
||||
pub data_used: u64, // bytes
|
||||
pub voice_limit: u64, // seconds
|
||||
pub voice_used: u64, // seconds
|
||||
pub sms_limit: u64, // count
|
||||
pub sms_used: u64, // count
|
||||
pub status: AccountStatus,
|
||||
pub last_updated: DateTime<Utc>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum AccountStatus {
|
||||
Active,
|
||||
Suspended,
|
||||
Terminated,
|
||||
Blocked,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct UsageEvent {
|
||||
pub imsi: String,
|
||||
pub session_id: String,
|
||||
pub usage_type: UsageType,
|
||||
pub volume: u64, // bytes, seconds, or count
|
||||
pub timestamp: DateTime<Utc>,
|
||||
pub rate: f64, // cost per unit
|
||||
pub cost: f64, // total cost
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub enum UsageType {
|
||||
Data,
|
||||
Voice,
|
||||
SMS,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct RatingPlan {
|
||||
pub plan_id: String,
|
||||
pub name: String,
|
||||
pub data_rate: f64, // cost per MB
|
||||
pub voice_rate: f64, // cost per minute
|
||||
pub sms_rate: f64, // cost per SMS
|
||||
pub monthly_fee: f64,
|
||||
pub data_limit: u64, // bytes
|
||||
pub voice_limit: u64, // seconds
|
||||
pub sms_limit: u64, // count
|
||||
}
|
||||
|
||||
pub struct ChargingEngine {
|
||||
redis_client: redis::Client,
|
||||
rating_plans: HashMap<String, RatingPlan>,
|
||||
sync_interval: Duration,
|
||||
}
|
||||
|
||||
impl ChargingEngine {
|
||||
pub fn new(redis_url: &str, sync_interval_secs: u64) -> ChargingResult<Self> {
|
||||
let redis_client = redis::Client::open(redis_url)
|
||||
.context("Failed to create Redis client")?;
|
||||
|
||||
let mut rating_plans = HashMap::new();
|
||||
|
||||
// Default rating plans
|
||||
rating_plans.insert("basic".to_string(), RatingPlan {
|
||||
plan_id: "basic".to_string(),
|
||||
name: "Basic Plan".to_string(),
|
||||
data_rate: 0.01, // $0.01 per MB
|
||||
voice_rate: 0.05, // $0.05 per minute
|
||||
sms_rate: 0.10, // $0.10 per SMS
|
||||
monthly_fee: 10.0, // $10 monthly
|
||||
data_limit: 1_000_000_000, // 1GB
|
||||
voice_limit: 300, // 300 minutes
|
||||
sms_limit: 100, // 100 SMS
|
||||
});
|
||||
|
||||
rating_plans.insert("premium".to_string(), RatingPlan {
|
||||
plan_id: "premium".to_string(),
|
||||
name: "Premium Plan".to_string(),
|
||||
data_rate: 0.005, // $0.005 per MB
|
||||
voice_rate: 0.02, // $0.02 per minute
|
||||
sms_rate: 0.05, // $0.05 per SMS
|
||||
monthly_fee: 25.0, // $25 monthly
|
||||
data_limit: 5_000_000_000, // 5GB
|
||||
voice_limit: 1000, // 1000 minutes
|
||||
sms_limit: 500, // 500 SMS
|
||||
});
|
||||
|
||||
Ok(Self {
|
||||
redis_client,
|
||||
rating_plans,
|
||||
sync_interval: Duration::from_secs(sync_interval_secs),
|
||||
})
|
||||
}
|
||||
|
||||
pub async fn start(&self) -> ChargingResult<()> {
|
||||
info!("Starting charging engine");
|
||||
|
||||
let mut conn = self.redis_client.get_multiplexed_async_connection().await
|
||||
.context("Failed to get Redis connection")?;
|
||||
|
||||
// Test connection
|
||||
let _: () = redis::cmd("PING").query_async(&mut conn).await?;
|
||||
info!("Connected to Redis successfully");
|
||||
|
||||
// Start sync loop
|
||||
let mut ticker = interval(self.sync_interval);
|
||||
|
||||
loop {
|
||||
ticker.tick().await;
|
||||
|
||||
if let Err(e) = self.sync_usage_data(&mut conn).await {
|
||||
error!("Failed to sync usage data: {}", e);
|
||||
}
|
||||
|
||||
if let Err(e) = self.check_credit_limits(&mut conn).await {
|
||||
error!("Failed to check credit limits: {}", e);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn sync_usage_data(&self, conn: &mut redis::aio::MultiplexedConnection) -> Result<()> {
|
||||
debug!("Syncing usage data from packet gateway");
|
||||
|
||||
// Get all usage keys from packet gateway
|
||||
let pattern = "usage:*";
|
||||
let keys: Vec<String> = redis::cmd("KEYS")
|
||||
.arg(pattern)
|
||||
.query_async(conn)
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
|
||||
for key in keys {
|
||||
if let Some(ip) = key.strip_prefix("usage:") {
|
||||
// Get IMSI mapping from IP -> IMSI
|
||||
let imsi_key = format!("ip2imsi:{}", ip);
|
||||
if let Ok(imsi): Result<String, _> = redis::cmd("GET")
|
||||
.arg(&imsi_key)
|
||||
.query_async(conn)
|
||||
.await
|
||||
{
|
||||
// Get usage data
|
||||
let usage: u64 = redis::cmd("GET")
|
||||
.arg(&key)
|
||||
.query_async(conn)
|
||||
.await
|
||||
.unwrap_or(0);
|
||||
|
||||
if usage > 0 {
|
||||
// Update subscriber usage
|
||||
self.update_subscriber_usage(conn, &imsi, usage).await?;
|
||||
|
||||
// Clear processed usage
|
||||
let _: () = redis::cmd("DEL").arg(&key).query_async(conn).await?;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn update_subscriber_usage(
|
||||
&self,
|
||||
conn: &mut redis::aio::MultiplexedConnection,
|
||||
imsi: &str,
|
||||
data_bytes: u64
|
||||
) -> Result<()> {
|
||||
let account_key = format!("account:{}", imsi);
|
||||
|
||||
// Get current account
|
||||
let account_data: Option<String> = redis::cmd("GET")
|
||||
.arg(&account_key)
|
||||
.query_async(conn)
|
||||
.await?;
|
||||
|
||||
if let Some(data) = account_data {
|
||||
let mut account: SubscriberAccount = serde_json::from_str(&data)
|
||||
.context("Failed to deserialize account")?;
|
||||
|
||||
// Update data usage
|
||||
account.data_used += data_bytes;
|
||||
account.last_updated = Utc::now();
|
||||
|
||||
// Calculate cost
|
||||
if let Some(plan) = self.rating_plans.get(&account.imsi) {
|
||||
let mb_used = data_bytes as f64 / (1024.0 * 1024.0);
|
||||
let cost = mb_used * plan.data_rate;
|
||||
|
||||
// Deduct from balance
|
||||
account.balance -= (cost * 100.0) as i64; // Convert to cents
|
||||
|
||||
// Create usage event
|
||||
let usage_event = UsageEvent {
|
||||
imsi: imsi.to_string(),
|
||||
session_id: "default".to_string(), // TODO: Get actual session
|
||||
usage_type: UsageType::Data,
|
||||
volume: data_bytes,
|
||||
timestamp: Utc::now(),
|
||||
rate: plan.data_rate,
|
||||
cost,
|
||||
};
|
||||
|
||||
// Store usage event
|
||||
let event_key = format!("events:{}:{}", imsi, Utc::now().timestamp());
|
||||
let event_json = serde_json::to_string(&usage_event)?;
|
||||
let _: () = redis::cmd("SET")
|
||||
.arg(&event_key)
|
||||
.arg(&event_json)
|
||||
.arg("EX")
|
||||
.arg(86400) // Expire in 24 hours
|
||||
.query_async(conn)
|
||||
.await?;
|
||||
|
||||
// Check if over limit
|
||||
if account.data_used > account.data_limit {
|
||||
warn!("Subscriber {} exceeded data limit: {} > {}",
|
||||
imsi, account.data_used, account.data_limit);
|
||||
|
||||
// Block subscriber if significantly over limit
|
||||
if account.data_used > account.data_limit * 110 / 100 {
|
||||
account.status = AccountStatus::Blocked;
|
||||
|
||||
// Update block list in packet gateway
|
||||
let ip_key = format!("imsi2ip:{}", imsi);
|
||||
if let Ok(ip): Result<String, _> = redis::cmd("GET")
|
||||
.arg(&ip_key)
|
||||
.query_async(conn)
|
||||
.await
|
||||
{
|
||||
let block_key = format!("block:{}", ip);
|
||||
let _: () = redis::cmd("SET")
|
||||
.arg(&block_key)
|
||||
.arg(1)
|
||||
.query_async(conn)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Save updated account
|
||||
let account_json = serde_json::to_string(&account)?;
|
||||
let _: () = redis::cmd("SET")
|
||||
.arg(&account_key)
|
||||
.arg(&account_json)
|
||||
.query_async(conn)
|
||||
.await?;
|
||||
|
||||
info!("Updated usage for {}: {} bytes, cost: ${:.4}, balance: ${:.2}",
|
||||
imsi, data_bytes, cost, account.balance as f64 / 100.0);
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn check_credit_limits(&self, conn: &mut redis::aio::MultiplexedConnection) -> Result<()> {
|
||||
debug!("Checking credit limits");
|
||||
|
||||
// Get all account keys
|
||||
let pattern = "account:*";
|
||||
let keys: Vec<String> = redis::cmd("KEYS")
|
||||
.arg(pattern)
|
||||
.query_async(conn)
|
||||
.await
|
||||
.unwrap_or_default();
|
||||
|
||||
for key in keys {
|
||||
if let Some(imsi) = key.strip_prefix("account:") {
|
||||
let account_data: Option<String> = redis::cmd("GET")
|
||||
.arg(&key)
|
||||
.query_async(conn)
|
||||
.await?;
|
||||
|
||||
if let Some(data) = account_data {
|
||||
let account: SubscriberAccount = serde_json::from_str(&data)
|
||||
.context("Failed to deserialize account")?;
|
||||
|
||||
// Check if balance is negative
|
||||
if account.balance < 0 {
|
||||
warn!("Subscriber {} has negative balance: ${:.2}",
|
||||
imsi, account.balance as f64 / 100.0);
|
||||
|
||||
// Block if significantly negative
|
||||
if account.balance < -1000 { // -$10
|
||||
account.status = AccountStatus::Blocked;
|
||||
|
||||
// Update block list
|
||||
let ip_key = format!("imsi2ip:{}", imsi);
|
||||
if let Ok(ip): Result<String, _> = redis::cmd("GET")
|
||||
.arg(&ip_key)
|
||||
.query_async(conn)
|
||||
.await
|
||||
{
|
||||
let block_key = format!("block:{}", ip);
|
||||
let _: () = redis::cmd("SET")
|
||||
.arg(&block_key)
|
||||
.arg(1)
|
||||
.query_async(conn)
|
||||
.await?;
|
||||
}
|
||||
|
||||
// Save updated account
|
||||
let account_json = serde_json::to_string(&account)?;
|
||||
let _: () = redis::cmd("SET")
|
||||
.arg(&key)
|
||||
.arg(&account_json)
|
||||
.query_async(conn)
|
||||
.await?;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn create_account(&self, imsi: &str, plan_id: &str, initial_balance: i64) -> Result<()> {
|
||||
let mut conn = self.redis_client.get_multiplexed_async_connection().await?;
|
||||
|
||||
let plan = self.rating_plans.get(plan_id)
|
||||
.ok_or_else(|| anyhow::anyhow!("Rating plan '{}' not found", plan_id))?;
|
||||
|
||||
let account = SubscriberAccount {
|
||||
imsi: imsi.to_string(),
|
||||
balance: initial_balance,
|
||||
data_limit: plan.data_limit,
|
||||
data_used: 0,
|
||||
voice_limit: plan.voice_limit,
|
||||
voice_used: 0,
|
||||
sms_limit: plan.sms_limit,
|
||||
sms_used: 0,
|
||||
status: AccountStatus::Active,
|
||||
last_updated: Utc::now(),
|
||||
};
|
||||
|
||||
let account_key = format!("account:{}", imsi);
|
||||
let account_json = serde_json::to_string(&account)?;
|
||||
|
||||
let _: () = redis::cmd("SET")
|
||||
.arg(&account_key)
|
||||
.arg(&account_json)
|
||||
.query_async(&mut conn)
|
||||
.await?;
|
||||
|
||||
info!("Created account for {} with plan {} and balance ${:.2}",
|
||||
imsi, plan_id, initial_balance as f64 / 100.0);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn get_account(&self, imsi: &str) -> Result<Option<SubscriberAccount>> {
|
||||
let mut conn = self.redis_client.get_multiplexed_async_connection().await?;
|
||||
|
||||
let account_key = format!("account:{}", imsi);
|
||||
let account_data: Option<String> = redis::cmd("GET")
|
||||
.arg(&account_key)
|
||||
.query_async(&mut conn)
|
||||
.await?;
|
||||
|
||||
if let Some(data) = account_data {
|
||||
let account: SubscriberAccount = serde_json::from_str(&data)
|
||||
.context("Failed to deserialize account")?;
|
||||
Ok(Some(account))
|
||||
} else {
|
||||
Ok(None)
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn top_up_balance(&self, imsi: &str, amount: i64) -> Result<()> {
|
||||
let mut conn = self.redis_client.get_multiplexed_async_connection().await?;
|
||||
|
||||
let account_key = format!("account:{}", imsi);
|
||||
let account_data: Option<String> = redis::cmd("GET")
|
||||
.arg(&account_key)
|
||||
.query_async(&mut conn)
|
||||
.await?;
|
||||
|
||||
if let Some(data) = account_data {
|
||||
let mut account: SubscriberAccount = serde_json::from_str(&data)
|
||||
.context("Failed to deserialize account")?;
|
||||
|
||||
account.balance += amount;
|
||||
account.last_updated = Utc::now();
|
||||
|
||||
// Unblock if balance is restored and status was blocked
|
||||
if account.status == AccountStatus::Blocked && account.balance > 0 {
|
||||
account.status = AccountStatus::Active;
|
||||
|
||||
// Remove from block list
|
||||
let ip_key = format!("imsi2ip:{}", imsi);
|
||||
if let Ok(ip): Result<String, _> = redis::cmd("GET")
|
||||
.arg(&ip_key)
|
||||
.query_async(&mut conn)
|
||||
.await
|
||||
{
|
||||
let block_key = format!("block:{}", ip);
|
||||
let _: () = redis::cmd("DEL").arg(&block_key).query_async(&mut conn).await?;
|
||||
}
|
||||
}
|
||||
|
||||
let account_json = serde_json::to_string(&account)?;
|
||||
let _: () = redis::cmd("SET")
|
||||
.arg(&account_key)
|
||||
.arg(&account_json)
|
||||
.query_async(&mut conn)
|
||||
.await?;
|
||||
|
||||
info!("Topped up balance for {} by ${:.2}, new balance: ${:.2}",
|
||||
imsi, amount as f64 / 100.0, account.balance as f64 / 100.0);
|
||||
} else {
|
||||
return Err(anyhow::anyhow!("Account not found for IMSI: {}", imsi));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// API methods for HTTP endpoints
|
||||
|
||||
pub async fn test_connection(&self) -> Result<()> {
|
||||
let mut conn = self.redis_client.get_multiplexed_async_connection().await
|
||||
.context("Failed to get Redis connection")?;
|
||||
|
||||
let _: () = redis::cmd("PING").query_async(&mut conn).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn check_credit(&self, ip: &str, bytes_requested: u64) -> Result<bool> {
|
||||
let mut conn = self.redis_client.get_multiplexed_async_connection().await
|
||||
.context("Failed to get Redis connection")?;
|
||||
|
||||
let key = format!("credit:{}", ip);
|
||||
let credit: i64 = conn.get(&key).await.unwrap_or(0);
|
||||
|
||||
Ok(credit >= bytes_requested as i64)
|
||||
}
|
||||
|
||||
pub async fn get_balance(&self, ip: &str) -> Result<i64> {
|
||||
let mut conn = self.redis_client.get_multiplexed_async_connection().await
|
||||
.context("Failed to get Redis connection")?;
|
||||
|
||||
let key = format!("credit:{}", ip);
|
||||
let balance: i64 = conn.get(&key).await.unwrap_or(0);
|
||||
Ok(balance)
|
||||
}
|
||||
|
||||
pub async fn deduct_credit(&self, ip: &str, bytes_used: u64) -> Result<i64> {
|
||||
let mut conn = self.redis_client.get_multiplexed_async_connection().await
|
||||
.context("Failed to get Redis connection")?;
|
||||
|
||||
let key = format!("credit:{}", ip);
|
||||
let new_balance: i64 = conn.decr(&key, bytes_used).await.unwrap_or(0);
|
||||
Ok(new_balance)
|
||||
}
|
||||
|
||||
pub async fn add_credit(&self, ip: &str, bytes_to_add: u64) -> Result<i64> {
|
||||
let mut conn = self.redis_client.get_multiplexed_async_connection().await
|
||||
.context("Failed to get Redis connection")?;
|
||||
|
||||
let key = format!("credit:{}", ip);
|
||||
let new_balance: i64 = conn.incr(&key, bytes_to_add).await.unwrap_or(0);
|
||||
Ok(new_balance)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,235 @@
|
||||
use axum::{
|
||||
http::StatusCode,
|
||||
response::{IntoResponse, Response},
|
||||
Json,
|
||||
};
|
||||
use serde_json::json;
|
||||
use thiserror::Error;
|
||||
use tracing::error;
|
||||
|
||||
#[derive(Error, Debug)]
|
||||
pub enum ChargingError {
|
||||
#[error("Redis connection failed: {0}")]
|
||||
RedisConnection(String),
|
||||
|
||||
#[error("Redis operation failed: {0}")]
|
||||
RedisOperation(String),
|
||||
|
||||
#[error("Account not found: {0}")]
|
||||
AccountNotFound(String),
|
||||
|
||||
#[error("Insufficient credit: available={available}, requested={requested}")]
|
||||
InsufficientCredit { available: i64, requested: u64 },
|
||||
|
||||
#[error("Invalid request: {0}")]
|
||||
InvalidRequest(String),
|
||||
|
||||
#[error("Rating plan not found: {0}")]
|
||||
RatingPlanNotFound(String),
|
||||
|
||||
#[error("Account operation failed: {0}")]
|
||||
AccountOperation(String),
|
||||
|
||||
#[error("Serialization error: {0}")]
|
||||
Serialization(String),
|
||||
|
||||
#[error("Internal server error: {0}")]
|
||||
Internal(String),
|
||||
}
|
||||
|
||||
impl IntoResponse for ChargingError {
|
||||
fn into_response(self) -> Response {
|
||||
let (status, error_message) = match self {
|
||||
ChargingError::RedisConnection(ref msg) => {
|
||||
error!("Redis connection error: {}", msg);
|
||||
(StatusCode::SERVICE_UNAVAILABLE, "Redis connection failed")
|
||||
}
|
||||
ChargingError::RedisOperation(ref msg) => {
|
||||
error!("Redis operation error: {}", msg);
|
||||
(StatusCode::INTERNAL_SERVER_ERROR, "Redis operation failed")
|
||||
}
|
||||
ChargingError::AccountNotFound(ref imsi) => {
|
||||
(StatusCode::NOT_FOUND, "Account not found")
|
||||
}
|
||||
ChargingError::InsufficientCredit { available, requested } => {
|
||||
(StatusCode::PAYMENT_REQUIRED, "Insufficient credit")
|
||||
}
|
||||
ChargingError::InvalidRequest(ref msg) => {
|
||||
(StatusCode::BAD_REQUEST, "Invalid request")
|
||||
}
|
||||
ChargingError::RatingPlanNotFound(ref plan_id) => {
|
||||
(StatusCode::NOT_FOUND, "Rating plan not found")
|
||||
}
|
||||
ChargingError::AccountOperation(ref msg) => {
|
||||
error!("Account operation error: {}", msg);
|
||||
(StatusCode::INTERNAL_SERVER_ERROR, "Account operation failed")
|
||||
}
|
||||
ChargingError::Serialization(ref msg) => {
|
||||
error!("Serialization error: {}", msg);
|
||||
(StatusCode::INTERNAL_SERVER_ERROR, "Data serialization failed")
|
||||
}
|
||||
ChargingError::Internal(ref msg) => {
|
||||
error!("Internal error: {}", msg);
|
||||
(StatusCode::INTERNAL_SERVER_ERROR, "Internal server error")
|
||||
}
|
||||
};
|
||||
|
||||
let body = Json(json!({
|
||||
"error": error_message,
|
||||
"details": self.to_string(),
|
||||
"timestamp": chrono::Utc::now().to_rfc3339()
|
||||
}));
|
||||
|
||||
(status, body).into_response()
|
||||
}
|
||||
}
|
||||
|
||||
impl From<redis::RedisError> for ChargingError {
|
||||
fn from(err: redis::RedisError) -> Self {
|
||||
match err.kind() {
|
||||
redis::ErrorKind::Io => {
|
||||
ChargingError::RedisConnection(err.to_string())
|
||||
}
|
||||
_ => {
|
||||
ChargingError::RedisOperation(err.to_string())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl From<serde_json::Error> for ChargingError {
|
||||
fn from(err: serde_json::Error) -> Self {
|
||||
ChargingError::Serialization(err.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
impl From<anyhow::Error> for ChargingError {
|
||||
fn from(err: anyhow::Error) -> Self {
|
||||
ChargingError::Internal(err.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
pub type ChargingResult<T> = Result<T, ChargingError>;
|
||||
|
||||
pub trait ErrorContext<T> {
|
||||
fn with_context(self, context: &str) -> ChargingResult<T>;
|
||||
}
|
||||
|
||||
impl<T, E> ErrorContext<T> for Result<T, E>
|
||||
where
|
||||
E: Into<ChargingError>,
|
||||
{
|
||||
fn with_context(self, context: &str) -> ChargingResult<T> {
|
||||
self.map_err(|e| {
|
||||
let charging_err = e.into();
|
||||
match charging_err {
|
||||
ChargingError::RedisOperation(ref msg) => {
|
||||
ChargingError::RedisOperation(format!("{}: {}", context, msg))
|
||||
}
|
||||
ChargingError::AccountOperation(ref msg) => {
|
||||
ChargingError::AccountOperation(format!("{}: {}", context, msg))
|
||||
}
|
||||
ChargingError::Internal(ref msg) => {
|
||||
ChargingError::Internal(format!("{}: {}", context, msg))
|
||||
}
|
||||
other => other,
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Validation errors
|
||||
#[derive(Error, Debug)]
|
||||
pub enum ValidationError {
|
||||
#[error("Invalid IP address: {0}")]
|
||||
InvalidIP(String),
|
||||
|
||||
#[error("Invalid byte value: {0}")]
|
||||
InvalidBytes(String),
|
||||
|
||||
#[error("Missing required field: {0}")]
|
||||
MissingField(String),
|
||||
|
||||
#[error("Field too long: {0} (max {1})")]
|
||||
FieldTooLong(String, usize),
|
||||
|
||||
#[error("Invalid format: {0}")]
|
||||
InvalidFormat(String),
|
||||
}
|
||||
|
||||
impl From<ValidationError> for ChargingError {
|
||||
fn from(err: ValidationError) -> Self {
|
||||
ChargingError::InvalidRequest(err.to_string())
|
||||
}
|
||||
}
|
||||
|
||||
// Validation helpers
|
||||
pub fn validate_ip(ip: &str) -> Result<(), ValidationError> {
|
||||
if ip.is_empty() {
|
||||
return Err(ValidationError::MissingField("ip".to_string()));
|
||||
}
|
||||
|
||||
// Basic IP validation - could use std::net::IpAddr for more thorough validation
|
||||
if ip.len() > 45 { // IPv6 max length
|
||||
return Err(ValidationError::FieldTooLong("ip".to_string(), 45));
|
||||
}
|
||||
|
||||
// Simple validation for now - could be enhanced
|
||||
if ip.contains("..") || ip.starts_with('.') || ip.ends_with('.') {
|
||||
return Err(ValidationError::InvalidFormat(ip.to_string()));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn validate_bytes(bytes: u64) -> Result<(), ValidationError> {
|
||||
if bytes == 0 {
|
||||
return Err(ValidationError::InvalidBytes("bytes cannot be zero".to_string()));
|
||||
}
|
||||
|
||||
if bytes > 1_000_000_000_000 { // 1TB limit
|
||||
return Err(ValidationError::InvalidBytes("bytes value too large".to_string()));
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn test_validate_ip() {
|
||||
assert!(validate_ip("192.168.1.1").is_ok());
|
||||
assert!(validate_ip("10.0.0.1").is_ok());
|
||||
|
||||
assert!(validate_ip("").is_err());
|
||||
assert!(validate_ip("192.168..1.1").is_err());
|
||||
assert!(validate_ip(".192.168.1.1").is_err());
|
||||
assert!(validate_ip("192.168.1.1.").is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_validate_bytes() {
|
||||
assert!(validate_bytes(1024).is_ok());
|
||||
assert!(validate_bytes(1_000_000_000).is_ok());
|
||||
|
||||
assert!(validate_bytes(0).is_err());
|
||||
assert!(validate_bytes(2_000_000_000_000).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_error_conversion() {
|
||||
let redis_err = redis::RedisError::from((
|
||||
redis::ErrorKind::Io,
|
||||
"Connection failed",
|
||||
"Detailed error"
|
||||
));
|
||||
let charging_err: ChargingError = redis_err.into();
|
||||
|
||||
match charging_err {
|
||||
ChargingError::RedisConnection(_) => {}
|
||||
_ => panic!("Expected RedisConnection error"),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,23 +1,17 @@
|
||||
use axum::{
|
||||
extract::{Path, State},
|
||||
http::StatusCode,
|
||||
response::IntoResponse,
|
||||
routing::{get, post},
|
||||
Json, Router,
|
||||
};
|
||||
use redis::AsyncCommands;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::sync::Arc;
|
||||
use tower_http::cors::{Any, CorsLayer};
|
||||
use tracing::{info, warn};
|
||||
mod api;
|
||||
mod charging;
|
||||
mod errors;
|
||||
mod models;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct AppState {
|
||||
redis: redis::Client,
|
||||
}
|
||||
use anyhow::Result;
|
||||
use std::sync::Arc;
|
||||
use tracing::info;
|
||||
|
||||
use api::create_router;
|
||||
use charging::ChargingEngine;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
async fn main() -> Result<()> {
|
||||
// Initialize tracing
|
||||
tracing_subscriber::fmt()
|
||||
.with_target(false)
|
||||
@@ -27,33 +21,26 @@ async fn main() -> anyhow::Result<()> {
|
||||
// Load environment variables
|
||||
dotenv::dotenv().ok();
|
||||
|
||||
// Connect to Redis
|
||||
// Initialize charging engine
|
||||
let redis_url = std::env::var("REDIS_URI").unwrap_or_else(|_| "redis://127.0.0.1/".to_string());
|
||||
let redis_client = redis::Client::open(redis_url)?;
|
||||
let sync_interval = std::env::var("SYNC_INTERVAL")
|
||||
.unwrap_or_else(|_| "1".to_string())
|
||||
.parse::<u64>()
|
||||
.unwrap_or(1);
|
||||
|
||||
let charging_engine = Arc::new(ChargingEngine::new(&redis_url, sync_interval)?);
|
||||
|
||||
// Test connection
|
||||
let mut conn = redis_client.get_multiplexed_async_connection().await?;
|
||||
let _: () = redis::cmd("PING").query_async(&mut conn).await?;
|
||||
// Test Redis connection
|
||||
charging_engine.test_connection().await?;
|
||||
info!("Connected to Redis successfully");
|
||||
|
||||
let state = Arc::new(AppState {
|
||||
redis: redis_client,
|
||||
});
|
||||
// Create application state
|
||||
let state = api::AppState {
|
||||
charging_engine,
|
||||
};
|
||||
|
||||
// Build HTTP API with CORS
|
||||
let cors = CorsLayer::new()
|
||||
.allow_origin(Any)
|
||||
.allow_methods(Any)
|
||||
.allow_headers(Any);
|
||||
|
||||
let app = Router::new()
|
||||
.route("/v1/credit/:ip/check", post(check_credit))
|
||||
.route("/v1/credit/:ip/deduct", post(deduct_credit))
|
||||
.route("/v1/credit/:ip/add", post(add_credit))
|
||||
.route("/v1/credit/:ip/balance", get(get_balance))
|
||||
.route("/health", get(health_check))
|
||||
.layer(cors)
|
||||
.with_state(state);
|
||||
// Create router
|
||||
let app = create_router(state);
|
||||
|
||||
// Start server
|
||||
let port = std::env::var("SERVER_PORT").unwrap_or_else(|_| "8080".to_string());
|
||||
@@ -66,163 +53,3 @@ async fn main() -> anyhow::Result<()> {
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct CreditCheckRequest {
|
||||
bytes_requested: u64,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct CreditCheckResponse {
|
||||
allowed: bool,
|
||||
remaining_bytes: i64,
|
||||
}
|
||||
|
||||
/// POST /v1/credit/:ip/check
|
||||
/// Check if user has enough credit for data usage
|
||||
async fn check_credit(
|
||||
Path(ip): Path<String>,
|
||||
State(state): State<Arc<AppState>>,
|
||||
Json(req): Json<CreditCheckRequest>,
|
||||
) -> impl IntoResponse {
|
||||
let mut conn = match state.redis.get_multiplexed_async_connection().await {
|
||||
Ok(c) => c,
|
||||
Err(e) => {
|
||||
warn!("Redis connection failed: {}", e);
|
||||
return (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(CreditCheckResponse {
|
||||
allowed: false,
|
||||
remaining_bytes: 0,
|
||||
}),
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
let key = format!("credit:{}", ip);
|
||||
|
||||
// Get current credit balance
|
||||
let credit: i64 = conn.get(&key).await.unwrap_or(0);
|
||||
|
||||
let allowed = credit >= req.bytes_requested as i64;
|
||||
|
||||
info!(
|
||||
"Credit check for {}: {} bytes requested, {} bytes available, allowed: {}",
|
||||
ip, req.bytes_requested, credit, allowed
|
||||
);
|
||||
|
||||
(
|
||||
StatusCode::OK,
|
||||
Json(CreditCheckResponse {
|
||||
allowed,
|
||||
remaining_bytes: credit,
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct DeductRequest {
|
||||
bytes_used: u64,
|
||||
}
|
||||
|
||||
/// POST /v1/credit/:ip/deduct
|
||||
/// Deduct bytes from user's credit balance
|
||||
async fn deduct_credit(
|
||||
Path(ip): Path<String>,
|
||||
State(state): State<Arc<AppState>>,
|
||||
Json(req): Json<DeductRequest>,
|
||||
) -> impl IntoResponse {
|
||||
let mut conn = match state.redis.get_multiplexed_async_connection().await {
|
||||
Ok(c) => c,
|
||||
Err(e) => {
|
||||
warn!("Redis connection failed: {}", e);
|
||||
return StatusCode::INTERNAL_SERVER_ERROR;
|
||||
}
|
||||
};
|
||||
|
||||
let key = format!("credit:{}", ip);
|
||||
|
||||
// Atomically subtract bytes
|
||||
let new_balance: i64 = conn.decr(&key, req.bytes_used).await.unwrap_or(0);
|
||||
|
||||
info!(
|
||||
"User {} deducted {} bytes, remaining: {}",
|
||||
ip, req.bytes_used, new_balance
|
||||
);
|
||||
|
||||
StatusCode::OK
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct AddCreditRequest {
|
||||
bytes_to_add: u64,
|
||||
}
|
||||
|
||||
/// POST /v1/credit/:ip/add
|
||||
/// Add bytes to user's credit balance
|
||||
async fn add_credit(
|
||||
Path(ip): Path<String>,
|
||||
State(state): State<Arc<AppState>>,
|
||||
Json(req): Json<AddCreditRequest>,
|
||||
) -> impl IntoResponse {
|
||||
let mut conn = match state.redis.get_multiplexed_async_connection().await {
|
||||
Ok(c) => c,
|
||||
Err(e) => {
|
||||
warn!("Redis connection failed: {}", e);
|
||||
return StatusCode::INTERNAL_SERVER_ERROR;
|
||||
}
|
||||
};
|
||||
|
||||
let key = format!("credit:{}", ip);
|
||||
|
||||
// Atomically add bytes
|
||||
let new_balance: i64 = conn.incr(&key, req.bytes_to_add).await.unwrap_or(0);
|
||||
|
||||
info!(
|
||||
"User {} added {} bytes, new balance: {}",
|
||||
ip, req.bytes_to_add, new_balance
|
||||
);
|
||||
|
||||
StatusCode::OK
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct BalanceResponse {
|
||||
ip: String,
|
||||
balance_bytes: i64,
|
||||
}
|
||||
|
||||
/// GET /v1/credit/:ip/balance
|
||||
/// Get current credit balance
|
||||
async fn get_balance(
|
||||
Path(ip): Path<String>,
|
||||
State(state): State<Arc<AppState>>,
|
||||
) -> impl IntoResponse {
|
||||
let mut conn = match state.redis.get_multiplexed_async_connection().await {
|
||||
Ok(c) => c,
|
||||
Err(_) => {
|
||||
return (
|
||||
StatusCode::INTERNAL_SERVER_ERROR,
|
||||
Json(BalanceResponse {
|
||||
ip: ip.clone(),
|
||||
balance_bytes: 0,
|
||||
}),
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
let key = format!("credit:{}", ip);
|
||||
let balance: i64 = conn.get(&key).await.unwrap_or(0);
|
||||
|
||||
(
|
||||
StatusCode::OK,
|
||||
Json(BalanceResponse {
|
||||
ip,
|
||||
balance_bytes: balance,
|
||||
}),
|
||||
)
|
||||
}
|
||||
|
||||
async fn health_check() -> &'static str {
|
||||
"OK"
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user