6 files changed, 176 insertions, 17 deletions

diff --git a/src/Makefile b/src/Makefile
index 5ded7d6..88dc08f 100644
--- a/src/Makefile
+++ b/src/Makefile
@@ -139,3 +139,4 @@ tests: test_bip32 bitflash-node
         ./test_bip32
         ./bitflash-node -selftest=wallet-keypool
         ./bitflash-node -selftest=wallet-hd
+        ./bitflash-node -selftest=net-message
diff --git a/src/main.cpp b/src/main.cpp
index 9082bd5..9bb3c24 100644
--- a/src/main.cpp
+++ b/src/main.cpp
@@ -181,6 +181,39 @@ static bool AddKeyIfMissing(const CKey& key)
     return AddKey(key);
 }
 
+static bool DeriveNextHDWalletKey(vector<unsigned char>& vchPubKeyRet, string& strErrorRet)
+{
+    vchPubKeyRet.clear();
+    CRITICAL_BLOCK(cs_keyPool)
+    {
+        if (!HaveHDSeed())
+        {
+            strErrorRet = "no seed";
+            return false;
+        }
+
+        CKey key;
+        if (!DeriveHDKey(nHDNext, key, strErrorRet))
+            return false;
+        if (!AddKeyIfMissing(key))
+        {
+            strErrorRet = "could not write the derived key to wallet.dat";
+            return false;
+        }
+
+        unsigned int nNext = nHDNext + 1;
+        if (!CWalletDB().WriteHDNext(nNext))
+        {
+            strErrorRet = "could not advance the derivation counter";
+            return false;
+        }
+
+        nHDNext = nNext;
+        vchPubKeyRet = key.GetPubKey();
+    }
+    return true;
+}
+
 vector<unsigned char> GenerateNewKey()
 {
     CKey key;
@@ -2310,27 +2343,55 @@ bool ProcessMessages(CNode* pfrom)
             if (LogAcceptsCategory("net")) printf("\n\nPROCESSMESSAGE SKIPPED %d BYTES\n\n", (int)(pstart - vRecv.begin()));
         vRecv.erase(vRecv.begin(), pstart);
 
-        // Read header
+        // Read the header from a copy first. The old path consumed the header,
+        // noticed the payload was incomplete, inserted the header back at the
+        // front of the receive buffer, then slept inside the shared message
+        // thread. A peer could announce a large payload and drip bytes forever,
+        // making every pass pay for an O(n) insert and a 100 ms stall.
         CMessageHeader hdr;
-        vRecv >> hdr;
+        CDataStream vHeader(vRecv.begin(), vRecv.begin() + sizeof(CMessageHeader),
+                            vRecv.nType, vRecv.nVersion);
+        vHeader >> hdr;
         if (!hdr.IsValid())
         {
             if (LogAcceptsCategory("net")) printf("\n\nPROCESSMESSAGE: ERRORS IN HEADER %s\n\n\n", hdr.GetCommand().c_str());
+            if (hdr.nMessageSize > MAX_PROTOCOL_MESSAGE_SIZE)
+                pfrom->fDisconnect = true;
+            pfrom->nIncompleteMessageStart = 0;
+            vRecv.ignore(sizeof(CMessageHeader));
             continue;
         }
         string strCommand = hdr.GetCommand();
 
         // Message size
         unsigned int nMessageSize = hdr.nMessageSize;
-        if (nMessageSize > vRecv.size())
+        unsigned int nPayloadAvailable = vRecv.size() - sizeof(CMessageHeader);
+        if (nMessageSize > nPayloadAvailable)
         {
-            // Rewind and wait for rest of message
-            ///// need a mechanism to give up waiting for overlong message size error
-            if (LogAcceptsCategory("net")) printf("MESSAGE-BREAK 2\n");
-            vRecv.insert(vRecv.begin(), BEGIN(hdr), END(hdr));
-            Sleep(100);
+            if (pfrom->nIncompleteMessageStart == 0 ||
+                pfrom->nIncompleteMessageSize != nMessageSize ||
+                pfrom->strIncompleteMessageCommand != strCommand)
+            {
+                pfrom->nIncompleteMessageStart = GetTime();
+                pfrom->nIncompleteMessageSize = nMessageSize;
+                pfrom->strIncompleteMessageCommand = strCommand;
+            }
+            else if (GetTime() - pfrom->nIncompleteMessageStart >
+                     BTF_INCOMPLETE_MESSAGE_TIMEOUT_SECS)
+            {
+                LogPrint("net",
+                         "disconnecting %s: incomplete %s message, %u of %u bytes after %d seconds\n",
+                         pfrom->addr.ToString().c_str(), strCommand.c_str(),
+                         nPayloadAvailable, nMessageSize,
+                         BTF_INCOMPLETE_MESSAGE_TIMEOUT_SECS);
+                pfrom->fDisconnect = true;
+            }
             break;
         }
+        pfrom->nIncompleteMessageStart = 0;
+        pfrom->nIncompleteMessageSize = 0;
+        pfrom->strIncompleteMessageCommand.clear();
+        vRecv.ignore(sizeof(CMessageHeader));
 
         // Copy message to its own buffer
         CDataStream vMsg(vRecv.begin(), vRecv.begin() + nMessageSize, vRecv.nType, vRecv.nVersion);
@@ -3774,13 +3835,28 @@ bool CreateTransaction(CScript scriptPubKey, int64 nValue, CWalletTx& wtxNew, in
                 // Fill vout[1] back to self with any change
                 if (nValueIn > nTotalValue)
                 {
-                    // Use the same key as one of the coins
                     vector<unsigned char> vchPubKey;
-                    CTransaction& txFirst = *(*setCoins.begin());
-                    foreach(const CTxOut& txout, txFirst.vout)
-                        if (txout.IsMine())
-                            if (ExtractPubKey(txout.scriptPubKey, true, vchPubKey))
-                                break;
+                    if (HaveHDSeed())
+                    {
+                        string strError;
+                        if (!DeriveNextHDWalletKey(vchPubKey, strError))
+                        {
+                            printf("CreateTransaction() : could not derive a phrase-backed change key: %s\n",
+                                   strError.c_str());
+                            return false;
+                        }
+                    }
+                    else
+                    {
+                        // Legacy wallets have no deterministic seed. Preserve
+                        // the 0.1.0 behaviour for them: return change to a key
+                        // already proven to belong to this wallet.
+                        CTransaction& txFirst = *(*setCoins.begin());
+                        foreach(const CTxOut& txout, txFirst.vout)
+                            if (txout.IsMine())
+                                if (ExtractPubKey(txout.scriptPubKey, true, vchPubKey))
+                                    break;
+                    }
                     if (vchPubKey.empty())
                         return false;
 
diff --git a/src/makefile.mingw b/src/makefile.mingw
index 5be3704..5025acb 100644
--- a/src/makefile.mingw
+++ b/src/makefile.mingw
@@ -111,3 +111,4 @@ tests: test_bip32.exe bitflash.exe
         ./test_bip32.exe
         ./bitflash.exe -selftest=wallet-keypool
         ./bitflash.exe -selftest=wallet-hd
+        ./bitflash.exe -selftest=net-message
diff --git a/src/net.h b/src/net.h
index 1c0c24e..37f335e 100644
--- a/src/net.h
+++ b/src/net.h
@@ -67,6 +67,13 @@ static const int          BTF_RECV_TIMEOUT_SECS    = 30 * 60;
 static const int          BTF_SEND_STALL_SECS      = 10 * 60;
 static const int          BTF_HANDSHAKE_GRACE_SECS = 60;
 
+// A serialized object cannot be larger than MAX_SIZE, so no valid network
+// message needs the old 256 MB header allowance. Keep the wire cap aligned with
+// what the deserializer can accept and bound how long one peer may sit on a
+// partial payload after a complete header has arrived.
+static const unsigned int MAX_PROTOCOL_MESSAGE_SIZE = MAX_SIZE;
+static const int          BTF_INCOMPLETE_MESSAGE_TIMEOUT_SECS = 2 * 60;
+
 // Ceiling on simultaneous connections.
 //
 // ThreadSocketHandler watches every peer through one select(), and select()
@@ -225,7 +232,7 @@ public:
         }
 
         // Message size
-        if (nMessageSize > 0x10000000)
+        if (nMessageSize > MAX_PROTOCOL_MESSAGE_SIZE)
         {
             if (LogAcceptsCategory("net")) printf("CMessageHeader::IsValid() : nMessageSize too large %u\n", nMessageSize);
             return false;
@@ -596,6 +603,9 @@ public:
     int64 nLastSend;
     int64 nLastRecv;
     int64 nLastSendEmpty;
+    int64 nIncompleteMessageStart;
+    unsigned int nIncompleteMessageSize;
+    string strIncompleteMessageCommand;
 
     // Height this peer announced in its version message, or -1 if it sent a
     // version message that predates the field. Knowing how far along everyone
@@ -626,6 +636,9 @@ public:
         nLastSend = 0;
         nLastRecv = 0;
         nLastSendEmpty = GetTime();
+        nIncompleteMessageStart = 0;
+        nIncompleteMessageSize = 0;
+        strIncompleteMessageCommand.clear();
         nStartingHeight = -1;
         vfSubscribe.assign(256, false);
 
diff --git a/src/selftest.cpp b/src/selftest.cpp
index 88e9863..c04adf4 100644
--- a/src/selftest.cpp
+++ b/src/selftest.cpp
@@ -470,6 +470,67 @@ static int RunWalletHDSelfTest()
     return nFail == 0 ? 0 : 1;
 }
 
+static int RunNetMessageSelfTest()
+{
+    fflush(stdout);
+    printf("net-message self-test\n");
+
+    int nFail = 0;
+    try
+    {
+        CNode complete(INVALID_SOCKET, CAddress("127.0.0.1"));
+        complete.nVersion = VERSION;
+        complete.vRecv << CMessageHeader("ping", 0);
+        nFail += Check(ProcessMessages(&complete), "a complete empty message processes") ? 0 : 1;
+        nFail += Check(complete.vRecv.empty(), "a complete message is consumed") ? 0 : 1;
+        nFail += Check(!complete.fDisconnect, "a valid ping does not disconnect the peer") ? 0 : 1;
+
+        CNode partial(INVALID_SOCKET, CAddress("127.0.0.1"));
+        partial.nVersion = VERSION;
+        partial.vRecv << CMessageHeader("block", 100);
+        unsigned int nPartialBefore = partial.vRecv.size();
+        nFail += Check(ProcessMessages(&partial), "an incomplete message returns cleanly") ? 0 : 1;
+        nFail += Check(partial.vRecv.size() == nPartialBefore,
+                       "an incomplete message keeps one header in the buffer") ? 0 : 1;
+        nFail += Check(partial.nIncompleteMessageStart != 0,
+                       "an incomplete message starts a timeout clock") ? 0 : 1;
+        nFail += Check(!partial.fDisconnect,
+                       "a fresh incomplete message does not disconnect immediately") ? 0 : 1;
+
+        CNode stale(INVALID_SOCKET, CAddress("127.0.0.1"));
+        stale.nVersion = VERSION;
+        stale.vRecv << CMessageHeader("block", 100);
+        stale.nIncompleteMessageStart = GetTime() - BTF_INCOMPLETE_MESSAGE_TIMEOUT_SECS - 1;
+        stale.nIncompleteMessageSize = 100;
+        stale.strIncompleteMessageCommand = "block";
+        ProcessMessages(&stale);
+        nFail += Check(stale.fDisconnect,
+                       "a stale incomplete message disconnects the peer") ? 0 : 1;
+
+        CNode oversized(INVALID_SOCKET, CAddress("127.0.0.1"));
+        oversized.nVersion = VERSION;
+        oversized.vRecv << CMessageHeader("block", MAX_PROTOCOL_MESSAGE_SIZE + 1);
+        ProcessMessages(&oversized);
+        nFail += Check(oversized.fDisconnect,
+                       "an oversized message header disconnects the peer") ? 0 : 1;
+    }
+    catch (const std::exception& e)
+    {
+        printf("  FAIL exception: %s\n", e.what());
+        nFail++;
+    }
+    catch (...)
+    {
+        printf("  FAIL unknown exception\n");
+        nFail++;
+    }
+
+    printf("%s (%d failure%s)\n", nFail == 0 ? "ALL TESTS PASSED" : "TESTS FAILED",
+           nFail, nFail == 1 ? "" : "s");
+    fflush(stdout);
+    return nFail == 0 ? 0 : 1;
+}
+
 int RunSelfTest(const std::string& name)
 {
     AttachTerminal();
@@ -478,8 +539,10 @@ int RunSelfTest(const std::string& name)
         return RunWalletKeyPoolSelfTest();
     if (name == "wallet-hd")
         return RunWalletHDSelfTest();
+    if (name == "net-message")
+        return RunNetMessageSelfTest();
 
     printf("Unknown self-test '%s'\n", name.c_str());
-    printf("Known self-tests: wallet-keypool, wallet-hd\n");
+    printf("Known self-tests: wallet-keypool, wallet-hd, net-message\n");
     return 1;
 }
diff --git a/src/walletcmd.cpp b/src/walletcmd.cpp
index 8eb3592..fae2095 100644
--- a/src/walletcmd.cpp
+++ b/src/walletcmd.cpp
@@ -25,6 +25,10 @@
 // would restore a wallet that looks empty.
 static const int RESTORE_BATCH   = 100;
 static const int RESTORE_MAX     = 10000;
+// The wallet may have a full unused key pool in front of a change address
+// created while spending pre-phrase coins. Looking only one empty batch ahead
+// can stop just before that change output.
+static const int RESTORE_MIN_SCAN = KEYPOOL_SIZE + RESTORE_BATCH;
 
 int CmdNewPhrase()
 {
@@ -105,6 +109,7 @@ bool RestoreFromPhrase(const std::string& strMnemonic,
     // up nothing. Every derived key is written to the wallet before the scan,
     // because the scan asks the wallet what belongs to it.
     int nTotalDerived = (int)nHDNext;
+    int nStopDepth = max(nMinDepth, RESTORE_MIN_SCAN);
     while (nTotalDerived < RESTORE_MAX)
     {
         // Count wallet transactions, not scan hits.
@@ -146,7 +151,7 @@ bool RestoreFromPhrase(const std::string& strMnemonic,
 
         // A whole batch with nothing in it means far enough -- unless the
         // caller asked to look deeper anyway.
-        if (nWalletAfter == nWalletBefore && nTotalDerived >= nMinDepth)
+        if (nWalletAfter == nWalletBefore && nTotalDerived >= nStopDepth)
             break;
     }