speedtest and ai for toxic comment moderate

proxy-ws/src/classifier.rs

@@ -0,0 +1,75 @@
+use anyhow::{anyhow, Result};
+use tokenizers::Tokenizer;
+use tract_onnx::prelude::*;
+
+pub struct ToxicityClassifier {
+    model: SimplePlan<TypedFact, Box<dyn TypedOp>, Graph<TypedFact, Box<dyn TypedOp>>>,
+    tokenizer: Tokenizer,
+}
+
+impl ToxicityClassifier {
+    pub fn new(model_path: &str, tokenizer_path: &str) -> Result<Self> {
+        let tokenizer = Tokenizer::from_file(tokenizer_path)
+            .map_err(|e| anyhow!("Tokenizer load error: {}", e))?;
+
+        let model = tract_onnx::onnx()
+            .model_for_path(model_path)?
+            .into_optimized()?
+            .into_runnable()?;
+
+        Ok(Self { model, tokenizer })
+    }
+
+    pub fn is_toxic(&self, text: &str) -> bool {
+        let encoding = match self.tokenizer.encode(text, true) {
+            Ok(e) => e,
+            Err(_) => return false,
+        };
+
+        let ids: Vec<i64> = encoding.get_ids().iter().map(|&x| x as i64).collect();
+        let mask: Vec<i64> = encoding
+            .get_attention_mask()
+            .iter()
+            .map(|&x| x as i64)
+            .collect();
+        let type_ids: Vec<i64> = encoding.get_type_ids().iter().map(|&x| x as i64).collect();
+
+        let seq_len = ids.len();
+
+        // В tract .into_shape() возвращает Tensor напрямую, unwrap не нужен
+        let ids_tensor = tensor1(&ids)
+            .into_shape(&[1, seq_len])
+            .expect("Shape error");
+        let mask_tensor = tensor1(&mask)
+            .into_shape(&[1, seq_len])
+            .expect("Shape error");
+        let type_tensor = tensor1(&type_ids)
+            .into_shape(&[1, seq_len])
+            .expect("Shape error");
+
+        let result = self.model.run(tvec![
+            ids_tensor.into(),
+            mask_tensor.into(),
+            type_tensor.into()
+        ]);
+
+        let outputs = match result {
+            Ok(out) => out,
+            Err(e) => {
+                println!("AI Link Error: {:?}", e);
+                return false;
+            }
+        };
+
+        // Извлекаем результат
+        let logits = match outputs[0].to_array_view::<f32>() {
+            Ok(view) => view,
+            Err(_) => return false,
+        };
+
+        logits.iter().any(|&x| {
+            let prob = 1.0 / (1.0 + (-x).exp());
+            prob > 0.8
+        })
+    }
+}

proxy-ws/src/main.rs

@@ -1,7 +1,7 @@
 use axum::{
     body::Bytes,
-    extract::DefaultBodyLimit,
-    http::{header, HeaderValue, Method, StatusCode},
+    extract::{DefaultBodyLimit, State},
+    http::{header, StatusCode},
     response::IntoResponse,
     routing::{get, post},
     Json, Router,
@@ -10,8 +10,10 @@ use regex::Regex;
 use reqwest::Client;
 use rustrict::CensorStr;
 use serde::{Deserialize, Serialize};
-use std::sync::LazyLock;
-use tower_http::cors::CorsLayer;
+use std::sync::{Arc, LazyLock};
+
+mod classifier;
+use classifier::ToxicityClassifier;
 
 #[derive(Deserialize, Serialize, Debug)]
 struct CommentPayload {
@@ -24,7 +26,6 @@ static BAD_WORDS_REGEX: LazyLock<Regex> = LazyLock::new(|| {
 });
 
 const SAFE_WORDS: &[&str] = &[
-    // Защита корня "ебл/ебат" (самые частые ложные срабатывания)
     "колебат",
     "оскорблят",
     "употреблят",
@@ -40,10 +41,8 @@ const SAFE_WORDS: &[&str] = &[
     "хлебл",
     "дебел",
     "колебл",
-    // Защита корня "хуй"
     "страхуй",
     "подстрахуй",
-    // Защита корня "бля"
     "рубл",
     "сабл",
     "челяб",
@@ -52,37 +51,39 @@ const SAFE_WORDS: &[&str] = &[
     "дубл",
     "гибл",
     "корабл",
-    // Защита корня "жоп"
     "джоп",
     "джобс",
     "прожора",
-    // Защита корня "пидар/пидор"
     "скипидар",
     "купидор",
     "коридор",
-    // Защита корня "чмо"
     "чмок",
     "ячмен",
     "черемух",
-    // Защита корня "мраз"
     "мразер",
     "замороз",
     "мороз",
     "мрамор",
-    // Защита корня "сука"
     "рисука",
     "присука",
     "бамбука",
-    // Защита корня "жид"
     "ожид",
     "жидкос",
     "жидк",
-    // Защита корня "даун"
     "даунгрейд",
     "адаун",
 ];
 
-// Заменяет легальные слова-исключения на плейсхолдер перед проверкой
+static RANDOM_PAYLOAD: LazyLock<Bytes> = LazyLock::new(|| {
+    let mut data = vec![0u8; 25 * 1024 * 1024];
+    let mut state: u32 = 123456789;
+    for byte in data.iter_mut() {
+        state = state.wrapping_mul(1103515245).wrapping_add(12345);
+        *byte = (state >> 16) as u8;
+    }
+    Bytes::from(data)
+});
+
 fn apply_allowlist(text: &str) -> String {
     let mut safe_text = text.to_string();
     for &safe_word in SAFE_WORDS {
@@ -91,24 +92,18 @@ fn apply_allowlist(text: &str) -> String {
     safe_text
 }
 
-// Функция жесткой нормализации текста перед проверкой
 fn normalize_for_filter(input: &str) -> String {
     let mut normalized = String::new();
     let mut prev_char: Option<char> = None;
-
-    let lower_input = input.to_lowercase();
-
-    for ch in lower_input.chars() {
+    for ch in input.to_lowercase().chars() {
         if ch.is_whitespace() {
             normalized.push(' ');
             prev_char = Some(' ');
             continue;
         }
-
         if !ch.is_alphanumeric() {
             continue;
         }
-
         let mapped = match ch {
             'a' | '@' => 'а',
             'b' => 'в',
@@ -125,17 +120,14 @@ fn normalize_for_filter(input: &str) -> String {
             '4' => 'ч',
             _ => ch,
         };
-
         if let Some(p) = prev_char {
             if p == mapped && mapped != ' ' {
                 continue;
             }
         }
-
         normalized.push(mapped);
         prev_char = Some(mapped);
     }
-
     normalized
 }
 
@@ -144,115 +136,106 @@ fn is_spam(input: &str) -> bool {
     if alpha_chars.is_empty() {
         return false;
     }
-
-    let uppercase_count = alpha_chars.chars().filter(|c| c.is_uppercase()).count();
-    let uppercase_ratio = uppercase_count as f64 / alpha_chars.len() as f64;
-
+    let uppercase_ratio =
+        alpha_chars.chars().filter(|c| c.is_uppercase()).count() as f64 / alpha_chars.len() as f64;
     if input.len() > 10 && uppercase_ratio > 0.7 {
         return true;
     }
+    input
+        .split_whitespace()
+        .any(|word| word.len() > 35 && !word.starts_with("http"))
+}
 
-    for word in input.split_whitespace() {
-        if word.len() > 35 && !word.starts_with("http") {
-            return true;
-        }
-    }
-
-    false
+struct AppState {
+    ai_classifier: ToxicityClassifier,
 }
 
 #[tokio::main]
 async fn main() {
-    // 🔥 Исправленный CORS: Явно указываем протоколы и разрешаем локалхост для тестов
-    let cors = CorsLayer::new()
-        .allow_origin([
-            "http://localhost:3000".parse::<HeaderValue>().unwrap(),
-            "https://netrunner-vpn.com".parse::<HeaderValue>().unwrap(),
-        ])
-        .allow_methods([Method::GET, Method::POST])
-        .allow_headers([header::CONTENT_TYPE]);
+    println!("🧠 Booting AI Neural Link (Pure Rust Engine)...");
+    let classifier = ToxicityClassifier::new("models/model.onnx", "models/tokenizer.json")
+        .expect("CRITICAL: Failed to load Neural Link model");
+
+    let state = Arc::new(AppState {
+        ai_classifier: classifier,
+    });
 
     let app = Router::new()
         .route("/api/comments", post(create_comment))
         .route("/api/speedtest/ping", get(|| async { "pong" }))
         .route("/api/speedtest/download", get(speedtest_download))
         .route("/api/speedtest/upload", post(speedtest_upload))
-        // 🔥 Увеличиваем лимит тела запроса до 50 Мегабайт для спидтеста (дефолт 2MB)
-        .layer(DefaultBodyLimit::max(50 * 1024 * 1024))
-        .layer(cors);
+        .layer(DefaultBodyLimit::max(100 * 1024 * 1024))
+        .with_state(state);
 
     println!("🚀 Netrunner Gateway started on 0.0.0.0:3001");
     let listener = tokio::net::TcpListener::bind("0.0.0.0:3001").await.unwrap();
     axum::serve(listener, app).await.unwrap();
 }
 
-async fn create_comment(Json(payload): Json<CommentPayload>) -> impl IntoResponse {
+async fn create_comment(
+    State(state): State<Arc<AppState>>,
+    Json(payload): Json<CommentPayload>,
+) -> impl IntoResponse {
     if is_spam(&payload.message) || is_spam(&payload.nickname) {
-        println!("🗑️ Blocked by SPAM filter: {}", payload.nickname);
-        return (
-            StatusCode::NOT_ACCEPTABLE,
-            "Transmission rejected: Spam pattern detected.",
-        );
+        return (StatusCode::NOT_ACCEPTABLE, "Spam detected.");
     }
-
     if payload.message.is_inappropriate() || payload.nickname.is_inappropriate() {
-        println!("🛡️ Blocked by rustrict: {}", payload.nickname);
-        return (
-            StatusCode::NOT_ACCEPTABLE,
-            "Transmission rejected: Policy violation.",
-        );
+        return (StatusCode::NOT_ACCEPTABLE, "Policy violation.");
     }
 
     let normalized_msg = normalize_for_filter(&payload.message);
     let normalized_nick = normalize_for_filter(&payload.nickname);
-
-    let safe_msg = apply_allowlist(&normalized_msg);
-    let safe_nick = apply_allowlist(&normalized_nick);
-
-    if BAD_WORDS_REGEX.is_match(&safe_msg) || BAD_WORDS_REGEX.is_match(&safe_nick) {
-        println!(
-            "🛡️ Blocked by Deep Regex: [Original: {}] -> [Normalized: {}]",
-            payload.message, safe_msg
-        );
-        return (
-            StatusCode::NOT_ACCEPTABLE,
-            "Transmission rejected: Policy violation.",
-        );
+    if BAD_WORDS_REGEX.is_match(&apply_allowlist(&normalized_msg))
+        || BAD_WORDS_REGEX.is_match(&apply_allowlist(&normalized_nick))
+    {
+        return (StatusCode::NOT_ACCEPTABLE, "Policy violation.");
     }
 
-    let pb_base_url =
-        std::env::var("PB_INTERNAL_URL").unwrap_or_else(|_| "http://localhost:8090".to_string());
-    let pb_url = format!("{}/api/collections/comments/records", pb_base_url);
+    let message = payload.message.clone();
+    let nickname = payload.nickname.clone();
+    let is_toxic = tokio::task::spawn_blocking(move || {
+        state.ai_classifier.is_toxic(&message) || state.ai_classifier.is_toxic(&nickname)
+    })
+    .await
+    .unwrap_or(false);
 
+    if is_toxic {
+        return (StatusCode::NOT_ACCEPTABLE, "Behavioral policy violation.");
+    }
+
+    let pb_url = format!(
+        "{}/api/collections/comments/records",
+        std::env::var("PB_INTERNAL_URL").unwrap_or_else(|_| "http://localhost:8090".to_string())
+    );
     let client = Client::new();
-    let pb_response = client.post(pb_url).json(&payload).send().await;
-
-    match pb_response {
-        Ok(res) if res.status().is_success() => {
-            (StatusCode::OK, "Transmission secured and delivered.")
-        }
-        Ok(res) => {
-            println!("⚠️ PocketBase rejected: {:?}", res.status());
-            (StatusCode::BAD_REQUEST, "Database rejected transmission.")
-        }
-        Err(e) => {
-            println!("💥 PB Connection error: {}", e);
-            (StatusCode::INTERNAL_SERVER_ERROR, "Uplink offline.")
-        }
+    match client.post(pb_url).json(&payload).send().await {
+        Ok(res) if res.status().is_success() => (StatusCode::OK, "Delivered."),
+        _ => (StatusCode::BAD_REQUEST, "Uplink offline."),
     }
 }
 
 async fn speedtest_download() -> impl IntoResponse {
-    let payload_size = 25 * 1024 * 1024; // 25 MB
-    let data = vec![0u8; payload_size];
+    let body = RANDOM_PAYLOAD.clone();
+    let len = body.len();
+
     (
-        [(axum::http::header::CONTENT_TYPE, "application/octet-stream")],
-        data,
+        [
+            (
+                header::CONTENT_TYPE,
+                header::HeaderValue::from_static("application/octet-stream"),
+            ),
+            (header::CONTENT_LENGTH, header::HeaderValue::from(len)),
+            (
+                header::CACHE_CONTROL,
+                header::HeaderValue::from_static("no-store, no-cache, must-revalidate, max-age=0"),
+            ),
+            (header::PRAGMA, header::HeaderValue::from_static("no-cache")),
+        ],
+        body,
     )
 }
 
-// 🔥 Возвращаем `body: Bytes`. Благодаря `DefaultBodyLimit` в main(),
-// Axum 0.7 успешно примет в память файл до 50МБ без ручной работы со стримами.
 async fn speedtest_upload(body: Bytes) -> impl IntoResponse {
     (StatusCode::OK, format!("Received {} bytes", body.len()))
 }