Deduplicate addrdb.cpp and use CHashWriter/Verifier

2 files changed, 72 insertions, 144 deletions

diff --git a/src/addrdb.cpp b/src/addrdb.cpp
index a3743cd0d4..7f85c16585 100644
--- a/src/addrdb.cpp
+++ b/src/addrdb.cpp
@@ -15,25 +15,31 @@
 #include "tinyformat.h"
 #include "util.h"
 
+namespace {
 
-CBanDB::CBanDB()
+template <typename Stream, typename Data>
+bool SerializeDB(Stream& stream, const Data& data)
 {
-    pathBanlist = GetDataDir() / "banlist.dat";
+    // Write and commit header, data
+    try {
+        CHashWriter hasher(SER_DISK, CLIENT_VERSION);
+        stream << FLATDATA(Params().MessageStart()) << data;
+        hasher << FLATDATA(Params().MessageStart()) << data;
+        stream << hasher.GetHash();
+    } catch (const std::exception& e) {
+        return error("%s: Serialize or I/O error - %s", __func__, e.what());
+    }
+
+    return true;
 }
 
-bool CBanDB::Write(const banmap_t& banSet)
+template <typename Data>
+bool SerializeFileDB(const std::string& prefix, const fs::path& path, const Data& data)
 {
     // Generate random temporary filename
     unsigned short randv = 0;
     GetRandBytes((unsigned char*)&randv, sizeof(randv));
-    std::string tmpfn = strprintf("banlist.dat.%04x", randv);
-
-    // serialize banlist, checksum data up to that point, then append csum
-    CDataStream ssBanlist(SER_DISK, CLIENT_VERSION);
-    ssBanlist << FLATDATA(Params().MessageStart());
-    ssBanlist << banSet;
-    uint256 hash = Hash(ssBanlist.begin(), ssBanlist.end());
-    ssBanlist << hash;
+    std::string tmpfn = strprintf("%s.%04x", prefix, randv);
 
     // open temp output file, and associate with CAutoFile
     fs::path pathTmp = GetDataDir() / tmpfn;
@@ -42,69 +48,41 @@ bool CBanDB::Write(const banmap_t& banSet)
     if (fileout.IsNull())
         return error("%s: Failed to open file %s", __func__, pathTmp.string());
 
-    // Write and commit header, data
-    try {
-        fileout << ssBanlist;
-    }
-    catch (const std::exception& e) {
-        return error("%s: Serialize or I/O error - %s", __func__, e.what());
-    }
+    // Serialize
+    if (!SerializeDB(fileout, data)) return false;
     FileCommit(fileout.Get());
     fileout.fclose();
 
-    // replace existing banlist.dat, if any, with new banlist.dat.XXXX
-    if (!RenameOver(pathTmp, pathBanlist))
+    // replace existing file, if any, with new file
+    if (!RenameOver(pathTmp, path))
         return error("%s: Rename-into-place failed", __func__);
 
     return true;
 }
 
-bool CBanDB::Read(banmap_t& banSet)
+template <typename Stream, typename Data>
+bool DeserializeDB(Stream& stream, Data& data, bool fCheckSum = true)
 {
-    // open input file, and associate with CAutoFile
-    FILE *file = fsbridge::fopen(pathBanlist, "rb");
-    CAutoFile filein(file, SER_DISK, CLIENT_VERSION);
-    if (filein.IsNull())
-        return error("%s: Failed to open file %s", __func__, pathBanlist.string());
-
-    // use file size to size memory buffer
-    uint64_t fileSize = fs::file_size(pathBanlist);
-    uint64_t dataSize = 0;
-    // Don't try to resize to a negative number if file is small
-    if (fileSize >= sizeof(uint256))
-        dataSize = fileSize - sizeof(uint256);
-    std::vector<unsigned char> vchData;
-    vchData.resize(dataSize);
-    uint256 hashIn;
-
-    // read data and checksum from file
-    try {
-        filein.read((char *)&vchData[0], dataSize);
-        filein >> hashIn;
-    }
-    catch (const std::exception& e) {
-        return error("%s: Deserialize or I/O error - %s", __func__, e.what());
-    }
-    filein.fclose();
-
-    CDataStream ssBanlist(vchData, SER_DISK, CLIENT_VERSION);
-
-    // verify stored checksum matches input data
-    uint256 hashTmp = Hash(ssBanlist.begin(), ssBanlist.end());
-    if (hashIn != hashTmp)
-        return error("%s: Checksum mismatch, data corrupted", __func__);
-
-    unsigned char pchMsgTmp[4];
     try {
+        CHashVerifier<Stream> verifier(&stream);
         // de-serialize file header (network specific magic number) and ..
-        ssBanlist >> FLATDATA(pchMsgTmp);
-
+        unsigned char pchMsgTmp[4];
+        verifier >> FLATDATA(pchMsgTmp);
         // ... verify the network matches ours
         if (memcmp(pchMsgTmp, Params().MessageStart(), sizeof(pchMsgTmp)))
             return error("%s: Invalid network magic number", __func__);
 
-        // de-serialize ban data
-        ssBanlist >> banSet;
+        // de-serialize data
+        verifier >> data;
+
+        // verify checksum
+        if (fCheckSum) {
+            uint256 hashTmp;
+            stream >> hashTmp;
+            if (hashTmp != verifier.GetHash()) {
+                return error("%s: Checksum mismatch, data corrupted", __func__);
+            }
+        }
     }
     catch (const std::exception& e) {
         return error("%s: Deserialize or I/O error - %s", __func__, e.what());
@@ -113,106 +91,56 @@ bool CBanDB::Read(banmap_t& banSet)
     return true;
 }
 
-CAddrDB::CAddrDB()
+template <typename Data>
+bool DeserializeFileDB(const fs::path& path, Data& data)
 {
-    pathAddr = GetDataDir() / "peers.dat";
+    // open input file, and associate with CAutoFile
+    FILE *file = fsbridge::fopen(path, "rb");
+    CAutoFile filein(file, SER_DISK, CLIENT_VERSION);
+    if (filein.IsNull())
+        return error("%s: Failed to open file %s", __func__, path.string());
+
+    return DeserializeDB(filein, data);
 }
 
-bool CAddrDB::Write(const CAddrMan& addr)
-{
-    // Generate random temporary filename
-    unsigned short randv = 0;
-    GetRandBytes((unsigned char*)&randv, sizeof(randv));
-    std::string tmpfn = strprintf("peers.dat.%04x", randv);
+}
 
-    // serialize addresses, checksum data up to that point, then append csum
-    CDataStream ssPeers(SER_DISK, CLIENT_VERSION);
-    ssPeers << FLATDATA(Params().MessageStart());
-    ssPeers << addr;
-    uint256 hash = Hash(ssPeers.begin(), ssPeers.end());
-    ssPeers << hash;
+CBanDB::CBanDB()
+{
+    pathBanlist = GetDataDir() / "banlist.dat";
+}
 
-    // open temp output file, and associate with CAutoFile
-    fs::path pathTmp = GetDataDir() / tmpfn;
-    FILE *file = fsbridge::fopen(pathTmp, "wb");
-    CAutoFile fileout(file, SER_DISK, CLIENT_VERSION);
-    if (fileout.IsNull())
-        return error("%s: Failed to open file %s", __func__, pathTmp.string());
+bool CBanDB::Write(const banmap_t& banSet)
+{
+    return SerializeFileDB("banlist", pathBanlist, banSet);
+}
 
-    // Write and commit header, data
-    try {
-        fileout << ssPeers;
-    }
-    catch (const std::exception& e) {
-        return error("%s: Serialize or I/O error - %s", __func__, e.what());
-    }
-    FileCommit(fileout.Get());
-    fileout.fclose();
+bool CBanDB::Read(banmap_t& banSet)
+{
+    return DeserializeFileDB(pathBanlist, banSet);
+}
 
-    // replace existing peers.dat, if any, with new peers.dat.XXXX
-    if (!RenameOver(pathTmp, pathAddr))
-        return error("%s: Rename-into-place failed", __func__);
+CAddrDB::CAddrDB()
+{
+    pathAddr = GetDataDir() / "peers.dat";
+}
 
-    return true;
+bool CAddrDB::Write(const CAddrMan& addr)
+{
+    return SerializeFileDB("peers", pathAddr, addr);
 }
 
 bool CAddrDB::Read(CAddrMan& addr)
 {
-    // open input file, and associate with CAutoFile
-    FILE *file = fsbridge::fopen(pathAddr, "rb");
-    CAutoFile filein(file, SER_DISK, CLIENT_VERSION);
-    if (filein.IsNull())
-        return error("%s: Failed to open file %s", __func__, pathAddr.string());
-
-    // use file size to size memory buffer
-    uint64_t fileSize = fs::file_size(pathAddr);
-    uint64_t dataSize = 0;
-    // Don't try to resize to a negative number if file is small
-    if (fileSize >= sizeof(uint256))
-        dataSize = fileSize - sizeof(uint256);
-    std::vector<unsigned char> vchData;
-    vchData.resize(dataSize);
-    uint256 hashIn;
-
-    // read data and checksum from file
-    try {
-        filein.read((char *)&vchData[0], dataSize);
-        filein >> hashIn;
-    }
-    catch (const std::exception& e) {
-        return error("%s: Deserialize or I/O error - %s", __func__, e.what());
-    }
-    filein.fclose();
-
-    CDataStream ssPeers(vchData, SER_DISK, CLIENT_VERSION);
-
-    // verify stored checksum matches input data
-    uint256 hashTmp = Hash(ssPeers.begin(), ssPeers.end());
-    if (hashIn != hashTmp)
-        return error("%s: Checksum mismatch, data corrupted", __func__);
-
-    return Read(addr, ssPeers);
+    return DeserializeFileDB(pathAddr, addr);
 }
 
 bool CAddrDB::Read(CAddrMan& addr, CDataStream& ssPeers)
 {
-    unsigned char pchMsgTmp[4];
-    try {
-        // de-serialize file header (network specific magic number) and ..
-        ssPeers >> FLATDATA(pchMsgTmp);
-
-        // ... verify the network matches ours
-        if (memcmp(pchMsgTmp, Params().MessageStart(), sizeof(pchMsgTmp)))
-            return error("%s: Invalid network magic number", __func__);
-
-        // de-serialize address data into one CAddrMan object
-        ssPeers >> addr;
-    }
-    catch (const std::exception& e) {
-        // de-serialization has failed, ensure addrman is left in a clean state
+    bool ret = DeserializeDB(ssPeers, addr, false);
+    if (!ret) {
+        // Ensure addrman is left in a clean state
         addr.Clear();
-        return error("%s: Deserialize or I/O error - %s", __func__, e.what());
     }
-
-    return true;
+    return ret;
 }
diff --git a/src/addrdb.h b/src/addrdb.h
index c3d509bd3a..6cb36dfac4 100644
--- a/src/addrdb.h
+++ b/src/addrdb.h
@@ -85,7 +85,7 @@ public:
     CAddrDB();
     bool Write(const CAddrMan& addr);
     bool Read(CAddrMan& addr);
-    bool Read(CAddrMan& addr, CDataStream& ssPeers);
+    static bool Read(CAddrMan& addr, CDataStream& ssPeers);
 };
 
 /** Access to the banlist database (banlist.dat) */