/*
 * This file is part of the Flowee project
 * Copyright (C) 2012-2016 The Bitcoin Core developers
 * Copyright (C) 2022 Tom Zander <tom@flowee.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
#include "address_manager_tests.h"

#include <clientversion.h>
#include <serialize.h>
#include <streaming/streams.h>

#include "net.h"
#include "chainparams.h"


static CDataStream addrmanToStream(const AddrManSerializationMock &addrman)
{
    CDataStream ssPeersIn(SER_DISK, CLIENT_VERSION);
    ssPeersIn << FLATDATA(Params().magic());
    ssPeersIn << addrman;
    std::string str = ssPeersIn.str();
    std::vector<unsigned char> vchData(str.begin(), str.end());
    return CDataStream(vchData, SER_DISK, CLIENT_VERSION);
}

void TestAddrMan::read()
{
    AddrManUncorrupted addrmanUncorrupted;
    addrmanUncorrupted.makeDeterministic();

    CService addr1 = CService("250.7.1.1", 8333);
    CService addr2 = CService("250.7.2.2", 9999);
    CService addr3 = CService("250.7.3.3", 9999);

    // Add three addresses to new table.
    addrmanUncorrupted.Add(CAddress(addr1), CService("252.5.1.1", 8333));
    addrmanUncorrupted.Add(CAddress(addr2), CService("252.5.1.1", 8333));
    addrmanUncorrupted.Add(CAddress(addr3), CService("252.5.1.1", 8333));

    // Test that the de-serialization does not throw an exception.
    CDataStream ssPeers1 = addrmanToStream(addrmanUncorrupted);
    bool exceptionThrown = false;
    CAddrMan addrman1;

    QCOMPARE(addrman1.size(),  0);
    try {
        unsigned char pchMsgTmp[4];
        ssPeers1 >> FLATDATA(pchMsgTmp);
        ssPeers1 >> addrman1;
    } catch (const std::exception& e) {
        exceptionThrown = true;
    }

    QCOMPARE(addrman1.size(), 3);
    QCOMPARE(exceptionThrown, false);

    // Test that CAddrDB::Read creates an addrman with the correct number of addrs.
    CDataStream ssPeers2 = addrmanToStream(addrmanUncorrupted);

    CAddrMan addrman2;
    CAddrDB adb;
    QCOMPARE(addrman2.size(), 0);
    adb.Read(addrman2, ssPeers2);
    QCOMPARE(addrman2.size(), 3);
}

void TestAddrMan::readCorrupted()
{
    AddrManCorrupted addrmanCorrupted;
    addrmanCorrupted.makeDeterministic();

    // Test that the de-serialization of corrupted addrman throws an exception.
    CDataStream ssPeers1 = addrmanToStream(addrmanCorrupted);
    bool exceptionThrown = false;
    CAddrMan addrman1;
    QCOMPARE(addrman1.size(), 0);
    try {
        unsigned char pchMsgTmp[4];
        ssPeers1 >> FLATDATA(pchMsgTmp);
        ssPeers1 >> addrman1;
    } catch (const std::exception& e) {
        exceptionThrown = true;
    }
    // Even through de-serialization failed addrman is not left in a clean state.
    QCOMPARE(addrman1.size(), 1);
    QVERIFY(exceptionThrown);

    // Test that CAddrDB::Read leaves addrman in a clean state if de-serialization fails.
    CDataStream ssPeers2 = addrmanToStream(addrmanCorrupted);

    CAddrMan addrman2;
    CAddrDB adb;
    QCOMPARE(addrman2.size(), 0);
    adb.Read(addrman2, ssPeers2);
    QCOMPARE(addrman2.size(), 0);
}

void TestAddrMan::simple()
{
    AddrManMock addrman;

    // Set addrman addr placement to be deterministic.
    addrman.makeDeterministic();

    CNetAddr source = CNetAddr("252.2.2.2:8333");

    // Test 1: Does Addrman respond correctly when empty.
    QCOMPARE(addrman.size(), 0);
    CAddrInfo addr_null = addrman.Select();
    QCOMPARE(addr_null.ToString(), std::string("[::]:0"));

    // Test 2: Does Addrman::Add work as expected.
    CService addr1 = CService("250.1.1.1:8333");
    addrman.Add(CAddress(addr1), source);
    QCOMPARE(addrman.size(), 1);
    CAddrInfo addr_ret1 = addrman.Select();
    QCOMPARE(addr_ret1.ToString(), std::string("250.1.1.1:8333"));

    // Test 3: Does IP address deduplication work correctly. 
    //  Expected dup IP should not be added.
    CService addr1_dup = CService("250.1.1.1:8333");
    addrman.Add(CAddress(addr1_dup), source);
    QCOMPARE(addrman.size(), 1);


    // Test 5: New table has one addr and we add a diff addr we should
    //  have two addrs.
    CService addr2 = CService("250.1.1.2:8333");
    addrman.Add(CAddress(addr2), source);
    QCOMPARE(addrman.size(), 2);

    // Test 6: AddrMan::Clear() should empty the new table. 
    addrman.Clear();
    QCOMPARE(addrman.size(), 0);
    CAddrInfo addr_null2 = addrman.Select();
    QCOMPARE(addr_null2.ToString(), "[::]:0");
}

void TestAddrMan::ports()
{
    AddrManMock addrman;

    // Set addrman addr placement to be deterministic.
    addrman.makeDeterministic();

    CNetAddr source = CNetAddr("252.2.2.2:8333");

    QCOMPARE(addrman.size(), 0);

    // Test 7; Addr with same IP but diff port does not replace existing addr.
    CService addr1 = CService("250.1.1.1:8333");
    addrman.Add(CAddress(addr1), source);
    QCOMPARE(addrman.size(), 1);

    CService addr1_port = CService("250.1.1.1:8334");
    addrman.Add(CAddress(addr1_port), source);
    QCOMPARE(addrman.size(), 1);
    CAddrInfo addr_ret2 = addrman.Select();
    QCOMPARE(addr_ret2.ToString(), "250.1.1.1:8333");

    // Test 8: Add same IP but diff port to tried table, it doesn't get added.
    //  Perhaps this is not ideal behavior but it is the current behavior.
    addrman.Good(CAddress(addr1_port));
    QCOMPARE(addrman.size(), 1);
    bool newOnly = true;
    CAddrInfo addr_ret3 = addrman.Select(newOnly);
    QCOMPARE(addr_ret3.ToString(), "250.1.1.1:8333");
}

void TestAddrMan::select()
{
    AddrManMock addrman;

    // Set addrman addr placement to be deterministic.
    addrman.makeDeterministic();

    CNetAddr source = CNetAddr("252.2.2.2:8333");

    // Test 9: Select from new with 1 addr in new.
    CService addr1 = CService("250.1.1.1:8333");
    addrman.Add(CAddress(addr1), source);
    QCOMPARE(addrman.size(), 1);

    bool newOnly = true;
    CAddrInfo addr_ret1 = addrman.Select(newOnly);
    QCOMPARE(addr_ret1.ToString(), "250.1.1.1:8333");


    // Test 10: move addr to tried, select from new expected nothing returned.
    addrman.Good(CAddress(addr1));
    QCOMPARE(addrman.size(), 1);
    CAddrInfo addr_ret2 = addrman.Select(newOnly);
    QCOMPARE(addr_ret2.ToString(), "[::]:0");

    CAddrInfo addr_ret3 = addrman.Select();
    QCOMPARE(addr_ret3.ToString(), "250.1.1.1:8333");
}

void TestAddrMan::new_collisions()
{
    AddrManMock addrman;

    // Set addrman addr placement to be deterministic.
    addrman.makeDeterministic();

    CNetAddr source = CNetAddr("252.2.2.2:8333");

    QCOMPARE(addrman.size(), 0);

    for (unsigned int i = 1; i < 4; i++){
        CService addr = CService("250.1.1."+std::to_string(i));
        addrman.Add(CAddress(addr), source);

        //Test 11: No collision in new table yet.
        QCOMPARE(addrman.size(), i);
    }

    //Test 12: new table collision!
    CService addr1 = CService("250.1.1.4");
    addrman.Add(CAddress(addr1), source);
    QCOMPARE(addrman.size(), 3);

    CService addr2 = CService("250.1.1.5");
    addrman.Add(CAddress(addr2), source);
    QCOMPARE(addrman.size(), 4);
}

void TestAddrMan::tried_collisions()
{
    AddrManMock addrman;

    // Set addrman addr placement to be deterministic.
    addrman.makeDeterministic();

    CNetAddr source = CNetAddr("252.2.2.2:8333");

    QCOMPARE(addrman.size(), 0);

    for (unsigned int i = 1; i < 75; i++){
        CService addr = CService("250.1.1."+std::to_string(i));
        addrman.Add(CAddress(addr), source);
        addrman.Good(CAddress(addr));

        //Test 13: No collision in tried table yet.
        // BOOST_TEST_MESSAGE(addrman.size());
        QCOMPARE(addrman.size(), i);
    }

    //Test 14: tried table collision!
    CService addr1 = CService("250.1.1.76");
    addrman.Add(CAddress(addr1), source);
    QCOMPARE(addrman.size(), 74);

    CService addr2 = CService("250.1.1.77");
    addrman.Add(CAddress(addr2), source);
    QCOMPARE(addrman.size(), 75);
}

void TestAddrMan::serialization()
{
    std::vector<char> data;
    {// save
        AddrManMock addrman;

        // Set addrman addr placement to be deterministic.
        addrman.makeDeterministic();

        CNetAddr source = CNetAddr("252.2.2.2:8333");

        for (unsigned int i = 1; i < 75; ++i) {
            CService addr = CService("250.1.1."+std::to_string(i));
            bool ok = addrman.Add(CAddress(addr), source);
            QVERIFY(ok);
            addrman.Good(CAddress(addr));
            if (i == 23 || i == 67) {
                struct in_addr pipv4Addr;
                addr.GetInAddr(&pipv4Addr);
                CAddrInfo *info = addrman.Find(CNetAddr(pipv4Addr));
                QVERIFY(info);
                info->setKnowsXThin(true);
            }
            //Test 13: No collision in tried table yet.
            // BOOST_TEST_MESSAGE(addrman.size());
            QCOMPARE(addrman.size(), i);
        }
        QCOMPARE(addrman.size(), 74);
        CDataStream ssPeers(SER_DISK, CLIENT_VERSION);
        // ssPeers << FLATDATA(Params().MessageStart());
        ssPeers << addrman;
        data = std::vector<char>(ssPeers.begin(), ssPeers.end());
    }

    { // load
        AddrManMock addrman;
        CDataStream ssPeers(data, SER_DISK, CLIENT_VERSION);
        ssPeers >> addrman;

        QCOMPARE(addrman.size(), 74);
        CNetAddr addr1("250.1.1.23");
        CNetAddr addr2("250.1.1.67");
        CNetAddr addr3("250.1.1.2");
        QVERIFY(addrman.Find(addr1));
        QCOMPARE(addrman.Find(addr1)->getKnowsXThin(), true);
        QVERIFY(addrman.Find(addr2));
        QCOMPARE(addrman.Find(addr2)->getKnowsXThin(), true);
        QVERIFY(addrman.Find(addr3));
        QCOMPARE(addrman.Find(addr3)->getKnowsXThin(), false);
    }
}


void AddrManMock::makeDeterministic()
{
    nKey.SetNull();
    seed_insecure_rand(true);
}

void AddrManUncorrupted::Serialize(CDataStream &s, int nType, int nVersionDummy) const
{
    CAddrMan::Serialize(s, nType, nVersionDummy);
}

void AddrManCorrupted::Serialize(CDataStream &s, int nType, int nVersionDummy) const
{
    Q_UNUSED(nType);
    Q_UNUSED(nVersionDummy);
    // Produces corrupt output that claims addrman has 20 addrs when it only has one addr.
    unsigned char nVersion = 1;
    s << nVersion;
    s << ((unsigned char)32);
    s << nKey;
    s << 10; // nNew
    s << 10; // nTried

    int nUBuckets = ADDRMAN_NEW_BUCKET_COUNT ^ (1 << 30);
    s << nUBuckets;

    CAddress addr = CAddress(CService("252.1.1.1", 7777));
    CAddrInfo info = CAddrInfo(addr, CNetAddr("252.2.2.2"));
    s << info;
}