fix #165, refine dh wrapper, ensure public key is 128bytes. 0.9.207.

winlin
Commit 9b6187c3d5f48f19c68edcb24cc611bbff249d0c 9b6187c3 1 parent cb311d99
trunk/src/core/srs_core.hpp
trunk/src/rtmp/srs_protocol_handshake.cpp
trunk/src/rtmp/srs_protocol_handshake.hpp
trunk/src/utest/srs_utest_protocol.cpp
--- a/trunk/src/core/srs_core.hpp
查看文件 @9b6187c
+++ b/trunk/src/core/srs_core.hpp
查看文件 @9b6187c
@@ -31,7 +31,7 @@ CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 // current release version
 #define VERSION_MAJOR "0"
 #define VERSION_MINOR "9"
-#define VERSION_REVISION "206"
+#define VERSION_REVISION "207"
 #define RTMP_SIG_SRS_VERSION VERSION_MAJOR"."VERSION_MINOR"."VERSION_REVISION
 // server info.
 #define RTMP_SIG_SRS_KEY "SRS"
--- a/trunk/src/rtmp/srs_protocol_handshake.cpp
查看文件 @9b6187c
+++ b/trunk/src/rtmp/srs_protocol_handshake.cpp
查看文件 @9b6187c
@@ -190,7 +190,7 @@ namespace _srs_internal
         return ret;
     }
-    int SrsDH::copy_public_key(char* pkey, int32_t* ppkey_size)
+    int SrsDH::copy_public_key(char* pkey, int32_t& pkey_size)
     {
         int ret = ERROR_SUCCESS;
@@ -199,20 +199,21 @@ namespace _srs_internal
         int32_t key_size = BN_num_bytes(pdh->pub_key);
         srs_assert(key_size > 0);
+        // maybe the key_size is 127, but dh will write all 128bytes pkey,
+        // so, donot need to set/initialize the pkey.
+        // @see https://github.com/winlinvip/simple-rtmp-server/issues/165
         key_size = BN_bn2bin(pdh->pub_key, (unsigned char*)pkey);
         srs_assert(key_size > 0);
-        if (ppkey_size != NULL) {
         // output the size of public key.
         // @see https://github.com/winlinvip/simple-rtmp-server/issues/165
-            srs_assert(key_size <= *ppkey_size);
-            *ppkey_size = key_size;
-        }
+        srs_assert(key_size <= pkey_size);
+        pkey_size = key_size;
         return ret;
     }
-    int SrsDH::copy_shared_key(const char* ppkey, int32_t ppkey_size, char* skey, int32_t* pskey_size)
+    int SrsDH::copy_shared_key(const char* ppkey, int32_t ppkey_size, char* skey, int32_t& skey_size)
     {
         int ret = ERROR_SUCCESS;
@@ -223,22 +224,19 @@ namespace _srs_internal
         }
         // if failed, donot return, do cleanup, @see ./test/dhtest.c:168
+        // maybe the key_size is 127, but dh will write all 128bytes skey,
+        // so, donot need to set/initialize the skey.
+        // @see https://github.com/winlinvip/simple-rtmp-server/issues/165
         int32_t key_size = DH_compute_key((unsigned char*)skey, ppk, pdh);
         if (key_size < ppkey_size) {
             srs_warn("shared key size=%d, ppk_size=%d", key_size, ppkey_size);
         }
-        if (key_size < 0) {
-            ret = ERROR_OpenSslComputeSharedKey;
-        } else {
-            if (pskey_size != NULL) {
-                if (key_size > *pskey_size) {
+        if (key_size < 0 || key_size > skey_size) {
             ret = ERROR_OpenSslComputeSharedKey;
         } else {
-                    *pskey_size = key_size;
-                }
-            }
+            skey_size = key_size;
         }
         if (ppk) {
@@ -936,29 +934,36 @@ namespace _srs_internal
         version = 0x01000504; // server s1 version
         SrsDH dh;
+        
+        // ensure generate 128bytes public key.
         if ((ret = dh.initialize(true)) != ERROR_SUCCESS) {
             return ret;
         }
+        
         if (schema == srs_schema0) {
             srs_key_block_init(&block0.key);
             srs_digest_block_init(&block1.digest);
             // directly generate the public key.
             // @see: https://github.com/winlinvip/simple-rtmp-server/issues/148
-            if ((ret = dh.copy_public_key((char*)block0.key.key, NULL)) != ERROR_SUCCESS) {
+            int pkey_size = 128;
+            if ((ret = dh.copy_public_key((char*)block0.key.key, pkey_size)) != ERROR_SUCCESS) {
                 srs_error("calc s1 key failed. ret=%d", ret);
                 return ret;
             }
+            srs_assert(pkey_size == 128);
         } else {
             srs_digest_block_init(&block0.digest);
             srs_key_block_init(&block1.key);
             // directly generate the public key.
             // @see: https://github.com/winlinvip/simple-rtmp-server/issues/148
-            if ((ret = dh.copy_public_key((char*)block1.key.key, NULL)) != ERROR_SUCCESS) {
+            int pkey_size = 128;
+            if ((ret = dh.copy_public_key((char*)block1.key.key, pkey_size)) != ERROR_SUCCESS) {
                 srs_error("calc s1 key failed. ret=%d", ret);
                 return ret;
             }
+            srs_assert(pkey_size == 128);
         }
         srs_verbose("calc s1 key success.");
--- a/trunk/src/rtmp/srs_protocol_handshake.hpp
查看文件 @9b6187c
+++ b/trunk/src/rtmp/srs_protocol_handshake.hpp
查看文件 @9b6187c
@@ -141,21 +141,21 @@ namespace _srs_internal
         /**
         * copy the public key.
         * @param pkey the bytes to copy the public key.
-        * @param ppkey_size the max public key size, output the actual public key size.
-        *       NULL to ignore.
+        * @param pkey_size the max public key size, output the actual public key size.
+        *       user should never ignore this size.
         * @remark, when ensure_128bytes_public_key, the size always 128.
         */
-        virtual int copy_public_key(char* pkey, int32_t* ppkey_size);
+        virtual int copy_public_key(char* pkey, int32_t& pkey_size);
         /**
         * generate and copy the shared key.
         * generate the shared key with peer public key.
         * @param ppkey peer public key.
         * @param ppkey_size the size of ppkey.
         * @param skey the computed shared key.
-        * @param pskey_size the max shared key size, output the actual shared key size.
-        *       NULL to ignore.
+        * @param skey_size the max shared key size, output the actual shared key size.
+        *       user should never ignore this size.
         */
-        virtual int copy_shared_key(const char* ppkey, int32_t ppkey_size, char* skey, int32_t* pskey_size);
+        virtual int copy_shared_key(const char* ppkey, int32_t ppkey_size, char* skey, int32_t& skey_size);
     private:
         virtual int do_initialize();
     };
--- a/trunk/src/utest/srs_utest_protocol.cpp
查看文件 @9b6187c
+++ b/trunk/src/utest/srs_utest_protocol.cpp
查看文件 @9b6187c
@@ -243,10 +243,13 @@ VOID TEST(ProtocolHandshakeTest, DHKey)
     ASSERT_TRUE(ERROR_SUCCESS == dh.initialize(true));
     char pub_key1[128];
-    EXPECT_TRUE(ERROR_SUCCESS == dh.copy_public_key(pub_key1, NULL));
+    int pkey_size = 128;
+    EXPECT_TRUE(ERROR_SUCCESS == dh.copy_public_key(pub_key1, pkey_size));
+    ASSERT_EQ(128, pkey_size);
     char pub_key2[128];
-    EXPECT_TRUE(ERROR_SUCCESS == dh.copy_public_key(pub_key2, NULL));
+    EXPECT_TRUE(ERROR_SUCCESS == dh.copy_public_key(pub_key2, pkey_size));
+    ASSERT_EQ(128, pkey_size);
     EXPECT_TRUE(srs_bytes_equals(pub_key1, pub_key2, 128));
@@ -255,7 +258,8 @@ VOID TEST(ProtocolHandshakeTest, DHKey)
     ASSERT_TRUE(ERROR_SUCCESS == dh0.initialize(true));
-    EXPECT_TRUE(ERROR_SUCCESS == dh0.copy_public_key(pub_key2, NULL));
+    EXPECT_TRUE(ERROR_SUCCESS == dh0.copy_public_key(pub_key2, pkey_size));
+    ASSERT_EQ(128, pkey_size);
     EXPECT_FALSE(srs_bytes_equals(pub_key1, pub_key2, 128));
 }