// Licensed to the Apache Software Foundation (ASF) under one // or more contributor license agreements. See the NOTICE file // distributed with this work for additional information // regarding copyright ownership. The ASF licenses this file // to you under the Apache License, Version 2.0 (the // "License"); you may not use this file except in compliance // with the License. You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, // software distributed under the License is distributed on an // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY // KIND, either express or implied. See the License for the // specific language governing permissions and limitations // under the License. // brpc - A framework to host and access services throughout Baidu. // Date: Fri May 20 15:52:22 CST 2016 #include #include #include #include #include #include #include #include "butil/time.h" #include "butil/macros.h" #include "brpc/socket.h" #include "brpc/acceptor.h" #include "brpc/server.h" #include "brpc/controller.h" #include "brpc/rtmp.h" #include "brpc/amf.h" namespace brpc { DECLARE_int32(amf_max_depth); DECLARE_int32(amf_max_string_size); DECLARE_int32(amf_max_array_size); } int main(int argc, char* argv[]) { testing::InitGoogleTest(&argc, argv); GFLAGS_NAMESPACE::ParseCommandLineFlags(&argc, &argv, true); return RUN_ALL_TESTS(); } namespace { class ScopedAMFMaxDepth { public: explicit ScopedAMFMaxDepth(int32_t depth) : _old_depth(brpc::FLAGS_amf_max_depth) { brpc::FLAGS_amf_max_depth = depth; } ~ScopedAMFMaxDepth() { brpc::FLAGS_amf_max_depth = _old_depth; } private: int32_t _old_depth; }; class ScopedAMFLimit { public: ScopedAMFLimit(int32_t* flag, int32_t value) : _flag(flag), _old_value(*flag) { *_flag = value; } ~ScopedAMFLimit() { *_flag = _old_value; } private: int32_t* _flag; int32_t _old_value; }; void AppendAMFStrictArrayHeader(std::string* out, uint32_t count) { out->push_back((char)brpc::AMF_MARKER_STRICT_ARRAY); out->push_back((char)((count >> 24) & 0xFF)); out->push_back((char)((count >> 16) & 0xFF)); out->push_back((char)((count >> 8) & 0xFF)); out->push_back((char)(count & 0xFF)); } void AppendAMFObjectHeader(std::string* out) { out->push_back((char)brpc::AMF_MARKER_OBJECT); } void AppendAMFEcmaArrayHeader(std::string* out, uint32_t count) { out->push_back((char)brpc::AMF_MARKER_ECMA_ARRAY); out->push_back((char)((count >> 24) & 0xFF)); out->push_back((char)((count >> 16) & 0xFF)); out->push_back((char)((count >> 8) & 0xFF)); out->push_back((char)(count & 0xFF)); } void AppendAMFShortStringBody(std::string* out, const char* name) { const uint16_t len = strlen(name); out->push_back((char)((len >> 8) & 0xFF)); out->push_back((char)(len & 0xFF)); out->append(name, len); } void AppendAMFLongStringHeader(std::string* out, uint32_t len) { out->push_back((char)brpc::AMF_MARKER_LONG_STRING); out->push_back((char)((len >> 24) & 0xFF)); out->push_back((char)((len >> 16) & 0xFF)); out->push_back((char)((len >> 8) & 0xFF)); out->push_back((char)(len & 0xFF)); } void AppendAMFObjectEnd(std::string* out) { AppendAMFShortStringBody(out, ""); out->push_back((char)brpc::AMF_MARKER_OBJECT_END); } std::string MakeNestedAMFObject(int depth) { std::string out; AppendAMFObjectHeader(&out); for (int i = 0; i < depth; ++i) { AppendAMFShortStringBody(&out, "x"); AppendAMFObjectHeader(&out); } for (int i = 0; i <= depth; ++i) { AppendAMFObjectEnd(&out); } return out; } std::string MakeNestedAMFEcmaArray(int depth) { std::string out; AppendAMFEcmaArrayHeader(&out, depth == 0 ? 0 : 1); for (int i = 0; i < depth; ++i) { AppendAMFShortStringBody(&out, "x"); AppendAMFEcmaArrayHeader(&out, i + 1 == depth ? 0 : 1); } return out; } } // namespace class TestRtmpClientStream : public brpc::RtmpClientStream { public: TestRtmpClientStream() : _called_on_stop(0) , _called_on_first_message(0) , _nvideomsg(0) , _naudiomsg(0) { LOG(INFO) << __FUNCTION__; } ~TestRtmpClientStream() { LOG(INFO) << __FUNCTION__; assertions_on_stop(); } void assertions_on_stop() { ASSERT_EQ(1, _called_on_stop); } void assertions_on_successful_play() { ASSERT_EQ(1, _called_on_first_message); ASSERT_LT(0, _nvideomsg); ASSERT_LT(0, _naudiomsg); } void assertions_on_failure() { ASSERT_EQ(0, _called_on_first_message); ASSERT_EQ(0, _nvideomsg); ASSERT_EQ(0, _naudiomsg); assertions_on_stop(); } void OnFirstMessage() { ++_called_on_first_message; } void OnStop() { ++_called_on_stop; } void OnVideoMessage(brpc::RtmpVideoMessage* msg) { ++_nvideomsg; // video data is ascii in UT, print it out. LOG(INFO) << remote_side() << "|stream=" << stream_id() << ": Got " << *msg << " data=" << msg->data; } void OnAudioMessage(brpc::RtmpAudioMessage* msg) { ++_naudiomsg; // audio data is ascii in UT, print it out. LOG(INFO) << remote_side() << "|stream=" << stream_id() << ": Got " << *msg << " data=" << msg->data; } private: int _called_on_stop; int _called_on_first_message; int _nvideomsg; int _naudiomsg; }; class TestRtmpRetryingClientStream : public brpc::RtmpRetryingClientStream { public: TestRtmpRetryingClientStream() : _called_on_stop(0) , _called_on_first_message(0) , _called_on_playable(0) { LOG(INFO) << __FUNCTION__; } ~TestRtmpRetryingClientStream() { LOG(INFO) << __FUNCTION__; assertions_on_stop(); } void assertions_on_stop() { ASSERT_EQ(1, _called_on_stop); } void OnStop() { ++_called_on_stop; } void OnFirstMessage() { ++_called_on_first_message; } void OnPlayable() { ++_called_on_playable; } void OnVideoMessage(brpc::RtmpVideoMessage* msg) { // video data is ascii in UT, print it out. LOG(INFO) << remote_side() << "|stream=" << stream_id() << ": Got " << *msg << " data=" << msg->data; } void OnAudioMessage(brpc::RtmpAudioMessage* msg) { // audio data is ascii in UT, print it out. LOG(INFO) << remote_side() << "|stream=" << stream_id() << ": Got " << *msg << " data=" << msg->data; } private: int _called_on_stop; int _called_on_first_message; int _called_on_playable; }; const char* UNEXIST_NAME = "unexist_stream"; class PlayingDummyStream : public brpc::RtmpServerStream { public: enum State { STATE_UNPLAYING, STATE_PLAYING, STATE_STOPPED }; PlayingDummyStream(int64_t sleep_ms) : _state(STATE_UNPLAYING), _sleep_ms(sleep_ms) { LOG(INFO) << __FUNCTION__ << "(" << this << ")"; } ~PlayingDummyStream() { LOG(INFO) << __FUNCTION__ << "(" << this << ")"; } void OnPlay(const brpc::RtmpPlayOptions& opt, butil::Status* status, google::protobuf::Closure* done) { brpc::ClosureGuard done_guard(done); LOG(INFO) << remote_side() << "|stream=" << stream_id() << ": Got play{stream_name=" << opt.stream_name << " start=" << opt.start << " duration=" << opt.duration << " reset=" << opt.reset << '}'; if (opt.stream_name == UNEXIST_NAME) { status->set_error(EPERM, "Unexist stream"); return; } if (_sleep_ms > 0) { LOG(INFO) << "Sleep " << _sleep_ms << " ms before responding play request"; bthread_usleep(_sleep_ms * 1000L); } int rc = bthread_start_background(&_play_thread, NULL, RunSendData, this); if (rc) { status->set_error(rc, "Fail to create thread"); return; } State expected = STATE_UNPLAYING; if (!_state.compare_exchange_strong(expected, STATE_PLAYING)) { if (expected == STATE_STOPPED) { bthread_stop(_play_thread); bthread_join(_play_thread, NULL); } else { CHECK(false) << "Impossible"; } } } void OnStop() { LOG(INFO) << "OnStop of PlayingDummyStream=" << this; if (_state.exchange(STATE_STOPPED) == STATE_PLAYING) { bthread_stop(_play_thread); bthread_join(_play_thread, NULL); } } void SendData(); private: static void* RunSendData(void* arg) { ((PlayingDummyStream*)arg)->SendData(); return NULL; } butil::atomic _state; bthread_t _play_thread; int64_t _sleep_ms; }; void PlayingDummyStream::SendData() { LOG(INFO) << "Enter SendData of PlayingDummyStream=" << this; brpc::RtmpVideoMessage vmsg; brpc::RtmpAudioMessage amsg; vmsg.timestamp = 1000; amsg.timestamp = 1000; for (int i = 0; !bthread_stopped(bthread_self()); ++i) { vmsg.timestamp += 20; amsg.timestamp += 20; vmsg.frame_type = brpc::FLV_VIDEO_FRAME_KEYFRAME; vmsg.codec = brpc::FLV_VIDEO_AVC; vmsg.data.clear(); vmsg.data.append(butil::string_printf("video_%d(ms_id=%u)", i, stream_id())); //failing to send is possible SendVideoMessage(vmsg); amsg.codec = brpc::FLV_AUDIO_AAC; amsg.rate = brpc::FLV_SOUND_RATE_44100HZ; amsg.bits = brpc::FLV_SOUND_16BIT; amsg.type = brpc::FLV_SOUND_STEREO; amsg.data.clear(); amsg.data.append(butil::string_printf("audio_%d(ms_id=%u)", i, stream_id())); SendAudioMessage(amsg); bthread_usleep(1000000); } LOG(INFO) << "Quit SendData of PlayingDummyStream=" << this; } class PlayingDummyService : public brpc::RtmpService { public: PlayingDummyService(int64_t sleep_ms = 0) : _sleep_ms(sleep_ms) {} private: // Called to create a server-side stream. virtual brpc::RtmpServerStream* NewStream( const brpc::RtmpConnectRequest&) { return new PlayingDummyStream(_sleep_ms); } int64_t _sleep_ms; }; class PublishStream : public brpc::RtmpServerStream { public: PublishStream(int64_t sleep_ms) : _sleep_ms(sleep_ms) , _called_on_stop(0) , _called_on_first_message(0) , _nvideomsg(0) , _naudiomsg(0) { LOG(INFO) << __FUNCTION__ << "(" << this << ")"; } ~PublishStream() { LOG(INFO) << __FUNCTION__ << "(" << this << ")"; assertions_on_stop(); } void assertions_on_stop() { ASSERT_EQ(1, _called_on_stop); } void OnPublish(const std::string& stream_name, brpc::RtmpPublishType publish_type, butil::Status* status, google::protobuf::Closure* done) { brpc::ClosureGuard done_guard(done); LOG(INFO) << remote_side() << "|stream=" << stream_id() << ": Got publish{stream_name=" << stream_name << " type=" << brpc::RtmpPublishType2Str(publish_type) << '}'; if (stream_name == UNEXIST_NAME) { status->set_error(EPERM, "Unexist stream"); return; } if (_sleep_ms > 0) { LOG(INFO) << "Sleep " << _sleep_ms << " ms before responding play request"; bthread_usleep(_sleep_ms * 1000L); } } void OnFirstMessage() { ++_called_on_first_message; } void OnStop() { LOG(INFO) << "OnStop of PublishStream=" << this; ++_called_on_stop; } void OnVideoMessage(brpc::RtmpVideoMessage* msg) { ++_nvideomsg; // video data is ascii in UT, print it out. LOG(INFO) << remote_side() << "|stream=" << stream_id() << ": Got " << *msg << " data=" << msg->data; } void OnAudioMessage(brpc::RtmpAudioMessage* msg) { ++_naudiomsg; // audio data is ascii in UT, print it out. LOG(INFO) << remote_side() << "|stream=" << stream_id() << ": Got " << *msg << " data=" << msg->data; } private: int64_t _sleep_ms; int _called_on_stop; int _called_on_first_message; int _nvideomsg; int _naudiomsg; }; class PublishService : public brpc::RtmpService { public: PublishService(int64_t sleep_ms = 0) : _sleep_ms(sleep_ms) { pthread_mutex_init(&_mutex, NULL); } ~PublishService() { pthread_mutex_destroy(&_mutex); } void move_created_streams( std::vector >* out) { out->clear(); BAIDU_SCOPED_LOCK(_mutex); out->swap(_created_streams); } private: // Called to create a server-side stream. virtual brpc::RtmpServerStream* NewStream( const brpc::RtmpConnectRequest&) { PublishStream* stream = new PublishStream(_sleep_ms); { BAIDU_SCOPED_LOCK(_mutex); _created_streams.push_back(stream); } return stream; } int64_t _sleep_ms; pthread_mutex_t _mutex; std::vector > _created_streams; }; class RtmpSubStream : public brpc::RtmpClientStream { public: explicit RtmpSubStream(brpc::RtmpMessageHandler* mh) : _message_handler(mh) {} // @RtmpStreamBase void OnMetaData(brpc::RtmpMetaData*, const butil::StringPiece&); void OnSharedObjectMessage(brpc::RtmpSharedObjectMessage* msg); void OnAudioMessage(brpc::RtmpAudioMessage* msg); void OnVideoMessage(brpc::RtmpVideoMessage* msg); void OnFirstMessage(); void OnStop(); private: std::unique_ptr _message_handler; }; void RtmpSubStream::OnFirstMessage() { _message_handler->OnPlayable(); } void RtmpSubStream::OnMetaData(brpc::RtmpMetaData* obj, const butil::StringPiece& name) { _message_handler->OnMetaData(obj, name); } void RtmpSubStream::OnSharedObjectMessage(brpc::RtmpSharedObjectMessage* msg) { _message_handler->OnSharedObjectMessage(msg); } void RtmpSubStream::OnAudioMessage(brpc::RtmpAudioMessage* msg) { _message_handler->OnAudioMessage(msg); } void RtmpSubStream::OnVideoMessage(brpc::RtmpVideoMessage* msg) { _message_handler->OnVideoMessage(msg); } void RtmpSubStream::OnStop() { _message_handler->OnSubStreamStop(this); } class RtmpSubStreamCreator : public brpc::SubStreamCreator { public: RtmpSubStreamCreator(const brpc::RtmpClient* client); ~RtmpSubStreamCreator(); // @SubStreamCreator void NewSubStream(brpc::RtmpMessageHandler* message_handler, butil::intrusive_ptr* sub_stream); void LaunchSubStream(brpc::RtmpStreamBase* sub_stream, brpc::RtmpRetryingClientStreamOptions* options); private: const brpc::RtmpClient* _client; }; RtmpSubStreamCreator::RtmpSubStreamCreator(const brpc::RtmpClient* client) : _client(client) {} RtmpSubStreamCreator::~RtmpSubStreamCreator() {} void RtmpSubStreamCreator::NewSubStream(brpc::RtmpMessageHandler* message_handler, butil::intrusive_ptr* sub_stream) { if (sub_stream) { (*sub_stream).reset(new RtmpSubStream(message_handler)); } return; } void RtmpSubStreamCreator::LaunchSubStream( brpc::RtmpStreamBase* sub_stream, brpc::RtmpRetryingClientStreamOptions* options) { brpc::RtmpClientStreamOptions client_options = *options; dynamic_cast(sub_stream)->Init(_client, client_options); } TEST(RtmpTest, parse_rtmp_url) { butil::StringPiece host; butil::StringPiece vhost; butil::StringPiece port; butil::StringPiece app; butil::StringPiece stream_name; brpc::ParseRtmpURL("rtmp://HOST/APP/STREAM", &host, &vhost, &port, &app, &stream_name); ASSERT_EQ("HOST", host); ASSERT_TRUE(vhost.empty()); ASSERT_EQ("1935", port); ASSERT_EQ("APP", app); ASSERT_EQ("STREAM", stream_name); brpc::ParseRtmpURL("HOST/APP/STREAM", &host, &vhost, &port, &app, &stream_name); ASSERT_EQ("HOST", host); ASSERT_TRUE(vhost.empty()); ASSERT_EQ("1935", port); ASSERT_EQ("APP", app); ASSERT_EQ("STREAM", stream_name); brpc::ParseRtmpURL("rtmp://HOST:8765//APP?vhost=abc///STREAM?queries", &host, &vhost, &port, &app, &stream_name); ASSERT_EQ("HOST", host); ASSERT_EQ("abc", vhost); ASSERT_EQ("8765", port); ASSERT_EQ("APP", app); ASSERT_EQ("STREAM?queries", stream_name); brpc::ParseRtmpURL("HOST:8765//APP?vhost=abc///STREAM?queries", &host, &vhost, &port, &app, &stream_name); ASSERT_EQ("HOST", host); ASSERT_EQ("abc", vhost); ASSERT_EQ("8765", port); ASSERT_EQ("APP", app); ASSERT_EQ("STREAM?queries", stream_name); brpc::ParseRtmpURL("HOST:8765//APP?vhost=abc///STREAM?queries/", &host, &vhost, &port, &app, &stream_name); ASSERT_EQ("HOST", host); ASSERT_EQ("abc", vhost); ASSERT_EQ("8765", port); ASSERT_EQ("APP", app); ASSERT_EQ("STREAM?queries/", stream_name); brpc::ParseRtmpURL("HOST:8765/APP?vhost=abc", &host, &vhost, &port, &app, &stream_name); ASSERT_EQ("HOST", host); ASSERT_EQ("abc", vhost); ASSERT_EQ("8765", port); ASSERT_EQ("APP", app); ASSERT_TRUE(stream_name.empty()); } TEST(RtmpTest, amf) { std::string req_buf; brpc::RtmpInfo info; brpc::AMFObject obj; std::string dummy = "_result"; { google::protobuf::io::StringOutputStream zc_stream(&req_buf); brpc::AMFOutputStream ostream(&zc_stream); brpc::WriteAMFString(dummy, &ostream); brpc::WriteAMFUint32(17, &ostream); info.set_code("NetConnection.Connect"); // TODO info.set_level("error"); info.set_description("heheda hello foobar"); brpc::WriteAMFObject(info, &ostream); ASSERT_TRUE(ostream.good()); obj.SetString("code", "foo"); obj.SetString("level", "bar"); obj.SetString("description", "heheda"); brpc::WriteAMFObject(obj, &ostream); ASSERT_TRUE(ostream.good()); } google::protobuf::io::ArrayInputStream zc_stream(req_buf.data(), req_buf.size()); brpc::AMFInputStream istream(&zc_stream); std::string result; ASSERT_TRUE(brpc::ReadAMFString(&result, &istream)); ASSERT_EQ(dummy, result); uint32_t num = 0; ASSERT_TRUE(brpc::ReadAMFUint32(&num, &istream)); ASSERT_EQ(17u, num); brpc::RtmpInfo info2; ASSERT_TRUE(brpc::ReadAMFObject(&info2, &istream)); ASSERT_EQ(info.code(), info2.code()); ASSERT_EQ(info.level(), info2.level()); ASSERT_EQ(info.description(), info2.description()); brpc::RtmpInfo info3; ASSERT_TRUE(brpc::ReadAMFObject(&info3, &istream)); ASSERT_EQ("foo", info3.code()); ASSERT_EQ("bar", info3.level()); ASSERT_EQ("heheda", info3.description()); } TEST(RtmpTest, amf_rejects_oversized_string_before_growing_output) { std::string req_buf; AppendAMFLongStringHeader(&req_buf, 16); req_buf.append("short", 5); google::protobuf::io::ArrayInputStream zc_stream(req_buf.data(), req_buf.size()); brpc::AMFInputStream istream(&zc_stream); std::string result = "unchanged"; EXPECT_FALSE(brpc::ReadAMFString(&result, &istream)); EXPECT_TRUE(result.empty()); ScopedAMFLimit scoped_limit(&brpc::FLAGS_amf_max_string_size, 4); std::string capped_buf; AppendAMFLongStringHeader(&capped_buf, 5); capped_buf.append("hello", 5); google::protobuf::io::ArrayInputStream zc_stream2(capped_buf.data(), capped_buf.size()); brpc::AMFInputStream istream2(&zc_stream2); EXPECT_FALSE(brpc::ReadAMFString(&result, &istream2)); } TEST(RtmpTest, amf_rejects_oversized_ecma_array_count) { ScopedAMFLimit scoped_limit(&brpc::FLAGS_amf_max_array_size, 1); std::string req_buf; AppendAMFEcmaArrayHeader(&req_buf, 2); AppendAMFShortStringBody(&req_buf, "x"); req_buf.push_back((char)brpc::AMF_MARKER_NULL); google::protobuf::io::ArrayInputStream zc_stream(req_buf.data(), req_buf.size()); brpc::AMFInputStream istream(&zc_stream); brpc::AMFObject obj; EXPECT_FALSE(brpc::ReadAMFObject(&obj, &istream)); } TEST(RtmpTest, amf_rejects_oversized_strict_array_count) { ScopedAMFLimit scoped_limit(&brpc::FLAGS_amf_max_array_size, 1); std::string req_buf; AppendAMFStrictArrayHeader(&req_buf, 2); req_buf.push_back((char)brpc::AMF_MARKER_NULL); google::protobuf::io::ArrayInputStream zc_stream(req_buf.data(), req_buf.size()); brpc::AMFInputStream istream(&zc_stream); brpc::AMFArray arr; EXPECT_FALSE(brpc::ReadAMFArray(&arr, &istream)); EXPECT_EQ(0u, arr.size()); } TEST(RtmpTest, amf_rejects_deep_nested_arrays) { ScopedAMFMaxDepth scoped_depth(4); std::string req_buf; for (int i = 0; i <= brpc::FLAGS_amf_max_depth + 1; ++i) { AppendAMFStrictArrayHeader(&req_buf, 1); } req_buf.push_back((char)brpc::AMF_MARKER_NULL); google::protobuf::io::ArrayInputStream zc_stream(req_buf.data(), req_buf.size()); brpc::AMFInputStream istream(&zc_stream); brpc::AMFArray arr; EXPECT_FALSE(brpc::ReadAMFArray(&arr, &istream)); req_buf.clear(); for (int i = 0; i < brpc::FLAGS_amf_max_depth; ++i) { AppendAMFStrictArrayHeader(&req_buf, 1); } req_buf.push_back((char)brpc::AMF_MARKER_NULL); google::protobuf::io::ArrayInputStream zc_stream2(req_buf.data(), req_buf.size()); brpc::AMFInputStream istream2(&zc_stream2); brpc::AMFArray valid_arr; EXPECT_TRUE(brpc::ReadAMFArray(&valid_arr, &istream2)); } TEST(RtmpTest, amf_rejects_deep_nested_objects) { ScopedAMFMaxDepth scoped_depth(4); std::string req_buf = MakeNestedAMFObject(brpc::FLAGS_amf_max_depth + 1); google::protobuf::io::ArrayInputStream zc_stream(req_buf.data(), req_buf.size()); brpc::AMFInputStream istream(&zc_stream); brpc::AMFObject obj; EXPECT_FALSE(brpc::ReadAMFObject(&obj, &istream)); req_buf = MakeNestedAMFObject(brpc::FLAGS_amf_max_depth); google::protobuf::io::ArrayInputStream zc_stream2(req_buf.data(), req_buf.size()); brpc::AMFInputStream istream2(&zc_stream2); brpc::AMFObject valid_obj; EXPECT_TRUE(brpc::ReadAMFObject(&valid_obj, &istream2)); } TEST(RtmpTest, amf_rejects_deep_nested_ecma_arrays) { ScopedAMFMaxDepth scoped_depth(4); std::string req_buf = MakeNestedAMFEcmaArray(brpc::FLAGS_amf_max_depth + 1); google::protobuf::io::ArrayInputStream zc_stream(req_buf.data(), req_buf.size()); brpc::AMFInputStream istream(&zc_stream); brpc::AMFObject obj; EXPECT_FALSE(brpc::ReadAMFObject(&obj, &istream)); req_buf = MakeNestedAMFEcmaArray(brpc::FLAGS_amf_max_depth); google::protobuf::io::ArrayInputStream zc_stream2(req_buf.data(), req_buf.size()); brpc::AMFInputStream istream2(&zc_stream2); brpc::AMFObject valid_obj; EXPECT_TRUE(brpc::ReadAMFObject(&valid_obj, &istream2)); } // Create() copies the record into a non-NUL-terminated buffer, so the SPS must // be parsed with an explicitly-sized view. A crafted sequence header whose SPS // body has no zero byte used to make ParseSPS run strlen off the end of that // buffer (an out-of-bounds read, caught here under ASan). TEST(RtmpTest, avc_seq_header_sps_without_zero_byte) { const uint16_t sps_length = 70; // keep the record above the small-array cap butil::IOBuf buf; const char head[6] = { 0x01, 0x64, 0x00, 0x28, (char)0xff, (char)0xe1 }; buf.append(head, sizeof(head)); // version/profile/level, lengthSizeMinus1=3, numSPS=1 const char len_be[2] = { (char)(sps_length >> 8), (char)(sps_length & 0xff) }; buf.append(len_be, sizeof(len_be)); std::string sps(sps_length, (char)0x67); // NAL header 0x67 then non-zero filler buf.append(sps); brpc::AVCDecoderConfigurationRecord avc; // Only requirement: the call must not read past the copied record. avc.Create(buf); } static void AppendBigEndian3Bytes(std::string* s, uint32_t v) { s->push_back((char)((v >> 16) & 0xFF)); s->push_back((char)((v >> 8) & 0xFF)); s->push_back((char)(v & 0xFF)); } // Build an FLV stream header followed by a single tag of `tag_type' whose // DataSize field is set to `data_size'. No tag body is appended, so a valid // tag needs data_size==0 to be rejected before any body is consumed. static std::string MakeFlvTagWithDataSize(char tag_type, uint32_t data_size) { std::string s; const char header[] = { 'F', 'L', 'V', 0x01, 0x05, 0, 0, 0, 0x09 }; s.append(header, sizeof(header)); s.append(4, '\0'); // PreviousTagSize0 s.push_back(tag_type); AppendBigEndian3Bytes(&s, data_size); // DataSize s.append(3, '\0'); // Timestamp s.push_back('\0'); // TimestampExtended s.append(3, '\0'); // StreamID s.append(4, '\0'); // PreviousTagSize return s; } TEST(RtmpTest, flv_reader_rejects_zero_datasize_video_tag) { std::string flv = MakeFlvTagWithDataSize((char)brpc::FLV_TAG_VIDEO, 0); butil::IOBuf buf; buf.append(flv); brpc::FlvReader reader(&buf); brpc::FlvTagType type; ASSERT_TRUE(reader.PeekMessageType(&type).ok()); ASSERT_EQ(brpc::FLV_TAG_VIDEO, type); const size_t before = buf.size(); brpc::RtmpVideoMessage vmsg; // A DataSize of 0 used to underflow `msg_size - 1' and drain the whole // buffer; it must now be rejected without consuming the tag. ASSERT_FALSE(reader.Read(&vmsg).ok()); ASSERT_EQ(before, buf.size()); } TEST(RtmpTest, flv_reader_rejects_zero_datasize_audio_tag) { std::string flv = MakeFlvTagWithDataSize((char)brpc::FLV_TAG_AUDIO, 0); butil::IOBuf buf; buf.append(flv); brpc::FlvReader reader(&buf); brpc::FlvTagType type; ASSERT_TRUE(reader.PeekMessageType(&type).ok()); ASSERT_EQ(brpc::FLV_TAG_AUDIO, type); const size_t before = buf.size(); brpc::RtmpAudioMessage amsg; ASSERT_FALSE(reader.Read(&amsg).ok()); ASSERT_EQ(before, buf.size()); } TEST(RtmpTest, successfully_play_streams) { PlayingDummyService rtmp_service; brpc::Server server; brpc::ServerOptions server_opt; server_opt.rtmp_service = &rtmp_service; ASSERT_EQ(0, server.Start(8571, &server_opt)); brpc::RtmpClientOptions rtmp_opt; rtmp_opt.app = "hello"; rtmp_opt.swfUrl = "anything"; rtmp_opt.tcUrl = "rtmp://heheda"; brpc::RtmpClient rtmp_client; ASSERT_EQ(0, rtmp_client.Init("localhost:8571", rtmp_opt)); // Create multiple streams. const int NSTREAM = 2; brpc::DestroyingPtr cstreams[NSTREAM]; for (int i = 0; i < NSTREAM; ++i) { cstreams[i].reset(new TestRtmpClientStream); brpc::RtmpClientStreamOptions opt; opt.play_name = butil::string_printf("play_name_%d", i); //opt.publish_name = butil::string_printf("pub_name_%d", i); opt.wait_until_play_or_publish_is_sent = true; cstreams[i]->Init(&rtmp_client, opt); } sleep(5); for (int i = 0; i < NSTREAM; ++i) { cstreams[i]->assertions_on_successful_play(); } LOG(INFO) << "Quiting program..."; } TEST(RtmpTest, fail_to_play_streams) { PlayingDummyService rtmp_service; brpc::Server server; brpc::ServerOptions server_opt; server_opt.rtmp_service = &rtmp_service; ASSERT_EQ(0, server.Start(8571, &server_opt)); brpc::RtmpClientOptions rtmp_opt; rtmp_opt.app = "hello"; rtmp_opt.swfUrl = "anything"; rtmp_opt.tcUrl = "rtmp://heheda"; brpc::RtmpClient rtmp_client; ASSERT_EQ(0, rtmp_client.Init("localhost:8571", rtmp_opt)); // Create multiple streams. const int NSTREAM = 2; brpc::DestroyingPtr cstreams[NSTREAM]; for (int i = 0; i < NSTREAM; ++i) { cstreams[i].reset(new TestRtmpClientStream); brpc::RtmpClientStreamOptions opt; opt.play_name = UNEXIST_NAME; opt.wait_until_play_or_publish_is_sent = true; cstreams[i]->Init(&rtmp_client, opt); } sleep(1); for (int i = 0; i < NSTREAM; ++i) { cstreams[i]->assertions_on_failure(); } LOG(INFO) << "Quiting program..."; } TEST(RtmpTest, successfully_publish_streams) { PublishService rtmp_service; brpc::Server server; brpc::ServerOptions server_opt; server_opt.rtmp_service = &rtmp_service; ASSERT_EQ(0, server.Start(8571, &server_opt)); brpc::RtmpClientOptions rtmp_opt; rtmp_opt.app = "hello"; rtmp_opt.swfUrl = "anything"; rtmp_opt.tcUrl = "rtmp://heheda"; brpc::RtmpClient rtmp_client; ASSERT_EQ(0, rtmp_client.Init("localhost:8571", rtmp_opt)); // Create multiple streams. const int NSTREAM = 2; brpc::DestroyingPtr cstreams[NSTREAM]; for (int i = 0; i < NSTREAM; ++i) { cstreams[i].reset(new TestRtmpClientStream); brpc::RtmpClientStreamOptions opt; opt.publish_name = butil::string_printf("pub_name_%d", i); opt.wait_until_play_or_publish_is_sent = true; cstreams[i]->Init(&rtmp_client, opt); } const int REP = 5; for (int i = 0; i < REP; ++i) { brpc::RtmpVideoMessage vmsg; vmsg.timestamp = 1000 + i * 20; vmsg.frame_type = brpc::FLV_VIDEO_FRAME_KEYFRAME; vmsg.codec = brpc::FLV_VIDEO_AVC; vmsg.data.append(butil::string_printf("video_%d", i)); for (int j = 0; j < NSTREAM; j += 2) { ASSERT_EQ(0, cstreams[j]->SendVideoMessage(vmsg)); } brpc::RtmpAudioMessage amsg; amsg.timestamp = 1000 + i * 20; amsg.codec = brpc::FLV_AUDIO_AAC; amsg.rate = brpc::FLV_SOUND_RATE_44100HZ; amsg.bits = brpc::FLV_SOUND_16BIT; amsg.type = brpc::FLV_SOUND_STEREO; amsg.data.append(butil::string_printf("audio_%d", i)); for (int j = 1; j < NSTREAM; j += 2) { ASSERT_EQ(0, cstreams[j]->SendAudioMessage(amsg)); } bthread_usleep(500000); } std::vector > created_streams; rtmp_service.move_created_streams(&created_streams); ASSERT_EQ(NSTREAM, (int)created_streams.size()); for (int i = 0; i < NSTREAM; ++i) { EXPECT_EQ(1, created_streams[i]->_called_on_first_message); } for (int j = 0; j < NSTREAM; j += 2) { ASSERT_EQ(REP, created_streams[j]->_nvideomsg); } for (int j = 1; j < NSTREAM; j += 2) { ASSERT_EQ(REP, created_streams[j]->_naudiomsg); } LOG(INFO) << "Quiting program..."; } TEST(RtmpTest, failed_to_publish_streams) { PublishService rtmp_service; brpc::Server server; brpc::ServerOptions server_opt; server_opt.rtmp_service = &rtmp_service; ASSERT_EQ(0, server.Start(8575, &server_opt)); brpc::RtmpClientOptions rtmp_opt; rtmp_opt.app = "hello"; rtmp_opt.swfUrl = "anything"; rtmp_opt.tcUrl = "rtmp://heheda"; brpc::RtmpClient rtmp_client; ASSERT_EQ(0, rtmp_client.Init("localhost:8575", rtmp_opt)); // Create multiple streams. const int NSTREAM = 2; brpc::DestroyingPtr cstreams[NSTREAM]; for (int i = 0; i < NSTREAM; ++i) { cstreams[i].reset(new TestRtmpClientStream); brpc::RtmpClientStreamOptions opt; opt.publish_name = UNEXIST_NAME; opt.wait_until_play_or_publish_is_sent = true; cstreams[i]->Init(&rtmp_client, opt); } sleep(1); for (int i = 0; i < NSTREAM; ++i) { cstreams[i]->assertions_on_failure(); } std::vector > created_streams; rtmp_service.move_created_streams(&created_streams); ASSERT_EQ(NSTREAM, (int)created_streams.size()); for (int i = 0; i < NSTREAM; ++i) { ASSERT_EQ(0, created_streams[i]->_called_on_first_message); ASSERT_EQ(0, created_streams[i]->_nvideomsg); ASSERT_EQ(0, created_streams[i]->_naudiomsg); } LOG(INFO) << "Quiting program..."; } TEST(RtmpTest, failed_to_connect_client_streams) { brpc::RtmpClientOptions rtmp_opt; rtmp_opt.app = "hello"; rtmp_opt.swfUrl = "anything"; rtmp_opt.tcUrl = "rtmp://heheda"; brpc::RtmpClient rtmp_client; ASSERT_EQ(0, rtmp_client.Init("localhost:8572", rtmp_opt)); // Create multiple streams. const int NSTREAM = 2; brpc::DestroyingPtr cstreams[NSTREAM]; for (int i = 0; i < NSTREAM; ++i) { cstreams[i].reset(new TestRtmpClientStream); brpc::RtmpClientStreamOptions opt; opt.play_name = butil::string_printf("play_name_%d", i); opt.wait_until_play_or_publish_is_sent = true; cstreams[i]->Init(&rtmp_client, opt); cstreams[i]->assertions_on_failure(); } LOG(INFO) << "Quiting program..."; } TEST(RtmpTest, destroy_client_streams_before_init) { brpc::RtmpClientOptions rtmp_opt; rtmp_opt.app = "hello"; rtmp_opt.swfUrl = "anything"; rtmp_opt.tcUrl = "rtmp://heheda"; brpc::RtmpClient rtmp_client; ASSERT_EQ(0, rtmp_client.Init("localhost:8573", rtmp_opt)); // Create multiple streams. const int NSTREAM = 2; butil::intrusive_ptr cstreams[NSTREAM]; for (int i = 0; i < NSTREAM; ++i) { cstreams[i].reset(new TestRtmpClientStream); cstreams[i]->Destroy(); ASSERT_EQ(1, cstreams[i]->_called_on_stop); ASSERT_EQ(brpc::RtmpClientStream::STATE_DESTROYING, cstreams[i]->_state); brpc::RtmpClientStreamOptions opt; opt.play_name = butil::string_printf("play_name_%d", i); opt.wait_until_play_or_publish_is_sent = true; cstreams[i]->Init(&rtmp_client, opt); cstreams[i]->assertions_on_failure(); } LOG(INFO) << "Quiting program..."; } TEST(RtmpTest, destroy_retrying_client_streams_before_init) { brpc::RtmpClientOptions rtmp_opt; rtmp_opt.app = "hello"; rtmp_opt.swfUrl = "anything"; rtmp_opt.tcUrl = "rtmp://heheda"; brpc::RtmpClient rtmp_client; ASSERT_EQ(0, rtmp_client.Init("localhost:8573", rtmp_opt)); // Create multiple streams. const int NSTREAM = 2; butil::intrusive_ptr cstreams[NSTREAM]; for (int i = 0; i < NSTREAM; ++i) { cstreams[i].reset(new TestRtmpRetryingClientStream); cstreams[i]->Destroy(); ASSERT_EQ(1, cstreams[i]->_called_on_stop); brpc::RtmpRetryingClientStreamOptions opt; opt.play_name = butil::string_printf("play_name_%d", i); brpc::SubStreamCreator* sc = new RtmpSubStreamCreator(&rtmp_client); cstreams[i]->Init(sc, opt); ASSERT_EQ(1, cstreams[i]->_called_on_stop); } LOG(INFO) << "Quiting program..."; } TEST(RtmpTest, destroy_client_streams_during_creation) { PlayingDummyService rtmp_service(2000/*sleep 2s*/); brpc::Server server; brpc::ServerOptions server_opt; server_opt.rtmp_service = &rtmp_service; ASSERT_EQ(0, server.Start(8574, &server_opt)); brpc::RtmpClientOptions rtmp_opt; rtmp_opt.app = "hello"; rtmp_opt.swfUrl = "anything"; rtmp_opt.tcUrl = "rtmp://heheda"; brpc::RtmpClient rtmp_client; ASSERT_EQ(0, rtmp_client.Init("localhost:8574", rtmp_opt)); // Create multiple streams. const int NSTREAM = 2; butil::intrusive_ptr cstreams[NSTREAM]; for (int i = 0; i < NSTREAM; ++i) { cstreams[i].reset(new TestRtmpClientStream); brpc::RtmpClientStreamOptions opt; opt.play_name = butil::string_printf("play_name_%d", i); cstreams[i]->Init(&rtmp_client, opt); ASSERT_EQ(0, cstreams[i]->_called_on_stop); usleep(500*1000); ASSERT_EQ(0, cstreams[i]->_called_on_stop); cstreams[i]->Destroy(); usleep(10*1000); ASSERT_EQ(1, cstreams[i]->_called_on_stop); } LOG(INFO) << "Quiting program..."; } TEST(RtmpTest, destroy_retrying_client_streams_during_creation) { PlayingDummyService rtmp_service(2000/*sleep 2s*/); brpc::Server server; brpc::ServerOptions server_opt; server_opt.rtmp_service = &rtmp_service; ASSERT_EQ(0, server.Start(8574, &server_opt)); brpc::RtmpClientOptions rtmp_opt; rtmp_opt.app = "hello"; rtmp_opt.swfUrl = "anything"; rtmp_opt.tcUrl = "rtmp://heheda"; brpc::RtmpClient rtmp_client; ASSERT_EQ(0, rtmp_client.Init("localhost:8574", rtmp_opt)); // Create multiple streams. const int NSTREAM = 2; butil::intrusive_ptr cstreams[NSTREAM]; for (int i = 0; i < NSTREAM; ++i) { cstreams[i].reset(new TestRtmpRetryingClientStream); brpc::RtmpRetryingClientStreamOptions opt; opt.play_name = butil::string_printf("play_name_%d", i); brpc::SubStreamCreator* sc = new RtmpSubStreamCreator(&rtmp_client); cstreams[i]->Init(sc, opt); ASSERT_EQ(0, cstreams[i]->_called_on_stop); usleep(500*1000); ASSERT_EQ(0, cstreams[i]->_called_on_stop); cstreams[i]->Destroy(); usleep(10*1000); ASSERT_EQ(1, cstreams[i]->_called_on_stop); } LOG(INFO) << "Quiting program..."; } TEST(RtmpTest, retrying_stream) { PlayingDummyService rtmp_service; brpc::Server server; brpc::ServerOptions server_opt; server_opt.rtmp_service = &rtmp_service; ASSERT_EQ(0, server.Start(8576, &server_opt)); brpc::RtmpClientOptions rtmp_opt; rtmp_opt.app = "hello"; rtmp_opt.swfUrl = "anything"; rtmp_opt.tcUrl = "rtmp://heheda"; brpc::RtmpClient rtmp_client; ASSERT_EQ(0, rtmp_client.Init("localhost:8576", rtmp_opt)); // Create multiple streams. const int NSTREAM = 2; brpc::DestroyingPtr cstreams[NSTREAM]; for (int i = 0; i < NSTREAM; ++i) { cstreams[i].reset(new TestRtmpRetryingClientStream); brpc::Controller cntl; brpc::RtmpRetryingClientStreamOptions opt; opt.play_name = butil::string_printf("name_%d", i); brpc::SubStreamCreator* sc = new RtmpSubStreamCreator(&rtmp_client); cstreams[i]->Init(sc, opt); } sleep(3); LOG(INFO) << "Stopping server"; server.Stop(0); server.Join(); LOG(INFO) << "Stopped server and sleep for a while"; sleep(3); ASSERT_EQ(0, server.Start(8576, &server_opt)); sleep(3); for (int i = 0; i < NSTREAM; ++i) { ASSERT_EQ(1, cstreams[i]->_called_on_first_message); ASSERT_EQ(2, cstreams[i]->_called_on_playable); } LOG(INFO) << "Quiting program..."; }