#include <cactus/test.h>

#include "socksd.h"

using namespace cactus;

const folly::SocketAddress kTestProxyAddress("127.0.0.1", 58432);
const folly::SocketAddress kTestServerAddress("127.0.0.1", 58433);

// Using fake client. Implementing real client is a separate task, and we don't want to show our solution here.
const std::string kFakeRequest = std::string("\x04\x01\xe4\x41\x7f\x00\x00\x01", 8) + std::string("prime\0", 6);

FIBER_TEST_CASE("Ping pong proxy") {
    SocksServer socks(kTestProxyAddress);
    auto pingPong = ListenTCP(kTestServerAddress);

    ServerGroup g;
    g.Spawn([&] {
        socks.Serve();
    });

    int clients = 0;
    g.Spawn([&] {
        std::shared_ptr<IConn> client = pingPong->Accept();
        ++clients;
        g.Spawn([client] {
            BufferedReader reader(client.get());
            reader.WriteTo(client.get());
        });
    });

    auto conn = DialTCP(kTestProxyAddress);
    conn->Write(View(kFakeRequest));

    std::array<char, 8> reply{};
    conn->ReadFull(View(reply));

    SECTION("Echo") {
        for (int i = 0; i < 1024; ++i) {
            std::array<int, 1> ping{i};
            conn->Write(View(ping));
            conn->ReadFull(View(ping));
            REQUIRE(ping[0] == i);
        }

        REQUIRE(1 == clients);
    }

    SECTION("Half open shutdown") {
        std::string hello = "hello";
        conn->Write(View(hello));
        conn->CloseWrite();

        REQUIRE("hello" == conn->ReadAllToString());
        REQUIRE(1 == clients);
    }
}

FIBER_TEST_CASE("Invalid inputs") {
    SocksServer socks(kTestProxyAddress);
    ServerGroup g;

    g.Spawn([&] {
        socks.Serve();
    });

    SECTION("Username overflow") {
        std::string fakeHeader("\x04\x01\x16\x00\x08\x08\x08\x08", 8);

        auto conn = DialTCP(kTestProxyAddress);
        conn->Write(View(fakeHeader));

        auto overflowBuffer = [&] {
            BufferedWriter buf(conn.get());
            for (;;) {
                std::string smiley = ":)";
                buf.Write(View(smiley));
            }
        };

        REQUIRE_THROWS_AS(overflowBuffer(), std::system_error);
    }
}