370 lines
14 KiB
C++
370 lines
14 KiB
C++
|
#include "StreamController.h"
|
||
|
|
#include "../services/DatabaseService.h"
|
||
|
|
#include "../services/StatsService.h"
|
||
|
|
#include "../services/RedisHelper.h"
|
||
|
|
#include "../services/OmeClient.h"
|
||
|
|
#include "../services/AuthService.h"
|
||
|
|
#include <drogon/utils/Utilities.h>
|
||
|
|
#include <drogon/Cookie.h>
|
||
|
|
#include <random>
|
||
|
|
#include <sstream>
|
||
|
|
#include <iomanip>
|
||
|
|
#include <chrono>
|
||
|
|
|
||
|
|
using namespace drogon::orm;
|
||
|
|
|
||
|
|
// Helper functions at the top
|
||
|
|
namespace {
|
||
|
|
// JSON response helper - saves 6-8 lines per endpoint
|
||
|
|
HttpResponsePtr jsonResp(const Json::Value& j, HttpStatusCode c = k200OK) {
|
||
|
|
auto r = HttpResponse::newHttpJsonResponse(j);
|
||
|
|
r->setStatusCode(c);
|
||
|
|
return r;
|
||
|
|
}
|
||
|
|
|
||
|
|
HttpResponsePtr jsonOk(const Json::Value& data) {
|
||
|
|
return jsonResp(data);
|
||
|
|
}
|
||
|
|
|
||
|
|
HttpResponsePtr jsonError(const std::string& error, HttpStatusCode code = k400BadRequest) {
|
||
|
|
Json::Value j;
|
||
|
|
j["success"] = false;
|
||
|
|
j["error"] = error;
|
||
|
|
return jsonResp(j, code);
|
||
|
|
}
|
||
|
|
|
||
|
|
// Quick JSON builder for common patterns
|
||
|
|
Json::Value json(std::initializer_list<std::pair<const char*, Json::Value>> items) {
|
||
|
|
Json::Value j;
|
||
|
|
for (const auto& [key, value] : items) {
|
||
|
|
j[key] = value;
|
||
|
|
}
|
||
|
|
return j;
|
||
|
|
}
|
||
|
|
|
||
|
|
UserInfo getUserFromRequest(const HttpRequestPtr &req) {
|
||
|
|
UserInfo user;
|
||
|
|
std::string auth = req->getHeader("Authorization");
|
||
|
|
|
||
|
|
if (auth.empty() || auth.substr(0, 7) != "Bearer ") {
|
||
|
|
return user;
|
||
|
|
}
|
||
|
|
|
||
|
|
std::string token = auth.substr(7);
|
||
|
|
AuthService::getInstance().validateToken(token, user);
|
||
|
|
return user;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Static member definitions
|
||
|
|
std::mutex StreamWebSocketController::connectionsMutex_;
|
||
|
|
std::unordered_map<std::string, std::unordered_set<WebSocketConnectionPtr>> StreamWebSocketController::tokenConnections_;
|
||
|
|
std::unordered_set<WebSocketConnectionPtr> StreamWebSocketController::connections_;
|
||
|
|
|
||
|
|
void StreamController::health(const HttpRequestPtr &,
|
||
|
|
std::function<void(const HttpResponsePtr &)> &&callback) {
|
||
|
|
callback(jsonOk(json({
|
||
|
|
{"status", "ok"},
|
||
|
|
{"timestamp", Json::Int64(std::chrono::duration_cast<std::chrono::milliseconds>(
|
||
|
|
std::chrono::system_clock::now().time_since_epoch()
|
||
|
|
).count())}
|
||
|
|
})));
|
||
|
|
}
|
||
|
|
|
||
|
|
void StreamController::validateStreamKey(const HttpRequestPtr &,
|
||
|
|
std::function<void(const HttpResponsePtr &)> &&callback,
|
||
|
|
const std::string &key) {
|
||
|
|
// Now validate against realms table
|
||
|
|
auto dbClient = app().getDbClient();
|
||
|
|
*dbClient << "SELECT 1 FROM realms WHERE stream_key = $1 AND is_active = true"
|
||
|
|
<< key
|
||
|
|
>> [callback, key](const Result &r) {
|
||
|
|
bool valid = !r.empty();
|
||
|
|
if (valid) {
|
||
|
|
// Store in Redis
|
||
|
|
RedisHelper::storeKey("stream_key:" + key, "1", 86400);
|
||
|
|
}
|
||
|
|
callback(jsonOk(json({{"valid", valid}})));
|
||
|
|
}
|
||
|
|
>> [callback](const DrogonDbException &e) {
|
||
|
|
LOG_ERROR << "Database error: " << e.base().what();
|
||
|
|
callback(jsonOk(json({{"valid", false}})));
|
||
|
|
};
|
||
|
|
}
|
||
|
|
|
||
|
|
void StreamController::disconnectStream(const HttpRequestPtr &req,
|
||
|
|
std::function<void(const HttpResponsePtr &)> &&callback,
|
||
|
|
const std::string &streamKey) {
|
||
|
|
UserInfo user = getUserFromRequest(req);
|
||
|
|
if (user.id == 0) {
|
||
|
|
callback(jsonError("Unauthorized", k401Unauthorized));
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
|
||
|
|
// Check if user owns this stream or is admin
|
||
|
|
auto dbClient = app().getDbClient();
|
||
|
|
*dbClient << "SELECT user_id FROM realms WHERE stream_key = $1 AND is_active = true"
|
||
|
|
<< streamKey
|
||
|
|
>> [user, callback, streamKey](const Result& r) {
|
||
|
|
if (r.empty() || (r[0]["user_id"].as<int64_t>() != user.id && !user.isAdmin)) {
|
||
|
|
callback(jsonError("Forbidden", k403Forbidden));
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
|
||
|
|
OmeClient::getInstance().disconnectStream(streamKey, [callback](bool success) {
|
||
|
|
if (success) {
|
||
|
|
callback(jsonOk(json({
|
||
|
|
{"success", true},
|
||
|
|
{"message", "Stream disconnected"}
|
||
|
|
})));
|
||
|
|
} else {
|
||
|
|
callback(jsonError("Failed to disconnect stream"));
|
||
|
|
}
|
||
|
|
});
|
||
|
|
}
|
||
|
|
>> [callback](const DrogonDbException& e) {
|
||
|
|
LOG_ERROR << "Database error: " << e.base().what();
|
||
|
|
callback(jsonError("Database error"));
|
||
|
|
};
|
||
|
|
}
|
||
|
|
|
||
|
|
void StreamController::getStreamStats(const HttpRequestPtr &,
|
||
|
|
std::function<void(const HttpResponsePtr &)> &&callback,
|
||
|
|
const std::string &streamKey) {
|
||
|
|
StatsService::getInstance().getStreamStats(streamKey,
|
||
|
|
[callback](bool success, const StreamStats& stats) {
|
||
|
|
if (success) {
|
||
|
|
Json::Value json;
|
||
|
|
json["success"] = true;
|
||
|
|
|
||
|
|
auto& s = json["stats"];
|
||
|
|
s["connections"] = static_cast<Json::Int64>(stats.uniqueViewers);
|
||
|
|
s["total_connections"] = static_cast<Json::Int64>(stats.totalConnections);
|
||
|
|
s["bytes_in"] = static_cast<Json::Int64>(stats.totalBytesIn);
|
||
|
|
s["bytes_out"] = static_cast<Json::Int64>(stats.totalBytesOut);
|
||
|
|
s["bitrate"] = stats.bitrate;
|
||
|
|
s["codec"] = stats.codec;
|
||
|
|
s["resolution"] = stats.resolution;
|
||
|
|
s["fps"] = stats.fps;
|
||
|
|
s["is_live"] = stats.isLive;
|
||
|
|
|
||
|
|
if (stats.totalBytesIn > 0) {
|
||
|
|
s["data_rate_in"] = stats.bitrate / 1000.0;
|
||
|
|
}
|
||
|
|
if (stats.totalBytesOut > 0) {
|
||
|
|
s["data_rate_out"] = stats.totalBytesOut / 1024.0 / 1024.0;
|
||
|
|
}
|
||
|
|
|
||
|
|
// Protocol breakdown
|
||
|
|
auto& pc = s["protocol_connections"];
|
||
|
|
pc["webrtc"] = static_cast<Json::Int64>(stats.protocolConnections.webrtc);
|
||
|
|
pc["hls"] = static_cast<Json::Int64>(stats.protocolConnections.hls);
|
||
|
|
pc["llhls"] = static_cast<Json::Int64>(stats.protocolConnections.llhls);
|
||
|
|
pc["dash"] = static_cast<Json::Int64>(stats.protocolConnections.dash);
|
||
|
|
|
||
|
|
callback(jsonResp(json));
|
||
|
|
} else {
|
||
|
|
callback(jsonError("Failed to retrieve stream stats"));
|
||
|
|
}
|
||
|
|
});
|
||
|
|
}
|
||
|
|
|
||
|
|
void StreamController::getActiveStreams(const HttpRequestPtr &,
|
||
|
|
std::function<void(const HttpResponsePtr &)> &&callback) {
|
||
|
|
OmeClient::getInstance().getActiveStreams([callback](bool success, const Json::Value& omeResponse) {
|
||
|
|
if (success) {
|
||
|
|
LOG_INFO << "Active streams: " << omeResponse["response"].toStyledString();
|
||
|
|
callback(jsonOk(json({
|
||
|
|
{"success", true},
|
||
|
|
{"streams", omeResponse["response"]}
|
||
|
|
})));
|
||
|
|
} else {
|
||
|
|
callback(jsonError("Failed to get active streams from OME"));
|
||
|
|
}
|
||
|
|
});
|
||
|
|
}
|
||
|
|
|
||
|
|
void StreamController::issueViewerToken(const HttpRequestPtr &,
|
||
|
|
std::function<void(const HttpResponsePtr &)> &&callback,
|
||
|
|
const std::string &streamKey) {
|
||
|
|
// Validate against realms
|
||
|
|
auto dbClient = app().getDbClient();
|
||
|
|
*dbClient << "SELECT 1 FROM realms WHERE stream_key = $1 AND is_active = true"
|
||
|
|
<< streamKey
|
||
|
|
>> [callback, streamKey](const Result& r) {
|
||
|
|
if (r.empty()) {
|
||
|
|
callback(jsonResp({}, k404NotFound));
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
|
||
|
|
auto bytes = drogon::utils::genRandomString(32);
|
||
|
|
std::string token = drogon::utils::base64Encode(
|
||
|
|
(const unsigned char*)bytes.data(), bytes.length()
|
||
|
|
);
|
||
|
|
|
||
|
|
RedisHelper::storeKeyAsync("viewer_token:" + token, streamKey, 30,
|
||
|
|
[callback, token](bool stored) {
|
||
|
|
if (!stored) {
|
||
|
|
callback(jsonResp({}, k500InternalServerError));
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
|
||
|
|
auto resp = HttpResponse::newHttpResponse();
|
||
|
|
|
||
|
|
Cookie cookie("viewer_token", token);
|
||
|
|
cookie.setPath("/");
|
||
|
|
cookie.setHttpOnly(true);
|
||
|
|
cookie.setSecure(false);
|
||
|
|
cookie.setMaxAge(300);
|
||
|
|
resp->addCookie(cookie);
|
||
|
|
|
||
|
|
Json::Value body;
|
||
|
|
body["success"] = true;
|
||
|
|
body["viewer_token"] = token;
|
||
|
|
body["expires_in"] = 30;
|
||
|
|
|
||
|
|
resp->setContentTypeCode(CT_APPLICATION_JSON);
|
||
|
|
resp->setBody(Json::FastWriter().write(body));
|
||
|
|
|
||
|
|
callback(resp);
|
||
|
|
}
|
||
|
|
);
|
||
|
|
}
|
||
|
|
>> [callback](const DrogonDbException& e) {
|
||
|
|
LOG_ERROR << "Database error: " << e.base().what();
|
||
|
|
callback(jsonResp({}, k500InternalServerError));
|
||
|
|
};
|
||
|
|
}
|
||
|
|
|
||
|
|
void StreamController::heartbeat(const HttpRequestPtr &req,
|
||
|
|
std::function<void(const HttpResponsePtr &)> &&callback,
|
||
|
|
const std::string &streamKey) {
|
||
|
|
auto token = req->getCookie("viewer_token");
|
||
|
|
if (token.empty()) {
|
||
|
|
callback(jsonResp({}, k403Forbidden));
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
|
||
|
|
RedisHelper::getKeyAsync("viewer_token:" + token,
|
||
|
|
[callback, streamKey, token](const std::string& storedStreamKey) {
|
||
|
|
if (storedStreamKey != streamKey) {
|
||
|
|
callback(jsonResp({}, k403Forbidden));
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
|
||
|
|
services::RedisHelper::instance().expireAsync("viewer_token:" + token, 30,
|
||
|
|
[callback](bool success) {
|
||
|
|
if (!success) {
|
||
|
|
callback(jsonResp({}, k500InternalServerError));
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
|
||
|
|
callback(jsonOk(json({
|
||
|
|
{"success", true},
|
||
|
|
{"renewed", true}
|
||
|
|
})));
|
||
|
|
}
|
||
|
|
);
|
||
|
|
}
|
||
|
|
);
|
||
|
|
}
|
||
|
|
|
||
|
|
// WebSocket implementation
|
||
|
|
void StreamWebSocketController::handleNewMessage(const WebSocketConnectionPtr&,
|
||
|
|
std::string &&message,
|
||
|
|
const WebSocketMessageType &type) {
|
||
|
|
if (type == WebSocketMessageType::Text) {
|
||
|
|
Json::Value msg;
|
||
|
|
Json::Reader reader;
|
||
|
|
if (reader.parse(message, msg) && msg["type"].asString() == "subscribe") {
|
||
|
|
LOG_INFO << "Client subscribed to stream updates";
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
void StreamWebSocketController::handleNewConnection(const HttpRequestPtr &req,
|
||
|
|
const WebSocketConnectionPtr& wsConnPtr) {
|
||
|
|
LOG_INFO << "New WebSocket connection established";
|
||
|
|
|
||
|
|
auto token = req->getCookie("viewer_token");
|
||
|
|
if (token.empty()) {
|
||
|
|
LOG_WARN << "WebSocket connection without viewer token";
|
||
|
|
wsConnPtr->shutdown();
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
|
||
|
|
RedisHelper::getKeyAsync("viewer_token:" + token,
|
||
|
|
[wsConnPtr, token](const std::string& streamKey) {
|
||
|
|
if (streamKey.empty()) {
|
||
|
|
LOG_WARN << "Invalid viewer token";
|
||
|
|
wsConnPtr->shutdown();
|
||
|
|
return;
|
||
|
|
}
|
||
|
|
|
||
|
|
std::lock_guard<std::mutex> lock(connectionsMutex_);
|
||
|
|
tokenConnections_[token].insert(wsConnPtr);
|
||
|
|
connections_.insert(wsConnPtr);
|
||
|
|
|
||
|
|
LOG_INFO << "WebSocket authenticated for stream: " << streamKey;
|
||
|
|
}
|
||
|
|
);
|
||
|
|
}
|
||
|
|
|
||
|
|
void StreamWebSocketController::handleConnectionClosed(const WebSocketConnectionPtr& wsConnPtr) {
|
||
|
|
LOG_INFO << "WebSocket connection closed";
|
||
|
|
|
||
|
|
std::lock_guard<std::mutex> lock(connectionsMutex_);
|
||
|
|
|
||
|
|
std::string tokenToDelete;
|
||
|
|
for (auto& [token, conns] : tokenConnections_) {
|
||
|
|
if (conns.erase(wsConnPtr)) {
|
||
|
|
if (conns.empty()) {
|
||
|
|
tokenToDelete = token;
|
||
|
|
}
|
||
|
|
break;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
connections_.erase(wsConnPtr);
|
||
|
|
|
||
|
|
if (!tokenToDelete.empty()) {
|
||
|
|
tokenConnections_.erase(tokenToDelete);
|
||
|
|
|
||
|
|
RedisHelper::deleteKeyAsync("viewer_token:" + tokenToDelete,
|
||
|
|
[tokenToDelete](bool success) {
|
||
|
|
if (success) {
|
||
|
|
LOG_INFO << "Deleted viewer token on disconnect: " << tokenToDelete;
|
||
|
|
} else {
|
||
|
|
LOG_WARN << "Failed to delete viewer token: " << tokenToDelete;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
void StreamWebSocketController::broadcastKeyUpdate(const std::string& userId, const std::string& newKey) {
|
||
|
|
Json::Value msg;
|
||
|
|
msg["type"] = "key_regenerated";
|
||
|
|
msg["user_id"] = userId;
|
||
|
|
msg["stream_key"] = newKey;
|
||
|
|
|
||
|
|
auto msgStr = Json::FastWriter().write(msg);
|
||
|
|
|
||
|
|
std::lock_guard<std::mutex> lock(connectionsMutex_);
|
||
|
|
for (const auto& conn : connections_) {
|
||
|
|
if (conn->connected()) {
|
||
|
|
conn->send(msgStr);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
void StreamWebSocketController::broadcastStatsUpdate(const Json::Value& stats) {
|
||
|
|
std::string jsonStr = Json::FastWriter().write(stats);
|
||
|
|
|
||
|
|
std::lock_guard<std::mutex> lock(connectionsMutex_);
|
||
|
|
for (const auto& conn : connections_) {
|
||
|
|
if (conn->connected()) {
|
||
|
|
conn->send(jsonStr);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|