This commit is contained in:
2026-06-22 15:19:46 +01:00
parent de96bc030d
commit 808d6ac5db
21 changed files with 948 additions and 282 deletions
+5 -1
View File
@@ -4,5 +4,9 @@ version = "0.1.0"
edition = "2024"
[dependencies]
sea-orm = { version = "1", features = ["sqlx-sqlite", "runtime-tokio-rustls", "macros"] }
sea-orm = { version = "1", features = [ "sqlx-sqlite", "runtime-tokio-rustls", "macros" ] }
serde = { version = "1", features = ["derive"] }
chrono = { version = "0.4", features = ["serde"] }
rand = "0.10.1"
argon2 = "0.5.3"
uuid = { version = "1.23.3", features = ["v4"] }
+54
View File
@@ -0,0 +1,54 @@
use sea_orm::{ActiveValue::Set, entity::prelude::*};
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "auth_session")]
pub struct Model {
#[sea_orm(primary_key)]
pub id: i32,
pub id_user: i32,
pub value: String,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
#[sea_orm(
belongs_to = "super::users::Entity",
from = "Column::IdUser",
to = "super::users::Column::Id"
)]
Users,
}
impl Related<super::users::Entity> for Entity {
fn to() -> RelationDef {
Relation::Users.def()
}
}
impl ActiveModelBehavior for ActiveModel {}
impl Entity {
pub async fn create(db: &DatabaseConnection, user_id: i32) -> Result<Model, DbErr> {
let value = uuid::Uuid::new_v4().to_string();
ActiveModel {
id_user: Set(user_id),
value: Set(value),
..Default::default()
}
.insert(db)
.await
}
pub async fn find_by_user_id(
db: &DatabaseConnection,
user_id: i32,
) -> Result<Option<Model>, DbErr> {
Entity::find()
.filter(Column::IdUser.eq(user_id))
.one(db)
.await
}
}
impl ActiveModel {}
+2
View File
@@ -1,3 +1,5 @@
pub mod auth_session;
pub mod prelude;
pub mod stream_key;
pub mod stream_session;
pub mod users;
+2
View File
@@ -1,2 +1,4 @@
pub use super::auth_session::Entity as AuthSession;
pub use super::stream_key::Entity as StreamKey;
pub use super::stream_session::Entity as StreamSession;
pub use super::users::Entity as Users;
+54 -1
View File
@@ -1,4 +1,4 @@
use sea_orm::entity::prelude::*;
use sea_orm::{Set, entity::prelude::*};
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
@@ -8,8 +8,10 @@ pub struct Model {
pub id: i32,
#[sea_orm(unique)]
pub key_value: String,
pub user_id: i32,
pub label: String,
pub is_active: bool,
pub is_unlisted: bool,
pub created_at: DateTimeUtc,
}
@@ -17,6 +19,12 @@ pub struct Model {
pub enum Relation {
#[sea_orm(has_many = "super::stream_session::Entity")]
StreamSession,
#[sea_orm(
belongs_to = "super::users::Entity",
from = "Column::UserId",
to = "super::users::Column::Id"
)]
User,
}
impl Related<super::stream_session::Entity> for Entity {
@@ -25,4 +33,49 @@ impl Related<super::stream_session::Entity> for Entity {
}
}
impl Related<super::users::Entity> for Entity {
fn to() -> RelationDef {
Relation::User.def()
}
}
impl ActiveModelBehavior for ActiveModel {}
impl Entity {
pub async fn create(
db: &DatabaseConnection,
user_id: i32,
key_value: String,
label: String,
is_unlisted: bool,
) -> Result<Model, DbErr> {
ActiveModel {
user_id: Set(user_id),
key_value: Set(key_value),
label: Set(label),
is_active: Set(false),
is_unlisted: Set(is_unlisted),
created_at: Set(chrono::Utc::now()),
..Default::default()
}
.insert(db)
.await
}
pub async fn find_by_key(
db: &DatabaseConnection,
key_value: &str,
) -> Result<Option<Model>, DbErr> {
Entity::find()
.filter(Column::KeyValue.eq(key_value))
.one(db)
.await
}
pub async fn find_by_user(db: &DatabaseConnection, user_id: i32) -> Result<Vec<Model>, DbErr> {
Entity::find()
.filter(Column::UserId.eq(user_id))
.all(db)
.await
}
}
+44 -1
View File
@@ -1,4 +1,8 @@
use sea_orm::entity::prelude::*;
use sea_orm::{
ActiveValue::{NotSet, Set},
entity::prelude::*,
sqlx::types::chrono::{self, Utc},
};
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
@@ -28,3 +32,42 @@ impl Related<super::stream_key::Entity> for Entity {
}
impl ActiveModelBehavior for ActiveModel {}
impl Model {
pub async fn create_stream_session(
db: &DatabaseConnection,
stream_key_id: i32,
started_at: chrono::DateTime<Utc>,
) -> Result<Model, DbErr> {
ActiveModel {
id: NotSet,
stream_key_id: Set(stream_key_id),
started_at: Set(started_at),
..Default::default()
}
.insert(db)
.await
}
pub async fn get_stream_session(
db: &DatabaseConnection,
stream_session_id: i32,
) -> Result<Model, DbErr> {
Entity::find_by_id(stream_session_id)
.one(db)
.await?
.ok_or(DbErr::RecordNotFound(format!(
"stream_session {stream_session_id}"
)))
}
}
impl ActiveModel {
pub async fn finish_stream_session(
mut self,
db: &DatabaseConnection,
ended_at: chrono::DateTime<Utc>,
) -> Result<Model, DbErr> {
self.ended_at = Set(Some(ended_at));
self.update(db).await
}
}
+92
View File
@@ -0,0 +1,92 @@
use sea_orm::{Set, entity::prelude::*};
use serde::{Deserialize, Serialize};
#[derive(Clone, Debug, PartialEq, DeriveEntityModel, Serialize, Deserialize)]
#[sea_orm(table_name = "users")]
pub struct Model {
#[sea_orm(primary_key)]
pub id: i32,
pub username: String,
pub hashed_password: String,
pub stream_key_limit: i32,
}
#[derive(Copy, Clone, Debug, EnumIter, DeriveRelation)]
pub enum Relation {
#[sea_orm(has_many = "super::stream_key::Entity")]
StreamKey,
#[sea_orm(has_many = "super::auth_session::Entity")]
AuthSession,
}
impl Related<super::stream_key::Entity> for Entity {
fn to() -> RelationDef {
Relation::StreamKey.def()
}
}
impl Related<super::auth_session::Entity> for Entity {
fn to() -> RelationDef {
Relation::AuthSession.def()
}
}
impl ActiveModelBehavior for ActiveModel {}
impl Entity {
pub async fn create(
db: &DatabaseConnection,
username: String,
password: String,
) -> Result<Model, DbErr> {
ActiveModel {
username: Set(username),
hashed_password: Set(password),
..Default::default()
}
.insert(db)
.await
}
pub async fn find_by_auth_session(
db: &DatabaseConnection,
auth_session: String,
) -> Result<Option<Model>, DbErr> {
let sessions = crate::auth_session::Entity::find()
.filter(crate::auth_session::Column::Value.eq(auth_session.to_string()))
.one(db)
.await?;
if let Some(x) = sessions {
Entity::find_by_id(x.id_user).one(db).await
} else {
return Ok(None);
}
}
}
impl ActiveModel {
pub async fn update_username(
mut self,
db: &DatabaseConnection,
username: String,
) -> Result<Model, DbErr> {
self.username = Set(username);
self.update(db).await
}
pub async fn update_password(
mut self,
db: &DatabaseConnection,
hashed_password: String,
) -> Result<Model, DbErr> {
self.hashed_password = Set(hashed_password);
self.update(db).await
}
pub async fn change_stream_key_limit(
mut self,
db: &DatabaseConnection,
new_limit: i32,
) -> Result<Model, DbErr> {
self.stream_key_limit = Set(new_limit);
self.update(db).await
}
}
+10 -2
View File
@@ -1,12 +1,20 @@
pub use sea_orm_migration::prelude::*;
mod m20260616_000001_create_tables;
mod m20260616_000001_create_users;
mod m20260616_000002_create_stream_key;
mod m20260616_000003_create_stream_session;
mod m20260616_000004_create_auth_session;
pub struct Migrator;
#[async_trait::async_trait]
impl MigratorTrait for Migrator {
fn migrations() -> Vec<Box<dyn MigrationTrait>> {
vec![Box::new(m20260616_000001_create_tables::Migration)]
vec![
Box::new(m20260616_000001_create_users::Migration),
Box::new(m20260616_000002_create_stream_key::Migration),
Box::new(m20260616_000003_create_stream_session::Migration),
Box::new(m20260616_000004_create_auth_session::Migration),
]
}
}
@@ -0,0 +1,48 @@
use sea_orm_migration::prelude::*;
#[derive(DeriveMigrationName)]
pub struct Migration;
#[async_trait::async_trait]
impl MigrationTrait for Migration {
async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
manager
.create_table(
Table::create()
.table(Users::Table)
.if_not_exists()
.col(
ColumnDef::new(Users::Id)
.integer()
.not_null()
.auto_increment()
.primary_key(),
)
.col(
ColumnDef::new(Users::Username)
.string()
.not_null()
.unique_key(),
)
.col(ColumnDef::new(Users::HashedPassword).string())
.col(ColumnDef::new(Users::StreamKeyLimit).integer().default(3))
.to_owned(),
)
.await
}
async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
manager
.drop_table(Table::drop().table(Users::Table).to_owned())
.await
}
}
#[derive(Iden)]
pub enum Users {
Table,
Id,
Username,
HashedPassword,
StreamKeyLimit,
}
@@ -1,5 +1,7 @@
use sea_orm_migration::prelude::*;
use super::m20260616_000001_create_users::Users;
#[derive(DeriveMigrationName)]
pub struct Migration;
@@ -18,6 +20,7 @@ impl MigrationTrait for Migration {
.auto_increment()
.primary_key(),
)
.col(ColumnDef::new(StreamKey::UserId).integer().not_null())
.col(
ColumnDef::new(StreamKey::KeyValue)
.string()
@@ -31,73 +34,38 @@ impl MigrationTrait for Migration {
.not_null()
.default(true),
)
.col(ColumnDef::new(StreamKey::CreatedAt).date_time().not_null())
.to_owned(),
)
.await?;
manager
.create_table(
Table::create()
.table(StreamSession::Table)
.if_not_exists()
.col(
ColumnDef::new(StreamSession::Id)
.integer()
ColumnDef::new(StreamKey::IsUnlisted)
.boolean()
.not_null()
.auto_increment()
.primary_key(),
.default(true),
)
.col(ColumnDef::new(StreamSession::StreamKeyId).integer().not_null())
.col(ColumnDef::new(StreamSession::StartedAt).date_time().not_null())
.col(ColumnDef::new(StreamSession::EndedAt).date_time().null())
.col(ColumnDef::new(StreamKey::CreatedAt).date_time().not_null())
.foreign_key(
ForeignKey::create()
.from(StreamSession::Table, StreamSession::StreamKeyId)
.to(StreamKey::Table, StreamKey::Id),
.from(StreamKey::Table, StreamKey::UserId)
.to(Users::Table, Users::Id),
)
.to_owned(),
)
.await?;
// Seed the default stream key so OBS can connect out of the box
manager
.get_connection()
.execute_unprepared(
"INSERT OR IGNORE INTO stream_key (key_value, label, is_active, created_at) \
VALUES ('test', 'Default', 1, datetime('now'))",
)
.await?;
Ok(())
.await
}
async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
manager
.drop_table(Table::drop().table(StreamSession::Table).to_owned())
.await?;
manager
.drop_table(Table::drop().table(StreamKey::Table).to_owned())
.await?;
Ok(())
.await
}
}
#[derive(Iden)]
enum StreamKey {
pub enum StreamKey {
Table,
Id,
UserId,
KeyValue,
Label,
IsActive,
IsUnlisted,
CreatedAt,
}
#[derive(Iden)]
enum StreamSession {
Table,
Id,
StreamKeyId,
StartedAt,
EndedAt,
}
@@ -0,0 +1,58 @@
use sea_orm_migration::prelude::*;
use super::m20260616_000002_create_stream_key::StreamKey;
#[derive(DeriveMigrationName)]
pub struct Migration;
#[async_trait::async_trait]
impl MigrationTrait for Migration {
async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
manager
.create_table(
Table::create()
.table(StreamSession::Table)
.if_not_exists()
.col(
ColumnDef::new(StreamSession::Id)
.integer()
.not_null()
.auto_increment()
.primary_key(),
)
.col(
ColumnDef::new(StreamSession::StreamKeyId)
.integer()
.not_null(),
)
.col(
ColumnDef::new(StreamSession::StartedAt)
.date_time()
.not_null(),
)
.col(ColumnDef::new(StreamSession::EndedAt).date_time().null())
.foreign_key(
ForeignKey::create()
.from(StreamSession::Table, StreamSession::StreamKeyId)
.to(StreamKey::Table, StreamKey::Id),
)
.to_owned(),
)
.await
}
async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
manager
.drop_table(Table::drop().table(StreamSession::Table).to_owned())
.await
}
}
#[derive(Iden)]
pub enum StreamSession {
Table,
Id,
StreamKeyId,
StartedAt,
EndedAt,
}
@@ -0,0 +1,48 @@
use sea_orm_migration::prelude::*;
use crate::m20260616_000001_create_users::Users;
#[derive(DeriveMigrationName)]
pub struct Migration;
#[async_trait::async_trait]
impl MigrationTrait for Migration {
async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
manager
.create_table(
Table::create()
.table(AuthSession::Table)
.if_not_exists()
.col(
ColumnDef::new(AuthSession::Id)
.integer()
.not_null()
.auto_increment()
.primary_key(),
)
.col(ColumnDef::new(AuthSession::IdUser).integer().not_null())
.col(ColumnDef::new(AuthSession::Value).string())
.foreign_key(
ForeignKey::create()
.from(AuthSession::Table, AuthSession::IdUser)
.to(Users::Table, Users::Id),
)
.to_owned(),
)
.await
}
async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
manager
.drop_table(Table::drop().table(AuthSession::Table).to_owned())
.await
}
}
#[derive(Iden)]
pub enum AuthSession {
Table,
Id,
IdUser,
Value,
}
+7 -4
View File
@@ -9,14 +9,17 @@ path = "src/main.rs"
[dependencies]
async-broadcast = "0.7.2"
bytes = "1.11.1"
bytes = "1"
dashmap = "6.2.1"
rand = "0.10.1"
rml_rtmp = "0.8.0"
str0m = "0.20.0"
tokio = { version = "1", features = ["full"] }
hyper = { version = "1", features = ["full"] }
http-body-util = "0.1"
hyper-util = { version = "0.1", features = ["full"] }
axum = "0.8"
serde = { version = "1.0.228", features = ["serde_derive"] }
serde_json = "1.0.150"
sea-orm = { version = "1", features = [ "sqlx-sqlite", "runtime-tokio-rustls", "macros" ] }
entity = {path = "../entity"}
migration = {path = "../migration"}
argon2 = "0.5.3"
uuid = { version = "1.23.3", features = ["v4"] }
+15
View File
@@ -0,0 +1,15 @@
use argon2::password_hash::{SaltString, rand_core::OsRng};
use argon2::{Argon2, PasswordHash, PasswordHasher, PasswordVerifier};
pub fn hash_password(password: &str) -> Result<String, argon2::password_hash::Error> {
let salt = SaltString::generate(&mut OsRng);
let hash = Argon2::default().hash_password(password.as_bytes(), &salt)?;
Ok(hash.to_string())
}
pub fn verify_password(password: &str, stored_hash: &str) -> bool {
let parsed = PasswordHash::new(stored_hash).unwrap();
Argon2::default()
.verify_password(password.as_bytes(), &parsed)
.is_ok()
}
+179 -40
View File
@@ -1,30 +1,30 @@
use std::{
hash::{DefaultHasher, Hash, Hasher},
net::SocketAddr,
sync::Arc,
};
use std::{net::SocketAddr, sync::Arc};
use axum::{
Json, Router,
body::Bytes,
extract::{Form, Query, State},
http::StatusCode,
body::{Body, Bytes},
extract::{Form, FromRequestParts, Path, Query, State},
http::{HeaderMap, StatusCode, header::SET_COOKIE, request::Parts},
response::{IntoResponse, Response},
routing::{get, post},
};
use entity::{auth_session, stream_key, users};
use sea_orm::DatabaseConnection;
use serde::{Deserialize, Serialize};
use tokio::{
net::TcpListener,
sync::{Mutex, mpsc::Sender},
};
use uuid::Uuid;
use crate::AppState;
use crate::{AppState, hash::hash_password};
pub struct HttpServer {
pub offer_tx: Sender<(i32, String, String)>,
pub accept_rx: async_broadcast::Receiver<(i32, String)>,
pub offer_tx: Sender<(i32, i32, String)>,
pub accept_rx: async_broadcast::Receiver<(i32, Option<String>)>,
pub appstate: Arc<Mutex<AppState>>,
pub request_count: Mutex<i32>,
pub db: DatabaseConnection,
}
impl HttpServer {
@@ -35,7 +35,12 @@ impl HttpServer {
let app = Router::new()
.route("/api/catalog", get(catalog_handler))
.route("/api/user", post(create_user_handler))
.route("/api/stream", get(stream_handler))
.route(
"/api/stream-key",
post(create_stream_key_handler).get(get_all_stream_keys),
)
.route("/api/login", post(login_handler))
.route("/api/stream/{slug}", post(stream_handler))
.route("/api/meow", get(meow_handler))
.with_state(state);
@@ -47,17 +52,103 @@ impl HttpServer {
}
}
fn calculate_hash<T: Hash>(t: &T) -> u64 {
let mut s = DefaultHasher::new();
t.hash(&mut s);
s.finish()
}
async fn catalog_handler(State(state): State<Arc<HttpServer>>) -> impl IntoResponse {
let catalog = catalog_from_state(&state.appstate).await;
Json(catalog)
}
#[derive(Deserialize)]
struct CreateStreamKeyBody {
label: String,
}
struct AuthUser(entity::users::Model);
impl FromRequestParts<Arc<HttpServer>> for AuthUser {
type Rejection = StatusCode;
async fn from_request_parts(
parts: &mut Parts,
state: &Arc<HttpServer>,
) -> Result<Self, StatusCode> {
let token = parts
.headers
.get("session")
// parse token...
.ok_or(StatusCode::UNAUTHORIZED)?;
let user =
users::Entity::find_by_auth_session(&state.db, token.to_str().unwrap().to_string())
.await
.map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?
.ok_or(StatusCode::UNAUTHORIZED)?;
Ok(AuthUser(user))
}
}
async fn create_stream_key_handler(
State(state): State<Arc<HttpServer>>,
auth: AuthUser,
Json(payload): Json<CreateStreamKeyBody>,
) -> impl IntoResponse {
let uuid = uuid::Uuid::new_v4();
let value = format!("stream-key-{uuid}");
let key = stream_key::Entity::create(&state.db, auth.0.id, value, payload.label, false).await;
if let Ok(_key) = key {
StatusCode::CREATED
} else {
StatusCode::INTERNAL_SERVER_ERROR
}
}
struct StreamKeys {
keys: Vec<StreamKey>,
}
struct StreamKey {
id: i32,
label: String,
value: String,
}
async fn get_all_stream_keys(
State(state): State<Arc<HttpServer>>,
auth: AuthUser,
) -> impl IntoResponse {
let user_keys = stream_key::Entity::find_by_user(&state.db, auth.0.id)
.await
.unwrap();
match stream_key::Entity::find_by_user(&state.db, auth.0.id).await {
Ok(keys) => Json(keys).into_response(),
Err(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
}
}
#[derive(Deserialize)]
struct LoginForm {
username: String,
password: String,
}
#[derive(Serialize)]
struct LoginResponse {
session_token: String,
}
async fn login_handler(
State(state): State<Arc<HttpServer>>,
Json(payload): Json<LoginForm>,
) -> StatusCode {
// TODO: hash payload.password with a real password hasher (argon2/bcrypt), look up user by username
// TODO: verify hash matches stored hashed_password
// TODO: call auth_session::Entity::create, return session token
// TODO: return 401 on bad credentials
todo!("login not yet implemented")
}
#[derive(Deserialize)]
struct CreateUserForm {
username: String,
@@ -65,44 +156,91 @@ struct CreateUserForm {
ref_token: String,
}
async fn create_user_handler(Form(form): Form<CreateUserForm>) -> StatusCode {
// TODO: ref_token validation, db integration
let _ = (
form.username,
calculate_hash(&form.password),
form.ref_token,
);
todo!("user creation not yet implemented")
}
async fn create_user_handler(
State(state): State<Arc<HttpServer>>,
Json(payload): Json<CreateUserForm>,
) -> (HeaderMap, StatusCode) {
if payload.ref_token != "TEST" {
return (HeaderMap::new(), StatusCode::UNAUTHORIZED);
}
#[derive(Deserialize)]
struct StreamQuery {
stream_label: String,
let meow = users::Entity::create(
&state.db,
payload.username,
hash_password(&payload.password).unwrap(),
)
.await;
// Create session
let session: auth_session::Model = if let Ok(meow) = meow {
auth_session::Entity::create(&state.db, meow.id)
.await
.unwrap() // This should be ok (hopefully)
} else {
return (HeaderMap::new(), StatusCode::CONFLICT);
};
let token = session.value;
let mut headers = HeaderMap::new();
headers.insert(
SET_COOKIE,
format!("session={token}; HttpOnly; SameSite=Strict; Path=/")
.parse()
.unwrap(),
);
(headers, StatusCode::OK)
}
async fn stream_handler(
State(state): State<Arc<HttpServer>>,
Query(query): Query<StreamQuery>,
body: Bytes,
Path(slug): Path<String>,
body: String,
) -> impl IntoResponse {
let body_str = String::from_utf8(body.to_vec()).unwrap_or_default();
let mut request_id = state.request_count.lock().await;
*request_id += 1;
let request_id_clone = *request_id;
drop(request_id);
let stream_key_id = {
let app = state.appstate.lock().await;
app.stream_sessions
.iter()
.find(|e| e.value().stream_key_label == slug)
.map(|e| *e.key())
};
let stream_key_id = if let Some(id) = stream_key_id {
id
} else {
return Response::builder()
.status(StatusCode::NOT_FOUND)
.header("content-type", "application/text")
.body("".to_string())
.unwrap();
};
let mut accept_rx = state.accept_rx.new_receiver();
let _ = state
.offer_tx
.send((request_id_clone, query.stream_label, body_str))
.await;
.send((request_id_clone, stream_key_id, body))
.await
.unwrap();
let mut reply_body = String::new();
while let Ok(answer) = accept_rx.recv().await {
if answer.0 == request_id_clone {
reply_body = answer.1;
break;
if let Some(reply) = answer.1 {
if answer.0 == request_id_clone {
Response::builder()
.status(StatusCode::CREATED)
.header("content-type", "application/sdp")
.body(reply)
.unwrap();
}
} else {
Response::builder()
.status(StatusCode::NOT_FOUND)
// .header("content-type", "application/sdp")
.body("")
.unwrap();
}
}
@@ -119,9 +257,10 @@ async fn meow_handler() -> &'static str {
#[derive(Serialize)]
struct StreamCatalog {
active_streams: Vec<String>,
active_streams: Vec<i32>,
}
// TODO: Make this return streak labels and their id instead of the whole stream key omfg
async fn catalog_from_state(state: &Arc<Mutex<AppState>>) -> StreamCatalog {
let app = state.lock().await;
let active_streams = app
+51 -8
View File
@@ -2,10 +2,12 @@ use std::{error::Error, sync::Arc};
use async_broadcast::broadcast;
use dashmap::DashMap;
use migration::{Migrator, MigratorTrait};
use rml_rtmp::{
handshake::{Handshake, HandshakeProcessResult, PeerType},
sessions::{ServerSession, ServerSessionConfig, ServerSessionEvent, ServerSessionResult},
};
use sea_orm::{Database, sqlx::types::chrono};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::TcpListener,
@@ -17,33 +19,49 @@ use crate::{
media::{H264Parser, VideoFrame},
};
mod hash;
mod http;
mod media;
mod rtmp;
mod webrtc;
pub struct AppState {
pub stream_sessions: Arc<DashMap<String, StreamSession>>,
pub stream_sessions: Arc<DashMap<i32, StreamSession>>,
}
pub struct StreamSession {
pub stream_key: String,
pub stream_key_id: i32,
pub stream_key_label: String,
pub frame_channel: async_broadcast::Sender<Arc<VideoFrame>>,
}
#[tokio::main]
async fn main() -> Result<(), Box<dyn Error>> {
let listener = TcpListener::bind("0.0.0.0:8123").await?;
let db = Database::connect("sqlite://./db.sqlite?mode=rwc")
.await
.unwrap();
Migrator::up(&db, None).await.unwrap();
let appstate = Arc::new(Mutex::new(AppState {
stream_sessions: Arc::new(DashMap::new()),
}));
let (offer_tx, offer_rx) = tokio::sync::mpsc::channel::<(String, String)>(4);
let (answer_tx, answer_rx) = tokio::sync::mpsc::channel::<(String, String)>(4);
let (offer_tx, offer_rx) = tokio::sync::mpsc::channel::<(i32, i32, String)>(4);
// Request_Id,
// String_Label,
// Offer_body
// --
// We are using broadcast instead of a regular channel because the way answers are processed are
// mulithreaded and we could easily process the answer of a completely different offer request
let (answer_tx, answer_rx) = broadcast::<(i32, Option<String>)>(4);
// Request_Id,
// Answer_body
let mut http = HttpServer {
let http = HttpServer {
offer_tx,
accept_rx: answer_rx,
appstate: appstate.clone(),
request_count: Mutex::new(0),
db: db.clone(),
};
http.start()?;
@@ -52,6 +70,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
offer_rx,
accept_tx: answer_tx,
sessions_ref: app.stream_sessions.clone(),
db: db.clone(),
};
webrtc.start()?;
drop(app);
@@ -59,6 +78,7 @@ async fn main() -> Result<(), Box<dyn Error>> {
loop {
let (stream, _) = listener.accept().await?;
let appstate = appstate.clone();
let db = db.clone();
tokio::spawn(async move {
let mut stream = stream;
@@ -108,22 +128,45 @@ async fn main() -> Result<(), Box<dyn Error>> {
}
ServerSessionResult::RaisedEvent(x) => match x {
ServerSessionEvent::PublishStreamFinished { stream_key, .. } => {
appstate.lock().await.stream_sessions.remove(&stream_key);
let key = entity::stream_key::Entity::find_by_key(&db, &stream_key)
.await
.unwrap()
.unwrap();
appstate.lock().await.stream_sessions.remove(&key.id);
}
ServerSessionEvent::PublishStreamRequested {
request_id,
stream_key,
..
} => {
let key =
entity::stream_key::Entity::find_by_key(&db, &stream_key).await;
let key = if let Ok(Some(key)) = key {
key
} else {
// Early return, gives error to client
let reply = rtmp_session
.reject_request(request_id, "", "Stream key invalid")
.unwrap();
for x in reply {
if let ServerSessionResult::OutboundResponse(y) = x {
stream.write_all(&y.bytes).await.unwrap();
}
}
break;
};
let session = StreamSession {
stream_key: stream_key.clone(),
stream_key_id: key.id,
stream_key_label: key.label,
frame_channel: video_channel.clone(),
};
appstate
.lock()
.await
.stream_sessions
.insert(stream_key.clone(), session);
.insert(key.id, session);
let reply = rtmp_session.accept_request(request_id).unwrap();
for x in reply {
+37 -21
View File
@@ -1,13 +1,15 @@
use dashmap::DashMap;
use tokio::{
net::UdpSocket,
sync::mpsc::{Receiver, Sender},
};
use entity::stream_key;
use sea_orm::DatabaseConnection;
use std::{
error::Error,
sync::Arc,
time::{Duration, Instant},
};
use tokio::{
net::UdpSocket,
sync::mpsc::{Receiver, Sender},
};
use str0m::{
Candidate, Event, Input, Output, Rtc,
@@ -19,16 +21,28 @@ use str0m::{
use crate::{StreamSession, media::VideoFrame};
pub struct Webrtc {
pub offer_rx: Receiver<(String, String)>,
pub accept_tx: Sender<(String, String)>,
pub sessions_ref: Arc<DashMap<String, StreamSession>>,
pub offer_rx: Receiver<(i32, i32, String)>,
pub accept_tx: async_broadcast::Sender<(i32, Option<String>)>,
pub sessions_ref: Arc<DashMap<i32, StreamSession>>,
pub db: DatabaseConnection,
}
const PER_CLIENT_CONNECTION_BUF: usize = 65535;
impl Webrtc {
pub fn start(mut self) -> Result<(), Box<dyn Error>> {
tokio::spawn(async move {
while let Some(offer) = self.offer_rx.recv().await {
let (stream_key, sdp_body) = offer;
let (request_id, stream_id, sdp_body) = offer;
let stream_key =
stream_key::Entity::find_by_key(&self.db, &stream_id.to_string()).await;
if let Ok(None) = stream_key {
self.accept_tx.broadcast((request_id, None)).await.unwrap();
continue;
}
let socket = UdpSocket::bind("127.0.0.1:0").await.unwrap();
let local_addr = socket.local_addr().unwrap();
@@ -50,7 +64,7 @@ impl Webrtc {
let mid = changes.add_media(
MediaKind::Video,
str0m::media::Direction::SendOnly,
Some(stream_key.clone()),
Some(stream_id.to_string()),
Some("video0".to_string()),
None,
);
@@ -64,13 +78,13 @@ impl Webrtc {
let answer_sdp = offer_answer.to_sdp_string();
self.accept_tx
.send((stream_key.clone(), answer_sdp))
.broadcast((request_id, Some(answer_sdp)))
.await
.unwrap();
let sessions_ref = self.sessions_ref.clone();
tokio::spawn(async move {
Webrtc::detach_connection(socket, rtc, sessions_ref, mid).await;
Webrtc::detach_connection(socket, rtc, sessions_ref, stream_id, mid).await;
});
}
});
@@ -80,7 +94,8 @@ impl Webrtc {
async fn detach_connection(
socket: UdpSocket,
mut rtc: Rtc,
sessions_ref: Arc<DashMap<String, StreamSession>>,
sessions_ref: Arc<DashMap<i32, StreamSession>>,
stream_id: i32,
_hint_mid: Mid,
) {
let mut video_mid: Option<Mid> = None;
@@ -88,7 +103,7 @@ impl Webrtc {
let mut connected = false;
let mut video_stream: Option<async_broadcast::Receiver<Arc<VideoFrame>>> = None;
let mut recv_buf = vec![0u8; 65535];
let mut recv_buf = vec![0u8; PER_CLIENT_CONNECTION_BUF];
let local_addr = socket.local_addr().unwrap();
loop {
@@ -104,11 +119,9 @@ impl Webrtc {
Event::MediaAdded(ma) => {
if ma.kind == MediaKind::Video {
if let Some(writer) = rtc.writer(ma.mid) {
let best = writer
.payload_params()
.max_by_key(|p| {
p.spec().format.profile_level_id.unwrap_or(0)
});
let best = writer.payload_params().max_by_key(|p| {
p.spec().format.profile_level_id.unwrap_or(0)
});
if let Some(params) = best {
println!("Selected PT {:?}", params.pt());
video_pt = Some(params.pt());
@@ -132,7 +145,7 @@ impl Webrtc {
if connected {
if video_stream.is_none() {
if let Some(session) = sessions_ref.get("test") {
if let Some(session) = sessions_ref.get(&stream_id) {
video_stream = Some(session.frame_channel.new_receiver());
}
}
@@ -141,7 +154,8 @@ impl Webrtc {
match stream.try_recv() {
Ok(frame) => {
let now = Instant::now();
let rtp_time = MediaTime::from_90khz(frame.timestamp_ms as u64 * 90);
let rtp_time =
MediaTime::from_90khz(frame.timestamp_ms as u64 * 90);
if let (Some(pt), Some(writer)) =
(video_pt, video_mid.and_then(|m| rtc.writer(m)))
{
@@ -160,7 +174,9 @@ impl Webrtc {
}
}
let wait_until = deadline.min(Instant::now() + Duration::from_millis(20)).max(Instant::now());
let wait_until = deadline
.min(Instant::now() + Duration::from_millis(20))
.max(Instant::now());
let sleep = tokio::time::sleep_until(wait_until.into());
tokio::select! {