Feat: webrtc.

SighingSnow · Nov 21, 2021 · ca458e9 · ca458e9
1 parent ca268c1
commit ca458e9
Show file tree

Hide file tree

Showing 8 changed files with 145 additions and 51 deletions.
diff --git a/README.md b/README.md
@@ -19,15 +19,15 @@ Lots of work still need to be done.
 
 - [x] Record OpenGL app screen and encode to H.264
 - [x] Rtmp Streamer, which can push the H.264 raw frame in buffer to server
-- [ ] Webrtc
+- [x] Webrtc
 - [ ] HEVC
-- [ ] Multi-thread
+- [x] Multi-thread
 - [ ] Parallel 
 - [ ] Integrate with Irrlicht Game Engine.
 
 ### 3 Next Steps
 2021.11.8
-- [ ] Use librtc(libdatachannel is the origin name) to send memory video.
+- [x] Use librtc(libdatachannel is the origin name) to send memory video.
 - [ ] Implement with HEVC
 
 **Please note: the repo's license is MIT, but the 3rd_party/librtc is GPL.**

diff --git a/src/encoder.cpp b/src/encoder.cpp
@@ -148,12 +148,15 @@ void Encoder::GenOnePkt(uint8_t* buffer,uint8_t** ret_buf,int& ret_buf_size)
 
 void Encoder::DumpLocalVideo()
 {
+    static int count = 2;
     pkt->stream_index = stream->index;
     int er = av_write_frame(ofctx,pkt);
     if(er < 0){
         std::cout<<"write frame error"<<std::endl;
     }
+    if(count == 0) fclose(fout);
     fwrite(pkt->data,1,pkt->size,fout);
+    //count--;
 }
 
 void Encoder::FlushEncoder(int streamIndex)

diff --git a/src/h264fileparser.cpp b/src/h264fileparser.cpp
@@ -29,7 +29,9 @@
 
 using namespace std;
 
-H264FileParser::H264FileParser(uint32_t fps, bool loop):sampleDuration_us(1000*1000/fps),StreamSource() {}
+H264FileParser::H264FileParser(uint32_t fps, bool loop):sampleDuration_us(1000*1000/fps),StreamSource() {
+    sample_prepared = false;
+}
 
 void H264FileParser::start(){
     sampleTime_us = std::numeric_limits<uint64_t>::max() - sampleDuration_us + 1;
@@ -39,39 +41,17 @@ void H264FileParser::stop(){
     StreamSource::stop();
 }
 
-void H264FileParser::loadNextSample(std::vector<std::byte>& buffer) {
+void H264FileParser::loadNextSample(std::vector<uint8_t>& buffer) {
     unsigned long long i = 0;
     sampleTime_us += sampleDuration_us;
-    int flag = 0;
-    sample = buffer;
-//    while (i < sample.size()) {
-//        assert(i + 4 < sample.size());
-//        auto lengthPtr = (uint32_t *) (sample.data() + i);
-//        uint32_t length = ntohl(*lengthPtr);
-//        auto naluStartIndex = i + 4;
-//        auto naluEndIndex = naluStartIndex + length;
-//        if(naluEndIndex > sample.size()){
-//            ofstream ofs("xxx.h264");
-//            std::vector<uint8_t> aloha = *reinterpret_cast<std::vector<uint8_t> *>(&buffer);
-//            for(int i = 0 ; i < buffer.size();i++)
-//                ofs<<aloha[i];
-//        }
-//        assert(naluEndIndex <= sample.size());
-//        auto header = reinterpret_cast<rtc::NalUnitHeader *>(sample.data() + naluStartIndex);
-//        auto type = header->unitType();
-//        switch (type) {
-//            case 7:
-//                previousUnitType7 = {sample.begin() + i, sample.begin() + naluEndIndex};
-//                break;
-//            case 8:
-//                previousUnitType8 = {sample.begin() + i, sample.begin() + naluEndIndex};;
-//                break;
-//            case 5:
-//                previousUnitType5 = {sample.begin() + i, sample.begin() + naluEndIndex};;
-//                break;
-//        }
-//        i = naluEndIndex;
-//    }
+    sample = *reinterpret_cast<std::vector<std::byte>*>(&buffer);
+}
+
+void H264FileParser::loadNalu(std::vector<uint8_t> &buffer, int begin, int end)
+{
+    std::vector<uint8_t> buf(buffer.begin()+begin,buffer.begin()+end);
+    auto tmp = *reinterpret_cast<std::vector<std::byte> *>(&buf);
+    sample.insert(sample.end(),tmp.begin(),tmp.end());
 }
 
 vector<byte> H264FileParser::initialNALUS() {
@@ -90,3 +70,73 @@ vector<byte> H264FileParser::initialNALUS() {
     }
     return units;
 }
+
+void H264FileParser::clearSentSample() {
+    sample.clear();
+    sample.shrink_to_fit();
+}
+
+NALU_t H264FileParser::getNalu(std::vector<uint8_t>& buffer,int begin)
+{
+    NALU_t nalu;
+    bool info2,info3;
+    int pos = begin;
+    bool nextNaluFound = false;
+    info2 = findStartCode2(buffer,begin);
+    if(!info2){
+        info3 = findStartCode3(buffer,begin);
+        if(info3){
+            nalu.startcodeprefix_len = 4;
+            pos = begin+4;
+        }
+        else{
+            exit(1);
+        }
+    }
+    else{
+        nalu.startcodeprefix_len = 3;
+        pos = begin+3;
+    }
+    info3 = false;
+    info2 = false;
+    while(!nextNaluFound){
+        pos++;
+        if(pos == buffer.size()){
+            nalu.len = (pos - begin) - nalu.startcodeprefix_len;
+            nalu.buf.insert(nalu.buf.end(),buffer.begin()+begin+nalu.startcodeprefix_len,buffer.end());
+            assert(nalu.buf.size() == nalu.len);
+            nalu.forbidden_bit = nalu.buf[0] & 0x80; //1 bit
+            nalu.nal_reference_idc = nalu.buf[0] & 0x60; // 2 bit
+            nalu.nal_unit_type = (nalu.buf[0]) & 0x1f;// 5 bit
+            return nalu;
+        }
+        info3 = findStartCode3(buffer,pos-4);
+        if(!info3){
+            info2 = findStartCode2(buffer,pos-3);
+        }
+        nextNaluFound = info3 || info2;
+    }
+    int rewind = info3 ? -4:-3;
+    nalu.len = (pos - begin + rewind) - nalu.startcodeprefix_len;
+    nalu.buf.insert(nalu.buf.end(),buffer.begin()+begin+nalu.startcodeprefix_len,buffer.begin()+begin+nalu.startcodeprefix_len+nalu.len);
+    assert(nalu.len == nalu.buf.size());
+    nalu.forbidden_bit = nalu.buf[0] & 0x80;
+    nalu.nal_reference_idc = nalu.buf[0] & 0x60;
+    nalu.nal_unit_type = nalu.buf[0] & 0x1f;
+    return nalu;
+}
+
+bool H264FileParser::findStartCode2(std::vector<uint8_t>& buf,int begin)
+{
+    if(buf[begin] == 0x00 && buf[begin+1] == 0x00 && buf[begin+2] == 0x01){
+        return true;
+    }
+    return false;
+}
+bool H264FileParser::findStartCode3(std::vector<uint8_t>& buf,int begin)
+{
+    if(buf[begin] == 0x00 && buf[begin+1] == 0x00 && buf[begin+2] == 0x00 && buf[begin+3] == 0x01){
+        return true;
+    }
+    return false;
+}
diff --git a/src/h264fileparser.hpp b/src/h264fileparser.hpp
@@ -23,18 +23,36 @@
 
 #include "stream.hpp"
 
+typedef struct
+{
+    int startcodeprefix_len;      //! 4 for parameter sets and first slice in picture, 3 for everything else (suggested)
+    unsigned len;                 //! Length of the NAL unit (Excluding the start code, which does not belong to the NALU)
+    unsigned max_size;            //! Nal Unit Buffer size
+    int forbidden_bit;            //! should be always FALSE
+    int nal_reference_idc;        //! NALU_PRIORITY_xxxx
+    int nal_unit_type;            //! NALU_TYPE_xxxx
+    std::vector<uint8_t> buf;     //! contains the first byte followed by the EBSP
+} NALU_t;
+
 class H264FileParser : public StreamSource{
     std::optional<std::vector<std::byte>> previousUnitType5 = std::nullopt;
     std::optional<std::vector<std::byte>> previousUnitType7 = std::nullopt;
     std::optional<std::vector<std::byte>> previousUnitType8 = std::nullopt;
+    std::vector<std::uint8_t> preparing_sample;
+    bool findStartCode2(std::vector<uint8_t>& buf,int begin);
+    bool findStartCode3(std::vector<uint8_t>& buf,int begin);
 
 public:
     void start() override;
     void stop() override;
     const uint64_t sampleDuration_us;
     H264FileParser(uint32_t fps, bool loop);
-    void loadNextSample(std::vector<std::byte>& buffer);
+    void loadNextSample(std::vector<uint8_t>& buffer);
+    void clearSentSample();
+    void loadNalu(std::vector<uint8_t>& buffer,int begin,int end);
     std::vector<std::byte> initialNALUS();
+    bool sample_prepared;
+    NALU_t getNalu(std::vector<uint8_t>& buffer,int begin);
 };
 
 #endif /* h264fileparser_hpp */
diff --git a/src/rtc_publisher.cpp b/src/rtc_publisher.cpp
@@ -48,19 +48,28 @@ void RtcPublisher::setUp()
 
 void RtcPublisher::publish(uint8_t *buf, int size) {
     //TODO: current buf is raw, no pps and sps data
-    if(avStream.has_value() == false)
+    int index = 0;
+    if(avStream.has_value() == false || !avStream.value()->isRunning)
         return;
     std::vector<uint8_t> buffer(buf,buf+size*sizeof(uint8_t));
-    auto sample = *reinterpret_cast<std::vector<std::byte>*>(&buffer);
-    assert(sample.size() == size);
-    video->loadNextSample(sample);
-    if(avStream.value()->isRunning)
-        avStream.value()->publishSample();
+    video->loadNextSample(buffer);
+    avStream.value()->publishSample();
+    video->clearSentSample();
+//    while(index < buffer.size())
+//    {
+//        NALU_t nalu = video->getNalu(buffer,index);
+//        if(nalu.nal_unit_type == 5){
+//            avStream.value()->publishSample();
+//            video->clearSentSample();
+//        }
+//        video->loadNalu(buffer,index+nalu.startcodeprefix_len,index+nalu.startcodeprefix_len+nalu.len-1);
+//        index = index + nalu.len + nalu.startcodeprefix_len;
+//    }
 }
 
 shared_ptr<Client> RtcPublisher::createPeerConnection(const Configuration &config, weak_ptr<WebSocket> wws, string id)
 {
-    auto pc = make_shared<PeerConnection>(config);
+    pc = make_shared<PeerConnection>(config);
     auto client = make_shared<Client>(pc);
 
     pc->onStateChange([this,id](PeerConnection::State state) {
@@ -103,7 +112,7 @@ shared_ptr<Client> RtcPublisher::createPeerConnection(const Configuration &confi
         cout << "Video from " << id << " opened" << endl;
     });
 
-    auto dc = pc->createDataChannel("ping-pong");
+    dc = pc->createDataChannel("ping-pong");
     dc->onOpen([id, wdc = make_weak_ptr(dc)]() {
         if (auto dc = wdc.lock()) {
             dc->send("Ping");
@@ -162,7 +171,7 @@ shared_ptr<Stream> RtcPublisher::createStream(const unsigned int fps) {
                 if (rtpConfig->timestampToSeconds(reportElapsedTimestamp) > 1) {
                     trackData->sender->setNeedsToReport();
                 }
-                cout << "Sending " << streamType << " frame with size: " << to_string(sample.size()) << " to " << client << endl;
+                cout << "Sending " << streamType << " sample with size: " << to_string(sample.size()) << " to " << client << endl;
                 bool send = false;
                 try {
                     // send sample
@@ -266,7 +275,8 @@ shared_ptr<ClientTrackData> RtcPublisher::addVideo(const shared_ptr<PeerConnecti
     // create RTP configuration
     auto rtpConfig = make_shared<RtpPacketizationConfig>(ssrc, cname, payloadType, H264RtpPacketizer::defaultClockRate);
     // create packetizer
-    auto packetizer = make_shared<H264RtpPacketizer>(H264RtpPacketizer::Separator::LongStartSequence, rtpConfig);
+    //auto packetizer = make_shared<H264RtpPacketizer>(H264RtpPacketizer::Separator::LongStartSequence, rtpConfig);
+    auto packetizer =make_shared<H264RtpPacketizer>(H264RtpPacketizer::Separator::StartSequence,rtpConfig);
     // create H264 handler
     auto h264Handler = make_shared<H264PacketizationHandler>(packetizer);
     // add RTCP SR handler

diff --git a/src/rtc_publisher.h b/src/rtc_publisher.h
@@ -21,19 +21,28 @@ template <class T> weak_ptr<T> make_weak_ptr(shared_ptr<T> ptr) { return ptr; }
 class RtcPublisher {
 public:
     RtcPublisher(){
+        cache_index = 0;
         ws = make_shared<WebSocket>();
         rtcThread = new DispatchQueue("RtcThread") ;
         video = make_shared<H264FileParser>(25,true);
     }
     void setUp();
     void publish(uint8_t *buf, int size);
+    ~RtcPublisher(){
+        dc->close();
+        pc->close();
+        ws->close();
+    }
 protected:
     Configuration rtc_config;
     shared_ptr<WebSocket> ws;
     std::shared_ptr<H264FileParser> video;
     optional<shared_ptr<Stream>> avStream;
     DispatchQueue* rtcThread;
-
+    shared_ptr<DataChannel> dc;
+    shared_ptr<PeerConnection> pc;
+    int cache_index;
+    std::vector<uint8_t> buffer_cache;
     unordered_map<string, shared_ptr<Client>> clients{};
 
     shared_ptr<Client> createPeerConnection(const Configuration &config,

diff --git a/src/stream.cpp b/src/stream.cpp
@@ -95,13 +95,16 @@ void Stream::publishSample() {
     //std::lock_guard lock(mutex);
     if(!isRunning)
         return;
-    auto ssSST = unsafePrepareForSample();
-    auto ss = ssSST.first;
-    auto sst = ssSST.second;
+    auto ss = video;
+    auto sst = StreamSourceType::Video;
     auto sample = ss->getSample();
     sampleHandler(sst,ss->getSampleTime_us(),sample);
 }
 
+void Stream::publishSample2(int id,uint8_t* data,int size) {
+    rtcSendMessage(id, reinterpret_cast<const char *>(data), size);
+}
+
 void Stream::onSample(std::function<void (StreamSourceType, uint64_t, rtc::binary)> handler) {
     sampleHandler = handler;
 }

diff --git a/src/stream.hpp b/src/stream.hpp
@@ -31,7 +31,7 @@ class StreamSource {
     StreamSource() { }
     virtual void start() = 0;
     virtual void stop();
-    virtual void loadNextSample(std::vector<std::byte>& buffer) = 0;
+    virtual void loadNextSample(std::vector<uint8_t>& buffer) = 0;
 
     inline uint64_t getSampleTime_us() { return sampleTime_us; }
     inline rtc::binary getSample() { return sample; }
@@ -62,6 +62,7 @@ class Stream: std::enable_shared_from_this<Stream> {
     void sendSample();
 public:
     void publishSample();
+    void publishSample2(int id,uint8_t* data,int size);
     void onSample(std::function<void (StreamSourceType, uint64_t, rtc::binary)> handler);
     void start();
     void stop();