chore: sync from fn-registry agent
This commit is contained in:
fn-registry agent
+404
@@ -0,0 +1,404 @@
|
||||
#include "ws_client.h"
|
||||
|
||||
#include <chrono>
|
||||
#include <cstdint>
|
||||
#include <cstdio>
|
||||
#include <cstring>
|
||||
#include <random>
|
||||
#include <sstream>
|
||||
#include <vector>
|
||||
|
||||
#ifdef _WIN32
|
||||
#include <winsock2.h>
|
||||
#include <ws2tcpip.h>
|
||||
#pragma comment(lib, "ws2_32.lib")
|
||||
typedef SOCKET sock_t;
|
||||
#define SOCK_INVALID INVALID_SOCKET
|
||||
#define SOCK_CLOSE closesocket
|
||||
#define SOCK_ERR WSAGetLastError()
|
||||
#define FN_SOCK_NONBLOCK(s) do { u_long m = 1; ioctlsocket((s), FIONBIO, &m); } while (0)
|
||||
#else
|
||||
#include <sys/socket.h>
|
||||
#include <netinet/in.h>
|
||||
#include <arpa/inet.h>
|
||||
#include <unistd.h>
|
||||
#include <fcntl.h>
|
||||
#include <errno.h>
|
||||
typedef int sock_t;
|
||||
#define SOCK_INVALID (-1)
|
||||
#define SOCK_CLOSE close
|
||||
#define SOCK_ERR errno
|
||||
#define FN_SOCK_NONBLOCK(s) do { int f = fcntl((s), F_GETFL, 0); fcntl((s), F_SETFL, f | O_NONBLOCK); } while (0)
|
||||
#endif
|
||||
|
||||
#ifdef _WIN32
|
||||
static bool wsa_init_ws() {
|
||||
static bool inited = false;
|
||||
static std::once_flag flag;
|
||||
std::call_once(flag, []() {
|
||||
WSADATA wsa;
|
||||
WSAStartup(MAKEWORD(2, 2), &wsa);
|
||||
});
|
||||
return true;
|
||||
}
|
||||
#endif
|
||||
|
||||
namespace {
|
||||
|
||||
// ----- Base64 (small, sufficient for 16-byte WS key) -----
|
||||
const char* kBase64 = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";
|
||||
|
||||
std::string base64_encode(const uint8_t* in, size_t len) {
|
||||
std::string out;
|
||||
out.reserve(((len + 2) / 3) * 4);
|
||||
size_t i = 0;
|
||||
for (; i + 3 <= len; i += 3) {
|
||||
uint32_t v = (uint32_t(in[i]) << 16) | (uint32_t(in[i + 1]) << 8) | uint32_t(in[i + 2]);
|
||||
out.push_back(kBase64[(v >> 18) & 63]);
|
||||
out.push_back(kBase64[(v >> 12) & 63]);
|
||||
out.push_back(kBase64[(v >> 6) & 63]);
|
||||
out.push_back(kBase64[v & 63]);
|
||||
}
|
||||
if (i < len) {
|
||||
uint32_t v = uint32_t(in[i]) << 16;
|
||||
if (i + 1 < len) v |= uint32_t(in[i + 1]) << 8;
|
||||
out.push_back(kBase64[(v >> 18) & 63]);
|
||||
out.push_back(kBase64[(v >> 12) & 63]);
|
||||
out.push_back(i + 1 < len ? kBase64[(v >> 6) & 63] : '=');
|
||||
out.push_back('=');
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
// Send exactly n bytes (blocking).
|
||||
bool send_all(sock_t sock, const char* data, size_t n) {
|
||||
size_t sent = 0;
|
||||
while (sent < n) {
|
||||
int k = send(sock, data + sent, static_cast<int>(n - sent), 0);
|
||||
if (k <= 0) return false;
|
||||
sent += k;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
// Receive exactly n bytes (blocking on a non-non-blocking socket).
|
||||
bool recv_all(sock_t sock, char* data, size_t n) {
|
||||
size_t got = 0;
|
||||
while (got < n) {
|
||||
int k = recv(sock, data + got, static_cast<int>(n - got), 0);
|
||||
if (k <= 0) return false;
|
||||
got += k;
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
// Receive up to n bytes; returns count, or -1 on error / -2 on would-block.
|
||||
int recv_some(sock_t sock, char* data, size_t n) {
|
||||
int k = recv(sock, data, static_cast<int>(n), 0);
|
||||
if (k > 0) return k;
|
||||
if (k == 0) return -1;
|
||||
#ifdef _WIN32
|
||||
if (WSAGetLastError() == WSAEWOULDBLOCK) return -2;
|
||||
#else
|
||||
if (errno == EAGAIN || errno == EWOULDBLOCK) return -2;
|
||||
#endif
|
||||
return -1;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
WsClient::WsClient() = default;
|
||||
|
||||
WsClient::~WsClient() {
|
||||
stop();
|
||||
}
|
||||
|
||||
void WsClient::start(const std::string& host, int port, const std::string& path) {
|
||||
State expected = State::Idle;
|
||||
if (!state_.compare_exchange_strong(expected, State::Connecting)) return;
|
||||
|
||||
host_ = host;
|
||||
port_ = port;
|
||||
path_ = path;
|
||||
stop_flag_.store(false);
|
||||
|
||||
worker_ = std::thread([this]() { this->run(); });
|
||||
}
|
||||
|
||||
void WsClient::stop() {
|
||||
stop_flag_.store(true);
|
||||
int s = sock_.exchange(-1);
|
||||
if (s != -1) SOCK_CLOSE(static_cast<sock_t>(s));
|
||||
out_cv_.notify_all();
|
||||
if (worker_.joinable()) worker_.join();
|
||||
state_.store(State::Stopped);
|
||||
}
|
||||
|
||||
int WsClient::drain(std::vector<std::string>& out, int max) {
|
||||
std::lock_guard<std::mutex> g(in_mu_);
|
||||
int n = 0;
|
||||
while (!in_queue_.empty() && n < max) {
|
||||
out.emplace_back(std::move(in_queue_.front()));
|
||||
in_queue_.pop_front();
|
||||
n++;
|
||||
}
|
||||
return n;
|
||||
}
|
||||
|
||||
bool WsClient::send_text(const std::string& payload) {
|
||||
if (state_.load() != State::Connected) return false;
|
||||
{
|
||||
std::lock_guard<std::mutex> g(out_mu_);
|
||||
out_queue_.push_back(payload);
|
||||
}
|
||||
out_cv_.notify_one();
|
||||
return true;
|
||||
}
|
||||
|
||||
void WsClient::run() {
|
||||
int backoff_ms = 500;
|
||||
while (!stop_flag_.load()) {
|
||||
state_.store(State::Connecting);
|
||||
if (connect_once()) {
|
||||
backoff_ms = 500; // reset on successful connect
|
||||
state_.store(State::Connected);
|
||||
read_loop();
|
||||
}
|
||||
state_.store(State::Backoff);
|
||||
if (stop_flag_.load()) break;
|
||||
|
||||
// Exponential backoff: 0.5s → 1s → 2s → 4s → 8s (cap).
|
||||
for (int slept = 0; slept < backoff_ms && !stop_flag_.load(); slept += 100) {
|
||||
std::this_thread::sleep_for(std::chrono::milliseconds(100));
|
||||
}
|
||||
backoff_ms = std::min(backoff_ms * 2, 8000);
|
||||
}
|
||||
state_.store(State::Stopped);
|
||||
}
|
||||
|
||||
bool WsClient::connect_once() {
|
||||
#ifdef _WIN32
|
||||
wsa_init_ws();
|
||||
#endif
|
||||
|
||||
sock_t sock = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
|
||||
if (sock == SOCK_INVALID) return false;
|
||||
|
||||
// 5s connect timeout via SO_*TIMEO. Stays blocking afterwards for the
|
||||
// handshake; read_loop switches to non-blocking with select().
|
||||
#ifdef _WIN32
|
||||
DWORD timeout_ms = 5000;
|
||||
setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, (const char*)&timeout_ms, sizeof(timeout_ms));
|
||||
setsockopt(sock, SOL_SOCKET, SO_SNDTIMEO, (const char*)&timeout_ms, sizeof(timeout_ms));
|
||||
#else
|
||||
struct timeval tv;
|
||||
tv.tv_sec = 5;
|
||||
tv.tv_usec = 0;
|
||||
setsockopt(sock, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
|
||||
setsockopt(sock, SOL_SOCKET, SO_SNDTIMEO, &tv, sizeof(tv));
|
||||
#endif
|
||||
|
||||
struct sockaddr_in addr;
|
||||
std::memset(&addr, 0, sizeof(addr));
|
||||
addr.sin_family = AF_INET;
|
||||
addr.sin_port = htons(static_cast<uint16_t>(port_));
|
||||
addr.sin_addr.s_addr = inet_addr(host_.c_str());
|
||||
|
||||
if (connect(sock, (struct sockaddr*)&addr, sizeof(addr)) != 0) {
|
||||
SOCK_CLOSE(sock);
|
||||
return false;
|
||||
}
|
||||
|
||||
sock_.store(static_cast<int>(sock));
|
||||
|
||||
if (!handshake()) {
|
||||
int s = sock_.exchange(-1);
|
||||
if (s != -1) SOCK_CLOSE(static_cast<sock_t>(s));
|
||||
return false;
|
||||
}
|
||||
|
||||
// Non-blocking for the read loop.
|
||||
FN_SOCK_NONBLOCK(sock);
|
||||
return true;
|
||||
}
|
||||
|
||||
bool WsClient::handshake() {
|
||||
sock_t sock = static_cast<sock_t>(sock_.load());
|
||||
|
||||
// 16 random bytes → base64 → Sec-WebSocket-Key.
|
||||
uint8_t key_raw[16];
|
||||
std::random_device rd;
|
||||
for (auto& b : key_raw) b = static_cast<uint8_t>(rd() & 0xff);
|
||||
std::string key_b64 = base64_encode(key_raw, sizeof(key_raw));
|
||||
|
||||
std::ostringstream req;
|
||||
req << "GET " << path_ << " HTTP/1.1\r\n";
|
||||
req << "Host: " << host_ << ":" << port_ << "\r\n";
|
||||
req << "Upgrade: websocket\r\n";
|
||||
req << "Connection: Upgrade\r\n";
|
||||
req << "Sec-WebSocket-Key: " << key_b64 << "\r\n";
|
||||
req << "Sec-WebSocket-Version: 13\r\n";
|
||||
req << "Origin: http://" << host_ << ":" << port_ << "\r\n";
|
||||
req << "\r\n";
|
||||
|
||||
std::string raw = req.str();
|
||||
if (!send_all(sock, raw.c_str(), raw.size())) return false;
|
||||
|
||||
// Read response headers (up to 4KB).
|
||||
std::string resp;
|
||||
char buf[1024];
|
||||
while (resp.find("\r\n\r\n") == std::string::npos) {
|
||||
int k = recv(sock, buf, sizeof(buf), 0);
|
||||
if (k <= 0) return false;
|
||||
resp.append(buf, k);
|
||||
if (resp.size() > 4096) return false;
|
||||
}
|
||||
|
||||
// Expect "HTTP/1.1 101".
|
||||
if (resp.compare(0, 12, "HTTP/1.1 101") != 0 &&
|
||||
resp.compare(0, 12, "HTTP/1.0 101") != 0) {
|
||||
fprintf(stderr, "[ws] handshake failed: %.*s\n",
|
||||
(int)std::min<size_t>(resp.size(), 120), resp.c_str());
|
||||
return false;
|
||||
}
|
||||
// We intentionally skip Sec-WebSocket-Accept verification — controlled
|
||||
// server, localhost-only, 101 status is enough for this app.
|
||||
return true;
|
||||
}
|
||||
|
||||
bool WsClient::read_loop() {
|
||||
sock_t sock = static_cast<sock_t>(sock_.load());
|
||||
|
||||
std::vector<uint8_t> rb; // accumulated read buffer
|
||||
rb.reserve(64 * 1024);
|
||||
|
||||
while (!stop_flag_.load()) {
|
||||
// Block on select() for up to 100ms so we can both read and check
|
||||
// the outgoing queue.
|
||||
fd_set rfds;
|
||||
FD_ZERO(&rfds);
|
||||
FD_SET(sock, &rfds);
|
||||
struct timeval tv;
|
||||
tv.tv_sec = 0;
|
||||
tv.tv_usec = 100 * 1000;
|
||||
int sel = select(static_cast<int>(sock) + 1, &rfds, nullptr, nullptr, &tv);
|
||||
if (sel < 0) return false;
|
||||
|
||||
if (sel > 0 && FD_ISSET(sock, &rfds)) {
|
||||
uint8_t tmp[8192];
|
||||
int k = recv_some(sock, reinterpret_cast<char*>(tmp), sizeof(tmp));
|
||||
if (k == -1) return false;
|
||||
if (k > 0) rb.insert(rb.end(), tmp, tmp + k);
|
||||
}
|
||||
|
||||
// Drain outgoing queue (text frames, masked).
|
||||
for (;;) {
|
||||
std::string payload;
|
||||
{
|
||||
std::lock_guard<std::mutex> g(out_mu_);
|
||||
if (out_queue_.empty()) break;
|
||||
payload = std::move(out_queue_.front());
|
||||
out_queue_.pop_front();
|
||||
}
|
||||
if (!send_frame(0x1, payload)) return false;
|
||||
}
|
||||
|
||||
// Parse frames. RFC6455 minimal: assume server never masks, no
|
||||
// continuation, opcodes: 0x1 text, 0x8 close, 0x9 ping, 0xA pong.
|
||||
while (rb.size() >= 2) {
|
||||
uint8_t b0 = rb[0];
|
||||
uint8_t b1 = rb[1];
|
||||
bool fin = (b0 & 0x80) != 0;
|
||||
(void)fin;
|
||||
int opcode = b0 & 0x0F;
|
||||
bool mask = (b1 & 0x80) != 0;
|
||||
uint64_t len = b1 & 0x7F;
|
||||
size_t pos = 2;
|
||||
if (len == 126) {
|
||||
if (rb.size() < pos + 2) break;
|
||||
len = (uint64_t(rb[pos]) << 8) | uint64_t(rb[pos + 1]);
|
||||
pos += 2;
|
||||
} else if (len == 127) {
|
||||
if (rb.size() < pos + 8) break;
|
||||
len = 0;
|
||||
for (int i = 0; i < 8; i++) len = (len << 8) | rb[pos + i];
|
||||
pos += 8;
|
||||
}
|
||||
uint8_t mkey[4] = {0, 0, 0, 0};
|
||||
if (mask) {
|
||||
if (rb.size() < pos + 4) break;
|
||||
for (int i = 0; i < 4; i++) mkey[i] = rb[pos + i];
|
||||
pos += 4;
|
||||
}
|
||||
if (rb.size() < pos + len) break;
|
||||
|
||||
std::string payload;
|
||||
payload.resize(static_cast<size_t>(len));
|
||||
for (size_t i = 0; i < len; i++) {
|
||||
uint8_t c = rb[pos + i];
|
||||
if (mask) c ^= mkey[i & 3];
|
||||
payload[i] = static_cast<char>(c);
|
||||
}
|
||||
rb.erase(rb.begin(), rb.begin() + pos + len);
|
||||
|
||||
switch (opcode) {
|
||||
case 0x1: { // text
|
||||
last_event_ts_.store(std::chrono::duration_cast<std::chrono::seconds>(
|
||||
std::chrono::system_clock::now().time_since_epoch()).count());
|
||||
std::lock_guard<std::mutex> g(in_mu_);
|
||||
in_queue_.emplace_back(std::move(payload));
|
||||
// Bound the queue. Drop oldest if too big — UI consumes
|
||||
// each frame so this should only kick in if the dashboard
|
||||
// is paused (window minimized, etc.).
|
||||
while (in_queue_.size() > 512) in_queue_.pop_front();
|
||||
break;
|
||||
}
|
||||
case 0x8: // close
|
||||
return false;
|
||||
case 0x9: // ping → reply with pong (same payload)
|
||||
if (!send_frame(0xA, payload)) return false;
|
||||
break;
|
||||
case 0xA: // pong, ignore
|
||||
break;
|
||||
default:
|
||||
// 0x0 continuation or unexpected opcode: bail (server
|
||||
// controlled by us, shouldn't happen).
|
||||
return false;
|
||||
}
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
bool WsClient::send_frame(int opcode, const std::string& payload) {
|
||||
sock_t sock = static_cast<sock_t>(sock_.load());
|
||||
if (sock == SOCK_INVALID || static_cast<int>(sock) < 0) return false;
|
||||
|
||||
std::vector<uint8_t> frame;
|
||||
frame.reserve(payload.size() + 16);
|
||||
frame.push_back(static_cast<uint8_t>(0x80 | (opcode & 0x0F))); // FIN + opcode
|
||||
|
||||
uint64_t len = payload.size();
|
||||
if (len < 126) {
|
||||
frame.push_back(static_cast<uint8_t>(0x80 | len)); // mask bit set
|
||||
} else if (len <= 0xFFFF) {
|
||||
frame.push_back(static_cast<uint8_t>(0x80 | 126));
|
||||
frame.push_back(static_cast<uint8_t>((len >> 8) & 0xFF));
|
||||
frame.push_back(static_cast<uint8_t>(len & 0xFF));
|
||||
} else {
|
||||
frame.push_back(static_cast<uint8_t>(0x80 | 127));
|
||||
for (int i = 7; i >= 0; i--) frame.push_back(static_cast<uint8_t>((len >> (8 * i)) & 0xFF));
|
||||
}
|
||||
|
||||
// Mask key (RFC requires mask for client → server frames).
|
||||
std::random_device rd;
|
||||
uint8_t mkey[4];
|
||||
for (auto& b : mkey) b = static_cast<uint8_t>(rd() & 0xff);
|
||||
for (int i = 0; i < 4; i++) frame.push_back(mkey[i]);
|
||||
|
||||
for (size_t i = 0; i < payload.size(); i++) {
|
||||
frame.push_back(static_cast<uint8_t>(payload[i]) ^ mkey[i & 3]);
|
||||
}
|
||||
|
||||
return send_all(sock, reinterpret_cast<const char*>(frame.data()), frame.size());
|
||||
}
|
||||
Reference in New Issue
Block a user