diff --git a/src/brpc/amf.cpp b/src/brpc/amf.cpp index 98025431ab..219199c78f 100644 --- a/src/brpc/amf.cpp +++ b/src/brpc/amf.cpp @@ -27,6 +27,10 @@ namespace brpc { DEFINE_int32(amf_max_depth, 128, "Maximum nesting depth for AMF objects and arrays"); +DEFINE_int32(amf_max_string_size, 64 * 1024 * 1024, + "Maximum byte size for AMF strings"); +DEFINE_int32(amf_max_array_size, 1024 * 1024, + "Maximum element count for AMF arrays"); static bool CheckAMFDepth(int depth) { if (depth > FLAGS_amf_max_depth) { @@ -37,6 +41,26 @@ static bool CheckAMFDepth(int depth) { return true; } +static bool CheckAMFStringSize(uint32_t len) { + if (FLAGS_amf_max_string_size < 0 || + len > (uint32_t)FLAGS_amf_max_string_size) { + LOG(ERROR) << "AMF string exceeds max size! max=" + << FLAGS_amf_max_string_size << ", actually=" << len; + return false; + } + return true; +} + +static bool CheckAMFArraySize(uint32_t count) { + if (FLAGS_amf_max_array_size < 0 || + count > (uint32_t)FLAGS_amf_max_array_size) { + LOG(ERROR) << "AMF array exceeds max size! max=" + << FLAGS_amf_max_array_size << ", actually=" << count; + return false; + } + return true; +} + const char* marker2str(AMFMarker marker) { switch (marker) { case AMF_MARKER_NUMBER: return "number"; @@ -268,8 +292,12 @@ static bool ReadAMFShortStringBody(std::string* str, AMFInputStream* stream) { LOG(ERROR) << "stream is not long enough"; return false; } + if (!CheckAMFStringSize(len)) { + return false; + } str->resize(len); if (len != 0 && stream->cutn(&(*str)[0], len) != len) { + str->clear(); LOG(ERROR) << "stream is not long enough"; return false; } @@ -282,8 +310,12 @@ static bool ReadAMFLongStringBody(std::string* str, AMFInputStream* stream) { LOG(ERROR) << "stream is not long enough"; return false; } + if (!CheckAMFStringSize(len)) { + return false; + } str->resize(len); if (len != 0 && stream->cutn(&(*str)[0], len) != len) { + str->clear(); LOG(ERROR) << "stream is not long enough"; return false; } @@ -584,6 +616,9 @@ static bool ReadAMFEcmaArrayBody(google::protobuf::Message* message, LOG(ERROR) << "stream is not long enough"; return false; } + if (!CheckAMFArraySize(count)) { + return false; + } const google::protobuf::Descriptor* desc = message->GetDescriptor(); std::string name; for (uint32_t i = 0; i < count; ++i) { @@ -760,6 +795,9 @@ static bool ReadAMFEcmaArrayBody(AMFObject* obj, AMFInputStream* stream, int dep LOG(ERROR) << "stream is not long enough"; return false; } + if (!CheckAMFArraySize(count)) { + return false; + } std::string name; for (uint32_t i = 0; i < count; ++i) { if (!ReadAMFShortStringBody(&name, stream)) { @@ -892,6 +930,9 @@ static bool ReadAMFArrayBody(AMFArray* arr, AMFInputStream* stream, int depth) { LOG(ERROR) << "stream is not long enough"; return false; } + if (!CheckAMFArraySize(count)) { + return false; + } for (uint32_t i = 0; i < count; ++i) { if (!ReadAMFArrayItem(stream, arr, depth)) { return false; diff --git a/test/brpc_rtmp_unittest.cpp b/test/brpc_rtmp_unittest.cpp index e6b9c62f67..6834036a3a 100644 --- a/test/brpc_rtmp_unittest.cpp +++ b/test/brpc_rtmp_unittest.cpp @@ -37,6 +37,8 @@ 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[]) { @@ -59,6 +61,20 @@ class ScopedAMFMaxDepth { 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)); @@ -86,6 +102,14 @@ void AppendAMFShortStringBody(std::string* out, const char* name) { 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); @@ -597,6 +621,56 @@ TEST(RtmpTest, amf) { 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);