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

README.md

@@ -208,6 +208,7 @@ Supported operating systems and hardware:
 * 2013-10-17, Created.<br/>
 
 ## History
+* v1.0, 2014-08-22, for [#165](https://github.com/winlinvip/simple-rtmp-server/issues/165), refine dh wrapper, ensure public key is 128bytes. 0.9.206.
 * v1.0, 2014-08-19, for [#160](https://github.com/winlinvip/simple-rtmp-server/issues/160), support forward/edge to flussonic, disable debug_srs_upnode to make flussonic happy. 0.9.201.
 * v1.0, 2014-08-17, for [#155](https://github.com/winlinvip/simple-rtmp-server/issues/155), refine for osx, with ssl/http, disable statistics. 0.9.198.
 * v1.0, 2014-08-06, fix [#148](https://github.com/winlinvip/simple-rtmp-server/issues/148), simplify the RTMP handshake key generation. 0.9.191.

trunk/src/core/srs_core.hpp

@@ -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 "205"
+#define VERSION_REVISION "206"
 #define RTMP_SIG_SRS_VERSION VERSION_MAJOR"."VERSION_MINOR"."VERSION_REVISION
 // server info.
 #define RTMP_SIG_SRS_KEY "SRS"

trunk/src/rtmp/srs_protocol_handshake.cpp

@@ -145,88 +145,100 @@ namespace _srs_internal
             "E485B576625E7EC6F44C42E9A637ED6B0BFF5CB6F406B7ED" \
             "EE386BFB5A899FA5AE9F24117C4B1FE649286651ECE65381" \
             "FFFFFFFFFFFFFFFF"
-    /**
+    
-    * initialize DH, create the public/private key.
+    SrsDH::SrsDH()
-    */
-    int __openssl_initialize_dh(DH* pdh, int32_t bits_count) 
     {
-        int ret = ERROR_SUCCESS;
+        pdh = NULL;
+    }
     
-        //2. Create his internal p and g
+    SrsDH::~SrsDH()
-        if ((pdh->p = BN_new()) == NULL) {
+    {
-            ret = ERROR_OpenSslCreateP; 
+        if (pdh != NULL) {
-            return ret;
+            if (pdh->p != NULL) {
+                BN_free(pdh->p);
+                pdh->p = NULL;
             }
-        if ((pdh->g = BN_new()) == NULL) {
+            if (pdh->g != NULL) {
-            ret = ERROR_OpenSslCreateG; 
+                BN_free(pdh->g);
-            return ret;
+                pdh->g = NULL;
             }
-    
+            DH_free(pdh);
-        //3. initialize p and g
+            pdh = NULL;
-        if (BN_hex2bn(&pdh->p, RFC2409_PRIME_1024) == 0) {
-            ret = ERROR_OpenSslParseP1024; 
-            return ret;
         }
-        if (BN_set_word(pdh->g, 2) != 1) {
-            ret = ERROR_OpenSslSetG;
-            return ret;
     }
     
-        //4. Set the key length
+    int SrsDH::initialize(bool ensure_128bytes_public_key)
-        pdh->length = bits_count;
+    {
+        int ret = ERROR_SUCCESS;
         
-        //5. Generate private and public key
+        for (;;) {
-        if (DH_generate_key(pdh) != 1) {
+            if ((ret = do_initialize()) != ERROR_SUCCESS) {
-            ret = ERROR_OpenSslGenerateDHKeys;
                 return ret;
             }
             
-        return ret;
+            if (ensure_128bytes_public_key) {
+                int32_t key_size = BN_num_bytes(pdh->pub_key);
+                if (key_size != 128) {
+                    srs_warn("regenerate 128B key, current=%dB", key_size);
+                    continue;
+                }
             }
-    /**
-    * create DH and copy the 128bytes public key.
-    */
-    int __openssl_copy_key(DH* pdh, char* public_key, int32_t size)
-    {
-        int ret = ERROR_SUCCESS;
             
-        int32_t bits_count = 1024; 
+            break;
+        }
         
-        // 2. generate the g, p, private/public key.
-        if ((ret = __openssl_initialize_dh(pdh, bits_count)) != ERROR_SUCCESS) {
         return ret;
     }
     
+    int SrsDH::copy_public_key(char* pkey, int32_t* ppkey_size)
+    {
+        int ret = ERROR_SUCCESS;
+        
         // copy public key to bytes.
         // sometimes, the key_size is 127, seems ok.
         int32_t key_size = BN_num_bytes(pdh->pub_key);
         srs_assert(key_size > 0);
         
-        if (BN_bn2bin(pdh->pub_key, (unsigned char*)public_key) != size) {
+        key_size = BN_bn2bin(pdh->pub_key, (unsigned char*)pkey);
-            //("Unable to copy key"); return ret;
+        srs_assert(key_size > 0);
-            ret = ERROR_OpenSslCopyKey;
+        
-            return ret;
+        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;
         }
         
         return ret;
     }
-    /**
+    
-    * use exists DH to create and copy the 128bytes shared key.
+    int SrsDH::copy_shared_key(const char* ppkey, int32_t ppkey_size, char* skey, int32_t* pskey_size)
-    * the peer public key used to generate the shared key.
-    */
-    int __openssl_copy_shared_key(DH* pdh, const char* peer_pub_key, int ppk_size, char* shared_key)
     {
         int ret = ERROR_SUCCESS;
         
         BIGNUM* ppk = NULL;
-        if ((ppk = BN_bin2bn((const unsigned char*)peer_pub_key, ppk_size, 0)) == NULL) {
+        if ((ppk = BN_bin2bn((const unsigned char*)ppkey, ppkey_size, 0)) == NULL) {
             ret = ERROR_OpenSslGetPeerPublicKey;
             return ret;
         }
         
-        // if failed, donot return, do cleanup.
+        // if failed, donot return, do cleanup, @see ./test/dhtest.c:168
-        if (DH_compute_key((unsigned char*)shared_key, ppk, pdh) < 0) {
+        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) {
+                    ret = ERROR_OpenSslComputeSharedKey;
+                } else {
+                    *pskey_size = key_size;
+                }
+            }
         }
         
         if (ppk) {
@@ -235,60 +247,50 @@ namespace _srs_internal
         
         return ret;
     }
-    /**
+    
-    * create DH and copy the 128bytes public key,
+    int SrsDH::do_initialize()
-    * generate and copy the shared key.
-    */
-    int __openssl_compute_key(DH* pdh, const char* peer_pub_key, int ppk_size, char* public_key, char* shared_key)
     {
         int ret = ERROR_SUCCESS;
         
-        // create DH and copy the 128bytes public key
+        int32_t bits_count = 1024; 
-        if ((ret = __openssl_copy_key(pdh, public_key, ppk_size)) != ERROR_SUCCESS) {
-            return ret;
-        }
         
-        // generate and copy the shared key
+        //1. Create the DH
-        if ((ret = __openssl_copy_shared_key(pdh, peer_pub_key, ppk_size, shared_key)) != ERROR_SUCCESS) {
+        if ((pdh = DH_new()) == NULL) {
+            ret = ERROR_OpenSslCreateDH; 
             return ret;
         }
     
+        //2. Create his internal p and g
+        if ((pdh->p = BN_new()) == NULL) {
+            ret = ERROR_OpenSslCreateP; 
             return ret;
         }
-    void __openssl_free(DH* pdh)
+        if ((pdh->g = BN_new()) == NULL) {
-    {
+            ret = ERROR_OpenSslCreateG; 
-        if (pdh != NULL) {
+            return ret;
-            if (pdh->p != NULL) {
-                BN_free(pdh->p);
-                pdh->p = NULL;
-            }
-            if (pdh->g != NULL) {
-                BN_free(pdh->g);
-                pdh->g = NULL;
         }
-            DH_free(pdh);
+    
-            pdh = NULL;
+        //3. initialize p and g, @see ./test/ectest.c:260
+        if (!BN_hex2bn(&pdh->p, RFC2409_PRIME_1024)) {
+            ret = ERROR_OpenSslParseP1024; 
+            return ret;
         }
+        // @see ./test/bntest.c:1764
+        if (!BN_set_word(pdh->g, 2)) {
+            ret = ERROR_OpenSslSetG;
+            return ret;
         }
-    int openssl_generate_key(char* public_key, int32_t size)
-    {
-        int ret = ERROR_SUCCESS;
     
-        // Initialize
+        // 4. Set the key length
-        DH* pdh = NULL;
+        pdh->length = bits_count;
     
-        //1. Create the DH
+        // 5. Generate private and public key
-        if ((pdh = DH_new()) == NULL) {
+        // @see ./test/dhtest.c:152
-            ret = ERROR_OpenSslCreateDH; 
+        if (!DH_generate_key(pdh)) {
+            ret = ERROR_OpenSslGenerateDHKeys;
             return ret;
         }
         
-        // generate and copy key.
-        ret = __openssl_copy_key(pdh, public_key, size);
-        
-        // cleanup
-        __openssl_free(pdh);
-    
         return ret;
     }
     
@@ -933,13 +935,17 @@ namespace _srs_internal
         time = ::time(NULL);
         version = 0x01000504; // server s1 version
         
+        SrsDH dh;
+        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 = openssl_generate_key(block0.key.key, 128)) != ERROR_SUCCESS) {
+            if ((ret = dh.copy_public_key((char*)block0.key.key, NULL)) != ERROR_SUCCESS) {
                 srs_error("calc s1 key failed. ret=%d", ret);
                 return ret;
             }
@@ -949,7 +955,7 @@ namespace _srs_internal
             
             // directly generate the public key.
             // @see: https://github.com/winlinvip/simple-rtmp-server/issues/148
-            if ((ret = openssl_generate_key(block1.key.key, 128)) != ERROR_SUCCESS) {
+            if ((ret = dh.copy_public_key((char*)block1.key.key, NULL)) != ERROR_SUCCESS) {
                 srs_error("calc s1 key failed. ret=%d", ret);
                 return ret;
             }

trunk/src/rtmp/srs_protocol_handshake.hpp

@@ -36,6 +36,9 @@ class SrsHandshakeBytes;
 
 #ifdef SRS_AUTO_SSL
 
+// for openssl.
+#include <openssl/hmac.h>
+
 namespace _srs_internal
 {
     /**
@@ -117,6 +120,46 @@ namespace _srs_internal
     int openssl_HMACsha256(const void* key, int key_size, const void* data, int data_size, void* digest);
     int openssl_generate_key(char* public_key, int32_t size);
     
+    /**
+    * the DH wrapper.
+    */
+    class SrsDH
+    {
+    private:
+        DH* pdh;
+    public:
+        SrsDH();
+        virtual ~SrsDH();
+    public:
+        /**
+        * initialize dh, generate the public and private key.
+        * @param ensure_128bytes_public_key whether ensure public key is 128bytes,
+        *       sometimes openssl generate 127bytes public key.
+        *       default to false to donot ensure.
+        */
+        virtual int initialize(bool ensure_128bytes_public_key = false);
+        /**
+        * 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.
+        * @remark, when ensure_128bytes_public_key, the size always 128.
+        */
+        virtual int copy_public_key(char* pkey, int32_t* ppkey_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.
+        */
+        virtual int copy_shared_key(const char* ppkey, int32_t ppkey_size, char* skey, int32_t* pskey_size);
+    private:
+        virtual int do_initialize();
+    };
+    
     // calc the offset of key,
     // the key->offset cannot be used as the offset of key.
     int srs_key_block_get_offset(key_block* key);

trunk/src/utest/srs_utest_protocol.cpp

@@ -238,11 +238,24 @@ VOID TEST(ProtocolHandshakeTest, OpensslSha256)
 // verify the dh key
 VOID TEST(ProtocolHandshakeTest, DHKey)
 {
+    _srs_internal::SrsDH dh;
+    
+    ASSERT_TRUE(ERROR_SUCCESS == dh.initialize(true));
+    
     char pub_key1[128];
-    openssl_generate_key(pub_key1, 128);
+    EXPECT_TRUE(ERROR_SUCCESS == dh.copy_public_key(pub_key1, NULL));
     
     char pub_key2[128];
-    openssl_generate_key(pub_key2, 128);
+    EXPECT_TRUE(ERROR_SUCCESS == dh.copy_public_key(pub_key2, NULL));
+    
+    EXPECT_TRUE(srs_bytes_equals(pub_key1, pub_key2, 128));
+    
+    // another dh
+    _srs_internal::SrsDH dh0;
+    
+    ASSERT_TRUE(ERROR_SUCCESS == dh0.initialize(true));
+    
+    EXPECT_TRUE(ERROR_SUCCESS == dh0.copy_public_key(pub_key2, NULL));
     
     EXPECT_FALSE(srs_bytes_equals(pub_key1, pub_key2, 128));
 }