Use std::unique_ptr for ARC4Encryptor in PeerConnection and MSEHandshake

src/InitiatorMSEHandshakeCommand.cc

@@ -153,13 +153,13 @@ bool InitiatorMSEHandshakeCommand::executeInternal() {
           (new PeerConnection(getCuid(), getPeer(), getSocket()));
         if(mseHandshake_->getNegotiatedCryptoType() ==
            MSEHandshake::CRYPTO_ARC4){
-          peerConnection->enableEncryption(mseHandshake_->getEncryptor(),
-                                           mseHandshake_->getDecryptor());
           size_t buflen = mseHandshake_->getBufferLength();
           mseHandshake_->getDecryptor()->encrypt(buflen,
                                                  mseHandshake_->getBuffer(),
                                                  mseHandshake_->getBuffer());
           peerConnection->presetBuffer(mseHandshake_->getBuffer(), buflen);
+          peerConnection->enableEncryption(mseHandshake_->popEncryptor(),
+                                           mseHandshake_->popDecryptor());
         } else {
           peerConnection->presetBuffer(mseHandshake_->getBuffer(),
                                        mseHandshake_->getBufferLength());

src/MSEHandshake.cc

@@ -193,7 +193,7 @@ void MSEHandshake::initCipher(const unsigned char* infoHash)
   sha1_->reset();
   message_digest::digest(localCipherKey, sizeof(localCipherKey),
                          sha1_.get(), s, sizeof(s));
-  encryptor_.reset(new ARC4Encryptor());
+  encryptor_ = make_unique<ARC4Encryptor>();
   encryptor_->init(localCipherKey, sizeof(localCipherKey));
 
   unsigned char peerCipherKey[20];
@@ -201,7 +201,7 @@ void MSEHandshake::initCipher(const unsigned char* infoHash)
   sha1_->reset();
   message_digest::digest(peerCipherKey, sizeof(peerCipherKey),
                          sha1_.get(), s, sizeof(s));
-  decryptor_.reset(new ARC4Encryptor());
+  decryptor_ = make_unique<ARC4Encryptor>();
   decryptor_->init(peerCipherKey, sizeof(peerCipherKey));
 
   // discard first 1024 bytes ARC4 output.
@@ -583,4 +583,14 @@ bool MSEHandshake::getWantWrite() const
   return !socketBuffer_.sendBufferIsEmpty();
 }
 
+std::unique_ptr<ARC4Encryptor> MSEHandshake::popEncryptor()
+{
+  return std::move(encryptor_);
+}
+
+std::unique_ptr<ARC4Encryptor> MSEHandshake::popDecryptor()
+{
+  return std::move(decryptor_);
+}
+
 } // namespace aria2

src/MSEHandshake.h

@@ -88,8 +88,8 @@ private:
 
   CRYPTO_TYPE negotiatedCryptoType_;
   DHKeyExchange* dh_;
-  std::shared_ptr<ARC4Encryptor> encryptor_;
+  std::unique_ptr<ARC4Encryptor> encryptor_;
-  std::shared_ptr<ARC4Encryptor> decryptor_;
+  std::unique_ptr<ARC4Encryptor> decryptor_;
   unsigned char infoHash_[INFO_HASH_LENGTH];
   unsigned char secret_[KEY_LENGTH];
   bool initiator_;
@@ -197,16 +197,20 @@ public:
     return negotiatedCryptoType_;
   }
 
-  const std::shared_ptr<ARC4Encryptor>& getEncryptor() const
+  const std::unique_ptr<ARC4Encryptor>& getEncryptor() const
   {
     return encryptor_;
   }
 
-  const std::shared_ptr<ARC4Encryptor>& getDecryptor() const
+  const std::unique_ptr<ARC4Encryptor>& getDecryptor() const
   {
     return decryptor_;
   }
 
+  std::unique_ptr<ARC4Encryptor> popEncryptor();
+
+  std::unique_ptr<ARC4Encryptor> popDecryptor();
+
   const unsigned char* getBuffer() const
   {
     return rbuf_;

src/PeerConnection.cc

@@ -232,11 +232,11 @@ void PeerConnection::readData
 }
 
 void PeerConnection::enableEncryption
-(const std::shared_ptr<ARC4Encryptor>& encryptor,
+(std::unique_ptr<ARC4Encryptor> encryptor,
- const std::shared_ptr<ARC4Encryptor>& decryptor)
+ std::unique_ptr<ARC4Encryptor> decryptor)
 {
-  encryptor_ = encryptor;
+  encryptor_ = std::move(encryptor);
-  decryptor_ = decryptor;
+  decryptor_ = std::move(decryptor);
 
   encryptionEnabled_ = true;
 }

src/PeerConnection.h

@@ -77,8 +77,8 @@ private:
   SocketBuffer socketBuffer_;
 
   bool encryptionEnabled_;
-  std::shared_ptr<ARC4Encryptor> encryptor_;
+  std::unique_ptr<ARC4Encryptor> encryptor_;
-  std::shared_ptr<ARC4Encryptor> decryptor_;
+  std::unique_ptr<ARC4Encryptor> decryptor_;
 
   bool prevPeek_;
 
@@ -111,8 +111,8 @@ public:
   bool receiveHandshake
   (unsigned char* data, size_t& dataLength, bool peek = false);
 
-  void enableEncryption(const std::shared_ptr<ARC4Encryptor>& encryptor,
+  void enableEncryption(std::unique_ptr<ARC4Encryptor> encryptor,
-                        const std::shared_ptr<ARC4Encryptor>& decryptor);
+                        std::unique_ptr<ARC4Encryptor> decryptor);
 
   void presetBuffer(const unsigned char* data, size_t length);
 

src/ReceiverMSEHandshakeCommand.cc

@@ -209,8 +209,8 @@ void ReceiverMSEHandshakeCommand::createCommand()
   std::shared_ptr<PeerConnection> peerConnection
     (new PeerConnection(getCuid(), getPeer(), getSocket()));
   if(mseHandshake_->getNegotiatedCryptoType() == MSEHandshake::CRYPTO_ARC4) {
-    peerConnection->enableEncryption(mseHandshake_->getEncryptor(),
+    peerConnection->enableEncryption(mseHandshake_->popEncryptor(),
-                                     mseHandshake_->getDecryptor());
+                                     mseHandshake_->popDecryptor());
   }
   // Since initiator cannot send payload stream before reading step2
   // from receiver, mseHandshake_->getBufferLength() should be 0.