ZaStoGram/TMessagesProj/jni/tgnet/wss/WssSocket.h

/*
 * This is the source code of tgnet library v. 1.1
 * It is licensed under GNU GPL v. 2 or later.
 */

#ifndef TGNET_WSS_SOCKET_H
#define TGNET_WSS_SOCKET_H

#include "../transport/TransportSocket.h"

#include <openssl/ssl.h>
#include <deque>
#include <memory>
#include <string>
#include <vector>

namespace tgnet {
namespace wss {

struct Route {
    std::string relayHost;
    std::string relayHostFallback;
    std::string connectHost;
    uint16_t relayPort = 443;
    std::string domain;
    std::string path = "/apiws";
    bool viaFallback = false;
    // The Worker tunnel reaches the DC over plain TCP, where bytes 60..61 of
    // the obfuscation header must name the DC and traffic class.
    bool tunnel = false;
    // Position in the Cloudflare front catalog (kwsN.<front domain>), or -1.
    // The front forwards the WebSocket to Telegram Web itself, so unlike the
    // tunnel it needs no destination address and carries the byte stream as is.
    int32_t cdnSlot = -1;
    int32_t cdnDcId = 0;
    // Which of the front's two Cloudflare addresses this route dials.
    int32_t cdnAddress = -1;
    // Key for failure counting and suppression when it must not follow the
    // domain: every front domain of a DC shares one health record.
    std::string healthDomain;
    // Network the route was chosen on (see SetNetworkType): its failures count
    // against that network only, even when reported after a switch.
    int32_t network = 0;
    // DC of an official relay route, which picks relayHost from the DC's list
    // of ingress addresses; 0 for fronts and the tunnel.
    int32_t relayDcId = 0;
    // Plain TCP to the DC's own address: the guarded step after the tunnel and
    // before the fronts that ignore suppression. Not a WSS socket: the caller
    // dials dcAddress itself and reports the verdict (NoteDirectStepClosed).
    bool direct = false;
};

// Wi-Fi and mobile data reach Telegram through different providers, so relay
// health, suppression and the IP/DNS preference are kept per network type.
void SetNetworkType(int32_t networkType);

// Telegram's public web relays cover production DC1-DC5. Media connections
// use the corresponding -1 relay, matching Telegram Web's transport catalog.
// While a DC's relay is suppressed as unreachable, the route switches to the
// Cloudflare front domains of tg-ws-proxy, then to the ZaStoGram Cloudflare
// Worker, which opens dcAddress (IPv4) over TCP itself, then to direct TCP
// (Route::direct) unless that is paused, and finally to the fronts again
// whatever their suppression.
bool OfficialRoute(int32_t dcId, bool mediaConnection, bool testBackend, const std::string &dcAddress, Route *route);

// Received bytes that prove a direct step: as with the fronts, networks that
// throttle Telegram let the first answers through and freeze the rest.
constexpr uint64_t kDirectProofBytes = 32 * 1024;

// Verdict on a direct step (Route::direct) once its socket closes. A failure
// (no TCP, no answer, or a session that ended short of kDirectProofBytes)
// pauses direct TCP for this DC and network at once; proof clears the pause
// history. Anything else (an idle close by the app) leaves it as it was.
void NoteDirectStepClosed(const Route &route, uint64_t receivedBytes, const char *failure);

// Whether OfficialRoute would still hand out this exact route: same network,
// not suppressed and not switched to the relay's DNS name.
bool RouteUsable(const Route &route);

// The Cloudflare front route that replaced a stale one after a failure moved
// the catalog position on; false if the route is not a front route.
bool FollowCdnRoute(const Route &stale, Route *fresh);

// Whether OfficialRoute would carry this DC through the Cloudflare tunnel now.
bool DatacenterTunneled(int32_t dcId, bool mediaConnection, bool testBackend);

// Keeps relay suppression across launches in this file (read once).
void SetRouteHealthPath(const std::string &path);

class Socket final : public transport::Socket {
public:
    explicit Socket(Route route);
    ~Socket() override;

    bool open(const struct sockaddr *address, socklen_t addressLength, std::string *diagnostic) override;
    int fd() const override;
    bool onEvent(uint32_t events, std::vector<std::vector<uint8_t>> &payloads, std::string *diagnostic) override;
    bool write(const uint8_t *data, uint32_t size, std::string *diagnostic) override;
    size_t queuedOutputBytes() const override;
    bool isReady() const override;
    bool wantsWrite() const override;
    bool canWriteApplicationData() const override;
    bool isClosed() const override;
    transport::HandshakePhase handshakePhase() const override;
    const char *transportName() const override;
    void timedOut() override;
    void timedOutMidPacket(bool midPacket) override;
    void noteAppDataTimeout() override;
    std::string takeSessionSummary() override;
    uint64_t receivedBytes() const override;
    void close() override;

    const Route &route() const;

    // A pool spare nobody is waiting for: its failures say too little about
    // the relay to move real connections to another address or the tunnel.
    void setSpeculative(bool value);

private:
    enum class State : uint8_t {
        TcpConnecting,
        TlsHandshake,
        HttpWrite,
        HttpRead,
        Ready,
        Closed,
    };

    enum class IoWait : uint8_t {
        None,
        Read,
        Write,
    };

    bool finishTcpConnect(std::string *diagnostic);
    bool startTls(std::string *diagnostic);
    bool pumpTls(std::string *diagnostic);
    bool queueHttpUpgrade(std::string *diagnostic);
    bool flushPending(std::string *diagnostic);
    bool readIntoBuffer(std::string *diagnostic);
    bool parseHttpResponse(std::string *diagnostic);
    bool parseFrames(std::vector<std::vector<uint8_t>> &payloads, std::string *diagnostic);
    bool queueFrame(uint8_t opcode, const uint8_t *data, uint32_t size, std::string *diagnostic);
    void setIoWait(IoWait wait, const char *operation);
    bool writesWaitForRead() const;
    void noteAttemptFailed();
    void noteUpgradeSucceeded();
    bool outputDrained() const;
    const char *stateName() const;
    const char *ioWaitName() const;
    Route routeConfig;
    // Numeric address this socket dialled, to tell a dropped flow from a dead relay.
    std::string peerAddress;
    SSL *ssl = nullptr;
    int socketFd = -1;
    State state = State::Closed;
    IoWait ioWait = IoWait::None;
    transport::HandshakePhase phase = transport::HandshakePhase::None;
    // Keep every TLS write retry in a stable, immutable allocation. OpenSSL
    // requires the same buffer and length after WANT_READ/WANT_WRITE; a single
    // appendable vector violated that contract under bursty uploads.
    std::deque<std::vector<uint8_t>> pendingOutput;
    size_t pendingOutputOffset = 0;
    size_t pendingOutputBytes = 0;
    std::vector<uint8_t> inputBuffer;
    std::string secWebSocketKey;
    // The relay parses the 64-byte obfuscation header out of the payload of the
    // FIRST binary frame alone — not out of the reassembled message and not out
    // of the TCP stream. A shorter first frame is fatal and, worse, silent: the
    // relay simply never answers. Hold the opening bytes back until the header
    // is complete so no caller chunking can violate that.
    std::vector<uint8_t> openingFrame;
    bool openingFrameSent = false;
    bool fragmentedMessage = false;
    bool failureRecorded = false;
    // Upgrades answered 503 on this TLS connection (Cloudflare front only).
    uint32_t upgradeRetries = 0;
    bool provenRecorded = false;
    bool speculative = false;
    bool fromPool = false;
    // For the wss_session summary logged on close.
    int64_t openedAtMs = 0;
    int64_t readyAtMs = 0;
    int64_t firstDataAtMs = 0;
    // With first_data_ms it gives throughput: rx / (last - first).
    int64_t lastDataAtMs = 0;
    // Set for the duration of timedOutMidPacket(true).
    bool timeoutMidPacket = false;
    uint64_t bytesOut = 0;
    uint64_t bytesIn = 0;
    bool summaryTaken = false;
    bool reachableRecorded = false;
    // SSL_write returned WANT_READ: the record can only continue after input.
    bool writeBlockedOnRead = false;
};

std::unique_ptr<transport::Socket> CreateSocket(Route route);

} // namespace wss
} // namespace tgnet

#endif