From 148f90d7df3784c1bdc374b00a98a9525c9caa1b Mon Sep 17 00:00:00 2001
From: amukhsimov <aka.mukhsimov@gmail.com>
Date: Mon, 22 Feb 2021 20:28:08 +0500
Subject: [PATCH] rtmp input added

---
 README.md       |  8 ++++++--
 headers/main.h  |  2 ++
 source/main.cpp | 42 +++++++++++++++++++++++++++++++++++++++---
 3 files changed, 47 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index a6f3e1b..6e4aefb 100644
--- a/README.md
+++ b/README.md
@@ -37,6 +37,10 @@ sudo apt update && sudo apt install -y tclsh pkg-config cmake libssl-dev build-e
 https://github.com/ALLATRA-IT/srt-server/releases
 
 ## Usage
-srt-server [PORT_RECEIVE PORT_SEND]
+srt-server [PORT_RECEIVE PORT_SEND [--rtmp]]
 
-*if PORT_RECEIVE and PORT_SEND are not specified, by default server receives on port 9000 and sends on 9001*
\ No newline at end of file
+*if PORT_RECEIVE and PORT_SEND are not specified, by default server receives on port 9000 and sends on 9001.* 
+
+*if **rtmp** is specified, rtmp stream is received on rtmp://0.0.0.0/rtmp/rtmp2srt*
+
+***ARGUMENTS ARE POSITION SENSITIVE***
\ No newline at end of file
diff --git a/headers/main.h b/headers/main.h
index e68d869..7a36cd6 100644
--- a/headers/main.h
+++ b/headers/main.h
@@ -79,6 +79,8 @@ uint64_t get_current_ms();
 
 SRTSOCKET create_starter_socket(string *service);
 
+void *begin_rtmp(void *opinfo);
+
 void *handle_data_transfer(void *opinfo);
 
 void *connections_handler(void *ptr);
diff --git a/source/main.cpp b/source/main.cpp
index 152b5f5..b568e05 100644
--- a/source/main.cpp
+++ b/source/main.cpp
@@ -6,18 +6,27 @@ int main(int argc, char *argv[]) {
     srt_startup();
     srt_setloglevel(srt_logging::LogLevel::fatal);
 
+    bool in_rtmp = false;
+    pthread_t rtmp_thread;
+
     int yes = 1, no = 0;
 
     service_rcv = string("9000");
     service_snd = string("9001");
 
-    if (argc == 3) {
+    if (argc >= 3) {
         if (strcmp(argv[1], argv[2]) == 0) {
             cout << "Argument 1 cannot be equal to argument 2" << endl;
             return 1;
         }
         service_rcv = argv[1];
         service_snd = argv[2];
+        if (argc > 3) {
+            if (strcmp(argv[3], "rtmp") == 0) {
+                cout << "RTMP input mode enabled" << endl;
+            }
+            in_rtmp = true;
+        }
     }
 
     // Open sockets
@@ -121,7 +130,7 @@ int main(int argc, char *argv[]) {
             .thread_info=&transferthread_infos[0]
     };
 
-    if (pthread_create(&transferthread_infos[0].thread, NULL, handle_data_transfer,
+    if (pthread_create(&transferthread_infos[0].thread, nullptr, handle_data_transfer,
                        (void *) (&transferthread_0_info)) != 0) {
         cout << "cannot create transfer thread 0: " << strerror(errno) << endl;
         return 7;
@@ -131,6 +140,17 @@ int main(int argc, char *argv[]) {
         return 8;
     }
 
+    if (in_rtmp) {
+        if (pthread_create(&rtmp_thread, nullptr, begin_rtmp, nullptr) != 0) {
+            cout << "cannot create rtmp thread" << strerror(errno) << endl;
+            return 7;
+        }
+        if (pthread_detach(rtmp_thread) != 0) {
+            cout << "cannot detach rtmp thread: " << strerror(errno) << endl;
+            return 8;
+        }
+    }
+
     transferthread_infos[0].is_alive = true;
 
     log(LOG_INFO, (char *) "Server was successfully initialized\n");
@@ -145,7 +165,7 @@ int main(int argc, char *argv[]) {
 
             pthread_mutex_lock(&stat_lock);
 
-            float timedelta = (float)(get_current_ms() - last_ms) / 1000;
+            float timedelta = (float) (get_current_ms() - last_ms) / 1000;
             total_rcv_bytes += temp_rcv_bytes;
             total_send_bytes += temp_send_bytes;
             print_stats(timedelta);
@@ -254,6 +274,22 @@ int main(int argc, char *argv[]) {
     return 0;
 }
 
+void *begin_rtmp(void *opinfo) {
+    char command[1000];
+    snprintf(command, 1000,
+             "ffmpeg -f flv -listen 1 -i rtmp://0.0.0.0:1935/rtmp/rtmp2srt "
+             "-c copy -f mpegts srt://127.0.0.1:%s?pkt_size=1316",
+             service_rcv.c_str());
+
+    while (true) {
+        cout << "Invoke cmd" << endl;
+        system(command);
+        cout << "RTMP error, restarting..." << endl;
+    }
+
+    return 0;
+}
+
 SRTSOCKET create_starter_socket(string *service) {
     addrinfo hints;
     addrinfo *res;