From c0d2223c778ae9719fd13d1fde6dc2fbf4faa82b Mon Sep 17 00:00:00 2001
From: Tatsuhiro Tsujikawa <tatsuhiro.t@gmail.com>
Date: Sun, 17 Feb 2008 15:56:47 +0000
Subject: 2008-02-16  Tatsuhiro Tsujikawa  <tujikawa at rednoah dot com>

	Added Message Stream Encryption(MSE) support.
	Currently, aria2 accepts incoming connections with Obfuscation Header
	and legacy BitTorrent Header and establishes connections with
	Obfuscation Header first and if failed then retry with legacy
	BitTorrent header. If plain text and ARC4 is provided, aria2 always
	choose ARC4. The new option to change the default behavior is planned.
	For tracker extension, "supportcrypto=1" is added statically.
	* src/PeerInitiateConnectionCommand.{h, cc}
	* src/PeerConnection.{h, cc}
	* src/HandleRegistry.h
	* src/SocketCore.h
	* src/PeerReceiveHandshakeCommand.{h, cc}
	* src/BtRegistry.{h, cc}
	* src/PeerListenCommand.cc
	* src/InitiatorMSEHandshakeCommand.{h, cc}
	* src/ReceiverMSEHandshakeCommand.{h, cc}
	* src/MSEHandshake.{h, cc}
	* src/ARC4Encryptor.h
	* src/ARC4Decryptor.h
	* src/LibgcryptARC4Encryptor.h
	* src/LibgcryptARC4Decryptor.h
	* src/LibgcryptARC4Context.h
	* src/LibsslARC4Encryptor.h
	* src/LibsslARC4Decryptor.h
	* src/LibsslARC4Context.h
	* src/DHKeyExchange.h
	* src/LibgcryptDHKeyExchange.h
	* src/LibsslDHKeyExchange.h
	* src/DefaultBtAnnounce.cc: Just added "supportcrypto=1" parameter.
	* test/DefaultBtAnnounceTest.cc
	* test/ARC4Test.cc
	* test/DHKeyExchangeTest.cc

	Removed prepareForRetry() because it is not used.
	* src/PeerAbstractCommand.{h, cc}
	* src/PeerInteractionCommand.{h, cc}
	* src/PeerInitiateConnectionCommand.{h, cc}
---
 src/PeerConnection.cc | 64 +++++++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 59 insertions(+), 5 deletions(-)

(limited to 'src/PeerConnection.cc')

diff --git a/src/PeerConnection.cc b/src/PeerConnection.cc
index 83c6737d..e528ec1c 100644
--- a/src/PeerConnection.cc
+++ b/src/PeerConnection.cc
@@ -40,7 +40,11 @@
 #include "BtHandshakeMessage.h"
 #include "Socket.h"
 #include "a2netcompat.h"
+#include "ARC4Encryptor.h"
+#include "ARC4Decryptor.h"
 #include <cstring>
+#include <cassert>
+#include <algorithm>
 
 namespace aria2 {
 
@@ -53,7 +57,10 @@ PeerConnection::PeerConnection(int32_t cuid,
    logger(LogFactory::getInstance()),
    resbufLength(0),
    currentPayloadLength(0),
-   lenbufLength(0)
+   lenbufLength(0),
+   _encryptionEnabled(false),
+   _encryptor(0),
+   _decryptor(0)
 {}
 
 PeerConnection::~PeerConnection() {}
@@ -62,7 +69,7 @@ int32_t PeerConnection::sendMessage(const unsigned char* data, int32_t dataLengt
   int32_t writtenLength = 0;
   if(socket->isWritable(0)) {
     // TODO fix this
-    socket->writeData((const char*)data, dataLength);
+    sendData(data, dataLength, _encryptionEnabled);
     writtenLength += dataLength;
   }
   return writtenLength;
@@ -76,7 +83,7 @@ bool PeerConnection::receiveMessage(unsigned char* data, int32_t& dataLength) {
     // read payload size, 32bit unsigned integer
     int32_t remaining = 4-lenbufLength;
     int32_t temp = remaining;
-    socket->readData(lenbuf+lenbufLength, remaining);
+    readData(lenbuf+lenbufLength, remaining, _encryptionEnabled);
     if(remaining == 0) {
       // we got EOF
       logger->debug("CUID#%d - In PeerConnection::receiveMessage(), remain=%d",
@@ -101,7 +108,7 @@ bool PeerConnection::receiveMessage(unsigned char* data, int32_t& dataLength) {
   int32_t remaining = currentPayloadLength-resbufLength;
   int32_t temp = remaining;
   if(remaining > 0) {
-    socket->readData(resbuf+resbufLength, remaining);
+    readData(resbuf+resbufLength, remaining, _encryptionEnabled);
     if(remaining == 0) {
       // we got EOF
       logger->debug("CUID#%d - In PeerConnection::receiveMessage(), payloadlen=%d, remaining=%d",
@@ -132,7 +139,7 @@ bool PeerConnection::receiveHandshake(unsigned char* data, int32_t& dataLength,
   bool retval = true;
   if(remaining > 0) {
     int32_t temp = remaining;
-    socket->readData(resbuf+resbufLength, remaining);
+    readData(resbuf+resbufLength, remaining, _encryptionEnabled);
     if(remaining == 0) {
       // we got EOF
       logger->debug("CUID#%d - In PeerConnection::receiveHandshake(), remain=%d",
@@ -153,4 +160,51 @@ bool PeerConnection::receiveHandshake(unsigned char* data, int32_t& dataLength,
   return retval;
 }
 
+void PeerConnection::readData(char* data, int32_t& length, bool encryption)
+{
+  if(encryption) {
+    unsigned char* cdata = reinterpret_cast<unsigned char*>(data);
+    unsigned char temp[MAX_PAYLOAD_LEN];
+    assert(MAX_PAYLOAD_LEN >= length);
+    socket->readData(temp, length);
+    _decryptor->decrypt(cdata, length, temp, length);
+  } else {
+    socket->readData(data, length);
+  }
+}
+
+void PeerConnection::sendData(const unsigned char* data, size_t length, bool encryption)
+{
+  if(encryption) {
+    unsigned char temp[4096];
+    const unsigned char* dptr = data;
+    size_t r = length;
+    while(r > 0) {
+      size_t s = std::min(r, sizeof(temp));
+      _encryptor->encrypt(temp, s, dptr, s);
+      socket->writeData(temp, s);
+      dptr += s;
+      r -= s;
+    }
+  } else {
+    socket->writeData(data, length);
+  }
+}
+
+void PeerConnection::enableEncryption(const SharedHandle<ARC4Encryptor>& encryptor,
+				      const SharedHandle<ARC4Decryptor>& decryptor)
+{
+  _encryptor = encryptor;
+  _decryptor = decryptor;
+
+  _encryptionEnabled = true;
+}
+
+void PeerConnection::presetBuffer(const unsigned char* data, size_t length)
+{
+  size_t nwrite = std::min((size_t)MAX_PAYLOAD_LEN, length);
+  memcpy(resbuf, data, nwrite);
+  resbufLength = length;
+}
+
 } // namespace aria2
-- 
cgit