[5/7] rtmp: use the new blowfish/rc4/xtea alloc functions

Message ID 1438032992-4680-6-git-send-email-jamrial@gmail.com
State New
Headers show

Commit Message

James Almer July 27, 2015, 9:36 p.m.
Signed-off-by: James Almer <jamrial@gmail.com>
---
 libavformat/rtmpcrypt.c | 66 +++++++++++++++++++++++++++++++------------------
 libavformat/rtmpcrypt.h |  5 ++--
 libavformat/rtmpproto.c | 15 ++++++-----
 3 files changed, 54 insertions(+), 32 deletions(-)

Patch

diff --git a/libavformat/rtmpcrypt.c b/libavformat/rtmpcrypt.c
index 887427d..9144938 100644
--- a/libavformat/rtmpcrypt.c
+++ b/libavformat/rtmpcrypt.c
@@ -28,6 +28,7 @@ 
 
 #include "libavutil/blowfish.h"
 #include "libavutil/intreadwrite.h"
+#include "libavutil/mem.h"
 #include "libavutil/opt.h"
 #include "libavutil/rc4.h"
 #include "libavutil/xtea.h"
@@ -43,8 +44,8 @@  typedef struct RTMPEContext {
     const AVClass *class;
     URLContext   *stream;            ///< TCP stream
     FF_DH        *dh;                ///< Diffie-Hellman context
-    struct AVRC4 key_in;             ///< RC4 key used for decrypt data
-    struct AVRC4 key_out;            ///< RC4 key used for encrypt data
+    struct AVRC4 *key_in;            ///< RC4 key used for decrypt data
+    struct AVRC4 *key_out;           ///< RC4 key used for encrypt data
     int          handshaked;         ///< flag indicating when the handshake is performed
     int          tunneling;          ///< use a HTTP connection (RTMPTE)
 } RTMPEContext;
@@ -170,54 +171,69 @@  int ff_rtmpe_compute_secret_key(URLContext *h, const uint8_t *serverdata,
     if ((ret = ff_rtmp_calc_digest(serverdata + server_pos, 128, 0, secret_key,
                                    128, digest)) < 0)
         return ret;
-    av_rc4_init(&rt->key_out, digest, 16 * 8, 1);
+
+    if (!(rt->key_out = av_rc4_alloc()))
+        return AVERROR(ENOMEM);
+    av_rc4_init(rt->key_out, digest, 16 * 8, 1);
 
     /* set input key */
     if ((ret = ff_rtmp_calc_digest(clientdata + client_pos, 128, 0, secret_key,
                                    128, digest)) < 0)
         return ret;
-    av_rc4_init(&rt->key_in, digest, 16 * 8, 1);
+
+    if (!(rt->key_in = av_rc4_alloc()))
+        return AVERROR(ENOMEM);
+    av_rc4_init(rt->key_in, digest, 16 * 8, 1);
 
     return 0;
 }
 
-static void rtmpe8_sig(const uint8_t *in, uint8_t *out, int key_id)
+static void rtmpe8_sig(struct AVXTEA *ctx, const uint8_t *in, uint8_t *out, int key_id)
 {
-    struct AVXTEA ctx;
-
-    av_xtea_init(&ctx, rtmpe8_keys[key_id]);
-    av_xtea_crypt(&ctx, out, in, 1, NULL, 0);
+    av_xtea_init(ctx, rtmpe8_keys[key_id]);
+    av_xtea_crypt(ctx, out, in, 1, NULL, 0);
 }
 
-static void rtmpe9_sig(const uint8_t *in, uint8_t *out, int key_id)
+static void rtmpe9_sig(struct AVBlowfish *ctx, const uint8_t *in, uint8_t *out, int key_id)
 {
-    struct AVBlowfish ctx;
     uint32_t xl, xr;
 
     xl = AV_RL32(in);
     xr = AV_RL32(in + 4);
 
-    av_blowfish_init(&ctx, rtmpe9_keys[key_id], 24);
-    av_blowfish_crypt_ecb(&ctx, &xl, &xr, 0);
+    av_blowfish_init(ctx, rtmpe9_keys[key_id], 24);
+    av_blowfish_crypt_ecb(ctx, &xl, &xr, 0);
 
     AV_WL32(out, xl);
     AV_WL32(out + 4, xr);
 }
 
-void ff_rtmpe_encrypt_sig(URLContext *h, uint8_t *sig, const uint8_t *digest,
+int ff_rtmpe_encrypt_sig(URLContext *h, uint8_t *sig, const uint8_t *digest,
                           int type)
 {
     int i;
 
-    for (i = 0; i < 32; i += 8) {
-        if (type == 8) {
+    if (type == 8) {
+        struct AVXTEA *ctx = av_xtea_alloc();
+
+        if (!ctx)
+            return AVERROR(ENOMEM);
+        for (i = 0; i < 32; i += 8)
             /* RTMPE type 8 uses XTEA on the signature */
-            rtmpe8_sig(sig + i, sig + i, digest[i] % 15);
-        } else if (type == 9) {
+            rtmpe8_sig(ctx, sig + i, sig + i, digest[i] % 15);
+        av_free(ctx);
+    } else if (type == 9) {
+        struct AVBlowfish *ctx = av_blowfish_alloc();
+
+        if (!ctx)
+            return AVERROR(ENOMEM);
+        for (i = 0; i < 32; i += 8)
             /* RTMPE type 9 uses Blowfish on the signature */
-            rtmpe9_sig(sig + i, sig + i, digest[i] % 15);
-        }
+            rtmpe9_sig(ctx, sig + i, sig + i, digest[i] % 15);
+        av_free(ctx);
     }
+
+    return 0;
 }
 
 int ff_rtmpe_update_keystream(URLContext *h)
@@ -226,8 +242,8 @@  int ff_rtmpe_update_keystream(URLContext *h)
     char buf[RTMP_HANDSHAKE_PACKET_SIZE];
 
     /* skip past 1536 bytes of the RC4 bytestream */
-    av_rc4_crypt(&rt->key_in, buf, NULL, sizeof(buf), NULL, 1);
-    av_rc4_crypt(&rt->key_out, buf, NULL, sizeof(buf), NULL, 1);
+    av_rc4_crypt(rt->key_in, buf, NULL, sizeof(buf), NULL, 1);
+    av_rc4_crypt(rt->key_out, buf, NULL, sizeof(buf), NULL, 1);
 
     /* the next requests will be encrypted using RC4 keys */
     rt->handshaked = 1;
@@ -239,6 +255,8 @@  static int rtmpe_close(URLContext *h)
 {
     RTMPEContext *rt = h->priv_data;
 
+    av_free(rt->key_in);
+    av_free(rt->key_out);
     ff_dh_free(rt->dh);
     ffurl_close(rt->stream);
 
@@ -287,7 +305,7 @@  static int rtmpe_read(URLContext *h, uint8_t *buf, int size)
 
     if (rt->handshaked && ret > 0) {
         /* decrypt data received by the server */
-        av_rc4_crypt(&rt->key_in, buf, buf, ret, NULL, 1);
+        av_rc4_crypt(rt->key_in, buf, buf, ret, NULL, 1);
     }
 
     return ret;
@@ -300,7 +318,7 @@  static int rtmpe_write(URLContext *h, const uint8_t *buf, int size)
 
     if (rt->handshaked) {
         /* encrypt data to send to the server */
-        av_rc4_crypt(&rt->key_out, buf, buf, size, NULL, 1);
+        av_rc4_crypt(rt->key_out, buf, buf, size, NULL, 1);
     }
 
     if ((ret = ffurl_write(rt->stream, buf, size)) < 0)
diff --git a/libavformat/rtmpcrypt.h b/libavformat/rtmpcrypt.h
index 2799433..7f8f09e 100644
--- a/libavformat/rtmpcrypt.h
+++ b/libavformat/rtmpcrypt.h
@@ -54,9 +54,10 @@  int ff_rtmpe_compute_secret_key(URLContext *h, const uint8_t *serverdata,
  * @param signature     the signature to encrypt
  * @param digest        the digest used for finding the encryption key
  * @param type          type of encryption (8 for XTEA, 9 for Blowfish)
+ * @return zero on success, negative value otherwise
  */
-void ff_rtmpe_encrypt_sig(URLContext *h, uint8_t *signature,
-                          const uint8_t *digest, int type);
+int ff_rtmpe_encrypt_sig(URLContext *h, uint8_t *signature,
+                         const uint8_t *digest, int type);
 
 /**
  * Update the keystream and set RC4 keys for encryption.
diff --git a/libavformat/rtmpproto.c b/libavformat/rtmpproto.c
index e75a519..78b353c 100644
--- a/libavformat/rtmpproto.c
+++ b/libavformat/rtmpproto.c
@@ -1322,7 +1322,8 @@  static int rtmp_handshake(URLContext *s, RTMPContext *rt)
                 return ret;
 
             /* Encrypt the signature received by the server. */
-            ff_rtmpe_encrypt_sig(rt->stream, signature, digest, serverdata[0]);
+            if ((ret = ff_rtmpe_encrypt_sig(rt->stream, signature, digest, serverdata[0])) < 0)
+                return ret;
         }
 
         if (memcmp(signature, clientdata + RTMP_HANDSHAKE_PACKET_SIZE - 32, 32)) {
@@ -1346,9 +1347,10 @@  static int rtmp_handshake(URLContext *s, RTMPContext *rt)
 
         if (CONFIG_FFRTMPCRYPT_PROTOCOL && rt->encrypted) {
             /* Encrypt the signature to be send to the server. */
-            ff_rtmpe_encrypt_sig(rt->stream, tosend +
-                                 RTMP_HANDSHAKE_PACKET_SIZE - 32, digest,
-                                 serverdata[0]);
+            if ((ret = ff_rtmpe_encrypt_sig(rt->stream, tosend +
+                                            RTMP_HANDSHAKE_PACKET_SIZE - 32, digest,
+                                            serverdata[0])) < 0)
+                return ret;
         }
 
         // write reply back to the server
@@ -1371,8 +1373,9 @@  static int rtmp_handshake(URLContext *s, RTMPContext *rt)
 
             if (serverdata[0] == 9) {
                 /* Encrypt the signature received by the server. */
-                ff_rtmpe_encrypt_sig(rt->stream, signature, digest,
-                                     serverdata[0]);
+                if ((ret = ff_rtmpe_encrypt_sig(rt->stream, signature, digest,
+                                                serverdata[0])) < 0)
+                    return ret;
             }
         }