From b7af94e2e4e2a1108cbb010338e775e217012443 Mon Sep 17 00:00:00 2001
From: levlam <levlam@telegram.org>
Date: Thu, 13 Feb 2020 03:49:59 +0300
Subject: [PATCH] Make RSA::encrypt safe.

GitOrigin-RevId: 0d83acb2f6c022af59320c3ea755257cd926cbe4
---
 td/mtproto/Handshake.cpp | 3 ++-
 td/mtproto/crypto.cpp    | 4 +++-
 td/mtproto/crypto.h      | 2 +-
 3 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/td/mtproto/Handshake.cpp b/td/mtproto/Handshake.cpp
index 2c8c2f66..feda2e10 100644
--- a/td/mtproto/Handshake.cpp
+++ b/td/mtproto/Handshake.cpp
@@ -109,7 +109,8 @@ Status AuthKeyHandshake::on_res_pq(Slice message, Callback *connection, PublicRs
   // encrypted_data := RSA (data_with_hash, server_public_key); a 255-byte long number (big endian)
   //   is raised to the requisite power over the requisite modulus, and the result is stored as a 256-byte number.
   string encrypted_data(256, 0);
-  rsa.encrypt(data_with_hash, size, reinterpret_cast<unsigned char *>(&encrypted_data[0]));
+  rsa.encrypt(data_with_hash, size, sizeof(data_with_hash), reinterpret_cast<unsigned char *>(&encrypted_data[0]),
+              encrypted_data.size());
 
   // req_DH_params#d712e4be nonce:int128 server_nonce:int128 p:string q:string public_key_fingerprint:long
   // encrypted_data:string = Server_DH_Params
diff --git a/td/mtproto/crypto.cpp b/td/mtproto/crypto.cpp
index 596a9263..487cfe97 100644
--- a/td/mtproto/crypto.cpp
+++ b/td/mtproto/crypto.cpp
@@ -95,13 +95,15 @@ size_t RSA::size() const {
   return 256;
 }
 
-size_t RSA::encrypt(unsigned char *from, size_t from_len, unsigned char *to) const {
+size_t RSA::encrypt(unsigned char *from, size_t from_len, size_t max_from_len, unsigned char *to, size_t to_len) const {
   CHECK(from_len > 0 && from_len <= 2550);
   size_t pad = (25500 - from_len - 32) % 255 + 32;
   size_t chunks = (from_len + pad) / 255;
   int bits = n_.get_num_bits();
   CHECK(bits >= 2041 && bits <= 2048);
   CHECK(chunks * 255 == from_len + pad);
+  CHECK(from_len + pad <= max_from_len);
+  CHECK(chunks * 256 <= to_len);
   Random::secure_bytes(from + from_len, pad);
 
   BigNumContext ctx;
diff --git a/td/mtproto/crypto.h b/td/mtproto/crypto.h
index 119815ab..dcf4b0a2 100644
--- a/td/mtproto/crypto.h
+++ b/td/mtproto/crypto.h
@@ -21,7 +21,7 @@ class RSA {
   RSA clone() const;
   int64 get_fingerprint() const;
   size_t size() const;
-  size_t encrypt(unsigned char *from, size_t from_len, unsigned char *to) const;
+  size_t encrypt(unsigned char *from, size_t from_len, size_t max_from_len, unsigned char *to, size_t to_len) const;
 
   void decrypt(Slice from, MutableSlice to) const;