banban/tools/test_wss_receive.py

"""Exercise the real receive/upgrade loop with delayed TLS records."""
from pathlib import Path
import subprocess
import os

root = Path(__file__).resolve().parents[1]
lib = root / 'ql-application/threadx/interface/custom_app/lib/librws'
source = (lib / 'src/rws_socketpriv.c').read_bytes().decode('latin1')
receive = source[source.index('rws_bool rws_socket_recv('):].split('/* RT-Thread Team add */', 1)[0]
upgrade = source[source.index('void rws_socket_wait_handshake_responce('):].split('void rws_socket_send_disconnect(', 1)[0]
harness = r'''
#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "rws_handshake.h"
#define LIBRWS_USING_MBED_TLS
#define MBEDTLS_ERR_SSL_WANT_READ -2
#define MBEDTLS_ERR_SSL_WANT_WRITE -3
#define MBEDTLS_ERR_NET_RECV_FAILED -4
#define WSAEWOULDBLOCK 11
#define WSAEINPROGRESS 115
#define RWS_RECEIVE_BUFF_SIZE 1024
#define COMMAND_WAIT_HANDSHAKE_RESPONCE 3
#define COMMAND_INFORM_CONNECTED 4
#define COMMAND_INFORM_DISCONNECTED 5
#define rws_error_code_read_write_socket 4
#define rws_error_code_parse_handshake 5
#define LOG_E(...) ((void)0)
#define LOG_D(...) ((void)0)
#define LOG_I(...) ((void)0)
typedef int rws_bool;
enum { rws_false, rws_true };
typedef struct { int ssl_ctx; } tls;
typedef struct {
    const char *scheme, *expected_accept;
    tls *ssl;
    char *payload_buff, *received;
    size_t received_size, received_len;
    int socket, command, is_connected;
    unsigned int diag_rx_calls, diag_rx_bytes;
    const char *error;
} _rws_socket;
static int replies[16], reply_count, calls, closes;
static size_t offset;
static unsigned int ticks;
static const char response[] =
    "HTTP/1.1 101 Switching Protocols\r\nUpgrade: websocket\r\n"
    "Connection: Upgrade\r\nSec-WebSocket-Accept: s3pPLMBiTxaQ9kYGzzhZRbK+xOo=\r\n\r\nDATA";
static unsigned int ql_rtos_get_systicks(void) { return ticks; }
static void rws_thread_sleep(int ms) { ticks += (unsigned int)ms; }
static void rws_error_delete_clean(const char **e) { *e = NULL; }
static const char *rws_error_new_code_descr(int code, const char *text) { (void)code; return text; }
static void rws_socket_close(_rws_socket *s) { (void)s; ++closes; }
static void rws_socket_resize_received(_rws_socket *s, size_t size) {
    s->received = (char *)realloc(s->received, size);
    assert(s->received); s->received_size = size;
}
static int recv(int fd, char *buf, int size, int flags) {
    (void)fd; (void)buf; (void)size; (void)flags;
    assert(0 && "WSS must use TLS receive"); return -1;
}
static int mbedtls_ssl_read(int *ctx, unsigned char *buf, size_t size) {
    int ret = calls < reply_count ? replies[calls] : MBEDTLS_ERR_SSL_WANT_READ;
    (void)ctx; ++calls;
    if (ret > 0) {
        assert((size_t)ret <= size && offset + (size_t)ret <= sizeof(response) - 1);
        memcpy(buf, response + offset, (size_t)ret); offset += (size_t)ret;
    }
    return ret;
}
'''
harness += receive + upgrade
harness += r'''
static void reset(_rws_socket *s, tls *ctx) {
    free(s->received); memset(s, 0, sizeof(*s));
    s->scheme = "wss"; s->ssl = ctx;
    s->expected_accept = "s3pPLMBiTxaQ9kYGzzhZRbK+xOo=";
    s->payload_buff = (char *)malloc(RWS_RECEIVE_BUFF_SIZE); assert(s->payload_buff);
    s->command = COMMAND_WAIT_HANDSHAKE_RESPONCE;
    calls = closes = reply_count = 0; offset = 0; ticks = 0;
}
int main(void) {
    tls ctx = {0}; _rws_socket s = {0}; int want;
    for (want = -2; want >= -3; --want) {
        reset(&s, &ctx); replies[0] = want; reply_count = 1;
        assert(!rws_socket_recv(&s));
        assert(!s.error && s.command == COMMAND_WAIT_HANDSHAKE_RESPONCE && closes == 0);
        replies[1] = 17; replies[2] = want; replies[3] = (int)sizeof(response) - 1 - 17;
        reply_count = 4;
        rws_socket_wait_handshake_responce(&s);
        assert(s.is_connected && !s.error && s.command == COMMAND_INFORM_CONNECTED);
        assert(calls == 4 && closes == 0 && s.received_len == 4);
        assert(!memcmp(s.received, "DATA", 4)); free(s.payload_buff);
    }
    reset(&s, &ctx);
    rws_socket_wait_handshake_responce(&s);
    assert(!s.is_connected && s.error && closes == 1 && ticks >= 2000 && calls < 300);
    free(s.payload_buff);
    reset(&s, &ctx); ticks = (unsigned int)-1000;
    rws_socket_wait_handshake_responce(&s);
    assert(!s.is_connected && s.error && closes == 1 && calls < 300); free(s.payload_buff);
    for (want = 0; want >= -99; want -= 99) {
        reset(&s, &ctx); replies[0] = want; reply_count = 1;
        rws_socket_wait_handshake_responce(&s);
        assert(!s.is_connected && s.error && closes == 1 && calls == 1);
        assert(s.command == COMMAND_INFORM_DISCONNECTED); free(s.payload_buff);
    }
    reset(&s, NULL); rws_socket_wait_handshake_responce(&s);
    assert(s.error && closes == 1 && calls == 0); free(s.payload_buff);
    reset(&s, &ctx); s.expected_accept = "wrong";
    replies[0] = (int)sizeof(response) - 1; reply_count = 1;
    rws_socket_wait_handshake_responce(&s);
    assert(!s.is_connected && s.error && closes == 1); free(s.payload_buff);
    free(s.received);
    puts("PASS: actual TLS receive/upgrade loop: WANT_READ/WRITE, split HTTP 101, trailing frame bytes, timeout/wrap, EOF/fatal, null TLS, invalid Accept");
    return 0;
}
'''
output = root / 'artifacts/wss'
output.mkdir(parents=True, exist_ok=True)
cfile = output / 'test_wss_receive.c'
cfile.write_text(harness, encoding='utf-8')
exe = output / 'test_wss_receive.exe'
if os.name == "nt":
    subprocess.run(['cl', '/nologo', '/W4', '/WX', '/wd4018', '/TC',
                    '/I' + str(lib / 'inc'), str(cfile), str(lib / 'src/rws_handshake.c'),
                    '/Fo' + str(output) + '/', '/Fe' + str(exe)], check=True)
else:
    subprocess.run(['cc', '-std=c99', '-Wall', '-Wextra', '-Werror', '-Wno-sign-compare', '-I' + str(lib / 'inc'), str(cfile), str(lib / 'src/rws_handshake.c'), '-o', str(exe)], check=True)
subprocess.run([str(exe)], check=True)