Major refactoring

This commit is contained in:
2026-02-25 10:44:43 +01:00
parent bac3477991
commit ef1d65ad2b
45 changed files with 5404 additions and 2120 deletions
+232
View File
@@ -0,0 +1,232 @@
#if defined(MAIN_SIMULATION) || defined(MAIN_TEST)
#define QUAKEY_ENABLE_MOCKS
#endif
#include <stdint.h>
#include <quakey.h>
#include "basic.h"
bool streq(string s1, string s2)
{
if (s1.len != s2.len)
return false;
for (int i = 0; i < s1.len; i++)
if (s1.ptr[i] != s2.ptr[i])
return false;
return true;
}
// Returns the current time in nanoseconds since
// an unspecified time in the past (useful to calculate
// elapsed time intervals)
Time get_current_time(void)
{
#ifdef _WIN32
{
int64_t count;
int64_t freq;
int ok;
ok = QueryPerformanceCounter((LARGE_INTEGER*) &count);
if (!ok) return INVALID_TIME;
ok = QueryPerformanceFrequency((LARGE_INTEGER*) &freq);
if (!ok) return INVALID_TIME;
uint64_t res = 1000000000 * (double) count / freq;
return res;
}
#else
{
struct timespec time;
if (clock_gettime(CLOCK_REALTIME, &time))
return INVALID_TIME;
uint64_t res;
uint64_t sec = time.tv_sec;
if (sec > UINT64_MAX / 1000000000)
return INVALID_TIME;
res = sec * 1000000000;
uint64_t nsec = time.tv_nsec;
if (res > UINT64_MAX - nsec)
return INVALID_TIME;
res += nsec;
return res;
}
#endif
}
void nearest_deadline(Time *a, Time b)
{
if (*a == INVALID_TIME || *a > b)
*a = b;
}
int deadline_to_timeout(Time deadline, Time current_time)
{
if (deadline == INVALID_TIME)
return -1;
return (deadline - current_time) / 1000000;
}
bool getargb(int argc, char **argv, char *name)
{
for (int i = 0; i < argc; i++)
if (!strcmp(argv[i], name))
return true;
return false;
}
string getargs(int argc, char **argv, char *name, char *fallback)
{
for (int i = 0; i < argc; i++)
if (!strcmp(argv[i], name)) {
i++;
if (i == argc)
break;
return (string) { argv[i], strlen(argv[i]) };
}
return (string) { fallback, strlen(fallback) };
}
int getargi(int argc, char **argv, char *name, int fallback)
{
for (int i = 0; i < argc; i++)
if (!strcmp(argv[i], name)) {
i++;
if (i == argc)
break;
errno = 0;
char *end;
long val = strtol(argv[i], &end, 10);
if (end == argv[i] || *end != '\0' || errno == ERANGE)
break;
if (val < INT_MIN || val > INT_MAX)
break;
return (int) val;
}
return fallback;
}
void append_hex_as_str(char *out, SHA256 hash)
{
char table[] = "0123456789abcdef";
for (int i = 0; i < (int) sizeof(hash); i++) {
out[(i << 1) + 0] = table[(uint8_t) hash.data[i] >> 4];
out[(i << 1) + 1] = table[(uint8_t) hash.data[i] & 0xF];
}
}
// TODO: check this function
bool addr_lower(Address a, Address b)
{
if (a.is_ipv4) {
if (!b.is_ipv4)
return true;
if (a.ipv4.data < b.ipv4.data)
return true;
if (a.ipv4.data == b.ipv4.data &&
a.port < b.port)
return true;
return false;
} else {
if (b.is_ipv4)
return false;
for (int i = 0; i < 8; i++) {
if (a.ipv6.data[i] < b.ipv6.data[i])
return true;
if (a.ipv6.data[i] > b.ipv6.data[i])
return false;
}
if (a.port < b.port)
return true;
return false;
}
}
bool addr_eql(Address a, Address b)
{
if (a.is_ipv4 != b.is_ipv4)
return false;
if (a.port != b.port)
return false;
if (a.is_ipv4) {
if (memcmp(&a.ipv4, &b.ipv4, sizeof(a.ipv4)))
return false;
} else {
if (memcmp(&a.ipv6, &b.ipv6, sizeof(a.ipv6)))
return false;
}
return true;
}
int parse_addr_arg(char *arg, Address *out)
{
int len = strlen(arg);
int i = 0;
while (i < len && arg[i] != ':')
i++;
if (i == len)
return -1; // No ':' character.
arg[i] = '\0';
IPv4 ipv4;
int ret = inet_pton(AF_INET, arg, &ipv4);
arg[i] = ':';
if (ret != 1)
return -1;
errno = 0;
ret = atoi(arg + i + 1);
if (ret == 0 && errno != 0)
return -1;
out->ipv4 = ipv4;
out->is_ipv4 = true;
out->port = ret;
return 0;
}
void addr_sort(Address *addrs, int count)
{
for (int i = 0; i < count; i++) {
int k = i; // Index of the lowest address in [i, num_nodes-1]
for (int j = i+1; j < count; j++) {
if (addr_lower(addrs[j], addrs[k]))
k = j;
}
Address tmp = addrs[i];
addrs[i] = addrs[k];
addrs[k] = tmp;
}
}
+61
View File
@@ -0,0 +1,61 @@
#ifndef BASIC_INCLUDED
#define BASIC_INCLUDED
#include <stdint.h>
#include <stdbool.h>
typedef struct {
char data[32];
} SHA256;
typedef struct {
uint32_t data;
} IPv4;
typedef struct {
uint16_t data[8];
} IPv6;
typedef struct {
union {
IPv4 ipv4;
IPv6 ipv6;
};
bool is_ipv4;
uint16_t port;
} Address;
typedef struct {
char *ptr;
int len;
} string;
typedef uint64_t Time;
#define INVALID_TIME ((Time) -1)
#define S(X) ((string) { (X), (int) sizeof(X)-1 })
#define MIN(X, Y) ((X) < (Y) ? (X) : (Y))
#define MAX(X, Y) ((X) > (Y) ? (X) : (Y))
#define UNREACHABLE __builtin_trap();
bool streq(string s1, string s2);
Time get_current_time(void);
void nearest_deadline(Time *a, Time b);
int deadline_to_timeout(Time deadline, Time current_time);
bool getargb(int argc, char **argv, char *name);
string getargs(int argc, char **argv, char *name, char *fallback);
int getargi(int argc, char **argv, char *name, int fallback);
void append_hex_as_str(char *out, SHA256 hash);
bool addr_eql(Address a, Address b);
bool addr_lower(Address a, Address b);
int parse_addr_arg(char *arg, Address *out);
void addr_sort(Address *addrs, int count);
#endif // BASIC_INCLUDED
+311
View File
@@ -0,0 +1,311 @@
#if defined(MAIN_SIMULATION) || defined(MAIN_TEST)
#define QUAKEY_ENABLE_MOCKS
#endif
#include <quakey.h>
#include <stdint.h>
#include <assert.h>
#include "byte_queue.h"
// This is the implementation of a byte queue useful
// for systems that need to process engs of bytes.
//
// It features sticky errors, a zero-copy interface,
// and a safe mechanism to patch previously written
// bytes.
//
// Only up to 4GB of data can be stored at once.
// Initialize the queue
void byte_queue_init(ByteQueue *queue, uint32_t limit)
{
queue->flags = 0;
queue->head = 0;
queue->size = 0;
queue->used = 0;
queue->curs = 0;
queue->limit = limit;
queue->data = NULL;
queue->read_target = NULL;
}
// Deinitialize the queue
void byte_queue_free(ByteQueue *queue)
{
if (queue->read_target) {
if (queue->read_target != queue->data)
free(queue->read_target);
queue->read_target = NULL;
queue->read_target_size = 0;
}
free(queue->data);
queue->data = NULL;
}
int byte_queue_error(ByteQueue *queue)
{
return queue->flags & BYTE_QUEUE_ERROR;
}
int byte_queue_empty(ByteQueue *queue)
{
return queue->used == 0;
}
int byte_queue_full(ByteQueue *queue)
{
return queue->used == queue->limit;
}
// Start a read operation on the queue.
//
// This function returnes the pointer to the memory region containing the bytes
// to read. Callers can't read more than [*len] bytes from it. To complete the
// read, the [byte_queue_read_ack] function must be called with the number of
// bytes that were acknowledged by the caller.
//
// Note:
// - You can't have more than one pending read.
ByteView byte_queue_read_buf(ByteQueue *queue)
{
if (queue->flags & BYTE_QUEUE_ERROR)
return (ByteView) {NULL, 0};
assert((queue->flags & BYTE_QUEUE_READ) == 0);
queue->flags |= BYTE_QUEUE_READ;
queue->read_target = queue->data;
queue->read_target_size = queue->size;
if (queue->data == NULL)
return (ByteView) {NULL, 0};
return (ByteView) { queue->data + queue->head, queue->used };
}
// Complete a previously started operation on the queue.
void byte_queue_read_ack(ByteQueue *queue, uint32_t num)
{
if (queue->flags & BYTE_QUEUE_ERROR)
return;
if ((queue->flags & BYTE_QUEUE_READ) == 0)
return;
queue->flags &= ~BYTE_QUEUE_READ;
assert((uint32_t) num <= queue->used);
queue->head += (uint32_t) num;
queue->used -= (uint32_t) num;
queue->curs += (uint32_t) num;
if (queue->read_target) {
if (queue->read_target != queue->data)
free(queue->read_target);
queue->read_target = NULL;
queue->read_target_size = 0;
}
}
bool byte_queue_reading(ByteQueue *queue)
{
return (queue->flags & BYTE_QUEUE_READ) != 0;
}
ByteView byte_queue_write_buf(ByteQueue *queue)
{
if ((queue->flags & BYTE_QUEUE_ERROR) || queue->data == NULL)
return (ByteView) {NULL, 0};
assert((queue->flags & BYTE_QUEUE_WRITE) == 0);
queue->flags |= BYTE_QUEUE_WRITE;
return (ByteView) {
queue->data + (queue->head + queue->used),
queue->size - (queue->head + queue->used),
};
}
void byte_queue_write_ack(ByteQueue *queue, uint32_t num)
{
if (queue->flags & BYTE_QUEUE_ERROR)
return;
if ((queue->flags & BYTE_QUEUE_WRITE) == 0)
return;
queue->flags &= ~BYTE_QUEUE_WRITE;
queue->used += num;
}
// Sets the minimum capacity for the next write operation
// and returns 1 if the content of the queue was moved, else
// 0 is returned.
//
// You must not call this function while a write is pending.
// In other words, you must do this:
//
// byte_queue_write_setmincap(queue, mincap);
// dst = byte_queue_write_buf(queue, &cap);
// ...
// byte_queue_write_ack(num);
//
// And NOT this:
//
// dst = byte_queue_write_buf(queue, &cap);
// byte_queue_write_setmincap(queue, mincap); <-- BAD
// ...
// byte_queue_write_ack(num);
//
int byte_queue_write_setmincap(ByteQueue *queue, uint32_t mincap)
{
// Sticky error
if (queue->flags & BYTE_QUEUE_ERROR)
return 0;
// In general, the queue's contents look like this:
//
// size
// v
// [___xxxxxxxxxxxx________]
// ^ ^ ^
// 0 head head + used
//
// This function needs to make sure that at least [mincap]
// bytes are available on the right side of the content.
//
// We have 3 cases:
//
// 1) If there is enough memory already, this function doesn't
// need to do anything.
//
// 2) If there isn't enough memory on the right but there is
// enough free memory if we cound the left unused region,
// then the content is moved back to the
// start of the buffer.
//
// 3) If there isn't enough memory considering both sides, this
// function needs to allocate a new buffer.
//
// If there are pending read or write operations, the application
// is holding pointers to the buffer, so we need to make sure
// to not invalidate them. The only real problem is pending reads
// since this function can only be called before starting a write
// opearation.
//
// To avoid invalidating the read pointer when we allocate a new
// buffer, we don't free the old buffer. Instead, we store the
// pointer in the "old" field so that the read ack function can
// free it.
//
// To avoid invalidating the pointer when we are moving back the
// content since there is enough memory at the start of the buffer,
// we just avoid that. Even if there is enough memory considering
// left and right free regions, we allocate a new buffer.
assert((queue->flags & BYTE_QUEUE_WRITE) == 0);
uint32_t total_free_space = queue->size - queue->used;
uint32_t free_space_after_data = queue->size - queue->used - queue->head;
int moved = 0;
if (free_space_after_data < mincap) {
if (total_free_space < mincap || (queue->read_target == queue->data)) {
// Resize required
if (queue->used + mincap > queue->limit) {
queue->flags |= BYTE_QUEUE_ERROR;
return 0;
}
uint32_t size;
if (queue->size > UINT32_MAX / 2)
size = UINT32_MAX;
else
size = 2 * queue->size;
if (size < queue->used + mincap)
size = queue->used + mincap;
if (size > queue->limit)
size = queue->limit;
uint8_t *data = malloc(size);
if (!data) {
queue->flags |= BYTE_QUEUE_ERROR;
return 0;
}
if (queue->used > 0)
memcpy(data, queue->data + queue->head, queue->used);
if (queue->read_target != queue->data)
free(queue->data);
queue->data = data;
queue->head = 0;
queue->size = size;
} else {
// Move required
memmove(queue->data, queue->data + queue->head, queue->used);
queue->head = 0;
}
moved = 1;
}
return moved;
}
void byte_queue_write(ByteQueue *queue, void *ptr, uint32_t len)
{
byte_queue_write_setmincap(queue, len);
ByteView dst = byte_queue_write_buf(queue);
if (dst.ptr) {
memcpy(dst.ptr, ptr, len);
byte_queue_write_ack(queue, len);
}
}
ByteQueueOffset byte_queue_offset(ByteQueue *queue)
{
if (queue->flags & BYTE_QUEUE_ERROR)
return (ByteQueueOffset) { 0 };
return (ByteQueueOffset) { queue->curs + queue->used };
}
void byte_queue_patch(ByteQueue *queue, ByteQueueOffset off,
void *src, uint32_t len)
{
if (queue->flags & BYTE_QUEUE_ERROR)
return;
// Check that the offset is in range
assert(off >= queue->curs && off - queue->curs < queue->used);
// Check that the length is in range
assert(len <= queue->used - (off - queue->curs));
// Perform the patch
uint8_t *dst = queue->data + queue->head + (off - queue->curs);
memcpy(dst, src, len);
}
uint32_t byte_queue_size_from_offset(ByteQueue *queue, ByteQueueOffset off)
{
return queue->curs + queue->used - off;
}
void byte_queue_remove_from_offset(ByteQueue *queue, ByteQueueOffset offset)
{
if (queue->flags & BYTE_QUEUE_ERROR)
return;
uint64_t num = (queue->curs + queue->used) - offset;
assert(num <= queue->used);
queue->used -= num;
}
+54
View File
@@ -0,0 +1,54 @@
#ifndef BYTE_QUEUE_INCLUDED
#define BYTE_QUEUE_INCLUDED
#include <stddef.h>
#include "basic.h"
typedef struct {
uint8_t *ptr;
size_t len;
} ByteView;
typedef struct {
uint64_t curs;
uint8_t* data;
uint32_t head;
uint32_t size;
uint32_t used;
uint32_t limit;
uint8_t* read_target;
uint32_t read_target_size;
int flags;
} ByteQueue;
typedef uint64_t ByteQueueOffset;
enum {
BYTE_QUEUE_ERROR = 1 << 0,
BYTE_QUEUE_READ = 1 << 1,
BYTE_QUEUE_WRITE = 1 << 2,
};
void byte_queue_init(ByteQueue *queue, uint32_t limit);
void byte_queue_free(ByteQueue *queue);
int byte_queue_error(ByteQueue *queue);
int byte_queue_empty(ByteQueue *queue);
int byte_queue_full(ByteQueue *queue);
ByteView byte_queue_read_buf(ByteQueue *queue);
void byte_queue_read_ack(ByteQueue *queue, uint32_t num);
bool byte_queue_reading(ByteQueue *queue);
ByteView byte_queue_write_buf(ByteQueue *queue);
void byte_queue_write_ack(ByteQueue *queue, uint32_t num);
int byte_queue_write_setmincap(ByteQueue *queue, uint32_t mincap);
void byte_queue_write(ByteQueue *queue, void *ptr, uint32_t len);
ByteQueueOffset byte_queue_offset(ByteQueue *queue);
void byte_queue_patch(ByteQueue *queue, ByteQueueOffset off, void *src, uint32_t len);
uint32_t byte_queue_size_from_offset(ByteQueue *queue, ByteQueueOffset off);
void byte_queue_remove_from_offset(ByteQueue *queue, ByteQueueOffset offset);
#endif // BYTE_QUEUE_INCLUDED
+459
View File
@@ -0,0 +1,459 @@
#if defined(MAIN_SIMULATION) || defined(MAIN_TEST)
#define QUAKEY_ENABLE_MOCKS
#endif
#include <stdint.h>
#include <quakey.h>
#include "file_system.h"
int rename_file_or_dir(string oldpath, string newpath);
bool file_exists(string path)
{
char zt[1<<10];
if (path.len >= (int) sizeof(zt))
return false;
memcpy(zt, path.ptr, path.len);
zt[path.len] = '\0';
#ifdef __linux__
return access(zt, F_OK) == 0;
#endif
#ifdef _WIN32
DWORD attrs = GetFileAttributesA(zt);
return attrs != INVALID_FILE_ATTRIBUTES;
#endif
}
int file_open(string path, Handle *fd)
{
#ifdef __linux__
char zt[1<<10];
if (path.len >= (int) sizeof(zt))
return -1;
memcpy(zt, path.ptr, path.len);
zt[path.len] = '\0';
int ret = open(zt, O_RDWR | O_CREAT | O_APPEND, 0644);
if (ret < 0)
return -1;
*fd = (Handle) { (uint64_t) ret };
return 0;
#endif
#ifdef _WIN32
WCHAR wpath[MAX_PATH];
MultiByteToWideChar(CP_UTF8, 0, path.ptr, path.len, wpath, MAX_PATH);
wpath[path.len] = L'\0';
HANDLE h = CreateFileW(
wpath,
GENERIC_WRITE | GENERIC_READ,
0,
NULL,
OPEN_ALWAYS,
FILE_ATTRIBUTE_NORMAL | FILE_FLAG_WRITE_THROUGH,
NULL
);
if (h == INVALID_HANDLE_VALUE)
return -1;
*fd = (Handle) { (uint64_t) h };
return 0;
#endif
}
void file_close(Handle fd)
{
#ifdef __linux__
close((int) fd.data);
#endif
#ifdef _WIN32
CloseHandle((HANDLE) fd.data);
#endif
}
int file_truncate(Handle fd, size_t new_size)
{
#ifdef __linux__
if (ftruncate((int) fd.data, new_size) < 0)
return -1;
return 0;
#endif
#ifdef _WIN32
return -1; // TODO: Not implemented
#endif
}
int file_set_offset(Handle fd, int off)
{
#ifdef __linux__
off_t ret = lseek((int) fd.data, off, SEEK_SET);
if (ret < 0)
return -1;
return 0;
#endif
#ifdef _WIN32
LARGE_INTEGER distance;
distance.QuadPart = off;
if (!SetFilePointer((HANDLE) fd.data, distance.LowPart, &distance.HighPart, FILE_BEGIN))
if (GetLastError() != 0)
return -1;
return 0;
#endif
}
int file_get_offset(Handle fd, int *off)
{
#ifdef __linux__
off_t ret = lseek((int) fd.data, 0, SEEK_CUR);
if (ret < 0)
return -1;
*off = (int) ret;
return 0;
#endif
#ifdef _WIN32
DWORD pos = SetFilePointer((HANDLE) fd.data, 0, NULL, FILE_CURRENT);
if (pos == INVALID_SET_FILE_POINTER && GetLastError() != 0)
return -1;
*off = (int) pos;
return 0;
#endif
}
int file_lock(Handle fd)
{
#ifdef __linux__
if (flock((int) fd.data, LOCK_EX) < 0)
return -1;
return 0;
#endif
#ifdef _WIN32
if (!LockFile((HANDLE) fd.data, 0, 0, MAXDWORD, MAXDWORD))
return -1;
return 0;
#endif
}
int file_unlock(Handle fd)
{
#ifdef __linux__
if (flock((int) fd.data, LOCK_UN) < 0)
return -1;
return 0;
#endif
#ifdef _WIN32
if (!UnlockFile((HANDLE) fd.data, 0, 0, MAXDWORD, MAXDWORD))
return -1;
return 0;
#endif
}
int file_sync(Handle fd)
{
#ifdef __linux__
if (fsync((int) fd.data) < 0)
return -1;
return 0;
#endif
#ifdef _WIN32
if (!FlushFileBuffers((HANDLE) fd.data))
return -1;
return 0;
#endif
}
int file_read(Handle fd, char *dst, int max)
{
#ifdef __linux__
return read((int) fd.data, dst, max);
#endif
#ifdef _WIN32
DWORD num;
if (!ReadFile((HANDLE) fd.data, dst, max, &num, NULL))
return -1;
if (num > INT_MAX)
return -1;
return num;
#endif
}
int file_write(Handle fd, char *src, int len)
{
#ifdef __linux__
return write((int) fd.data, src, len);
#endif
#ifdef _WIN32
DWORD num;
if (!WriteFile((HANDLE) fd.data, src, len, &num, NULL))
return -1;
if (num > INT_MAX)
return -1;
return num;
#endif
}
int file_size(Handle fd, size_t *len)
{
#ifdef __linux__
struct stat buf;
if (fstat((int) fd.data, &buf) < 0)
return -1;
if (buf.st_size < 0 || (uint64_t) buf.st_size > SIZE_MAX)
return -1;
*len = (size_t) buf.st_size;
return 0;
#endif
#ifdef _WIN32
LARGE_INTEGER buf;
if (!GetFileSizeEx((HANDLE) fd.data, &buf))
return -1;
if (buf.QuadPart < 0 || (uint64_t) buf.QuadPart > SIZE_MAX)
return -1;
*len = buf.QuadPart;
return 0;
#endif
}
int create_dir(string path)
{
char zt[PATH_MAX];
if (path.len >= (int) sizeof(zt))
return -1;
memcpy(zt, path.ptr, path.len);
zt[path.len] = '\0';
#ifdef _WIN32
if (_mkdir(zt) < 0)
return -1;
#else
if (mkdir(zt, 0766))
return -1;
#endif
return 0;
}
int rename_file_or_dir(string oldpath, string newpath)
{
char oldpath_zt[PATH_MAX];
if (oldpath.len >= (int) sizeof(oldpath_zt))
return -1;
memcpy(oldpath_zt, oldpath.ptr, oldpath.len);
oldpath_zt[oldpath.len] = '\0';
char newpath_zt[PATH_MAX];
if (newpath.len >= (int) sizeof(newpath_zt))
return -1;
memcpy(newpath_zt, newpath.ptr, newpath.len);
newpath_zt[newpath.len] = '\0';
if (rename(oldpath_zt, newpath_zt))
return -1;
return 0;
}
int remove_file_or_dir(string path)
{
char path_zt[PATH_MAX];
if (path.len >= (int) sizeof(path_zt))
return -1;
memcpy(path_zt, path.ptr, path.len);
path_zt[path.len] = '\0';
if (remove(path_zt))
return -1;
return 0;
}
int get_full_path(string path, char *dst)
{
char path_zt[PATH_MAX];
if (path.len >= (int) sizeof(path_zt))
return -1;
memcpy(path_zt, path.ptr, path.len);
path_zt[path.len] = '\0';
#ifdef __linux__
if (realpath(path_zt, dst) == NULL)
return -1;
#endif
#ifdef _WIN32
if (_fullpath(dst, path_zt, PATH_MAX) == NULL)
return -1;
#endif
size_t path_len = strlen(dst);
if (path_len > 0 && dst[path_len-1] == '/')
dst[path_len-1] = '\0';
return 0;
}
int file_read_all(string path, string *data)
{
Handle fd;
int ret = file_open(path, &fd);
if (ret < 0)
return -1;
size_t len;
ret = file_size(fd, &len);
if (ret < 0) {
file_close(fd);
return -1;
}
char *dst = malloc(len);
if (dst == NULL) {
file_close(fd);
return -1;
}
int copied = 0;
while ((size_t) copied < len) {
ret = file_read(fd, dst + copied, len - copied);
if (ret < 0) {
free(dst);
file_close(fd);
return -1;
}
copied += ret;
}
*data = (string) { dst, len };
file_close(fd);
return 0;
}
#ifdef _WIN32
int directory_scanner_init(DirectoryScanner *scanner, string path)
{
char pattern[PATH_MAX];
int ret = snprintf(pattern, sizeof(pattern), "%.*s\\*", path.len, path.ptr);
if (ret < 0 || ret >= (int) sizeof(pattern))
return -1;
scanner->handle = FindFirstFileA(pattern, &scanner->find_data);
if (scanner->handle == INVALID_HANDLE_VALUE) {
if (GetLastError() == ERROR_FILE_NOT_FOUND) {
scanner->done = true;
return 0;
}
return -1;
}
scanner->done = false;
scanner->first = true;
return 0;
}
int directory_scanner_next(DirectoryScanner *scanner, string *name)
{
if (scanner->done)
return 1;
if (!scanner->first) {
BOOL ok = FindNextFileA(scanner->handle, &scanner->find_data);
if (!ok) {
scanner->done = true;
if (GetLastError() == ERROR_NO_MORE_FILES)
return 1;
return -1;
}
} else {
scanner->first = false;
}
char *p = scanner->find_data.cFileName;
*name = (string) { p, strlen(p) };
return 0;
}
void directory_scanner_free(DirectoryScanner *scanner)
{
FindClose(scanner->handle);
}
#else
int directory_scanner_init(DirectoryScanner *scanner, string path)
{
char path_copy[PATH_MAX];
if (path.len >= PATH_MAX)
return -1;
memcpy(path_copy, path.ptr, path.len);
path_copy[path.len] = '\0';
scanner->d = opendir(path_copy);
if (scanner->d == NULL) {
scanner->done = true;
return -1;
}
scanner->done = false;
return 0;
}
int directory_scanner_next(DirectoryScanner *scanner, string *name)
{
if (scanner->done)
return 1;
scanner->e = readdir(scanner->d);
if (scanner->e == NULL) {
scanner->done = true;
return 1;
}
*name = (string) { scanner->e->d_name, strlen(scanner->e->d_name) };
return 0;
}
void directory_scanner_free(DirectoryScanner *scanner)
{
closedir(scanner->d);
}
#endif
int file_read_exact(Handle handle, char *dst, int len)
{
int copied = 0;
while (copied < len) {
int ret = file_read(handle, dst + copied, len - copied);
if (ret < 0)
return -1;
if (ret == 0)
return 0; // EOF
copied += ret;
}
return copied;
}
int file_write_exact(Handle handle, char *src, int len)
{
int copied = 0;
while (copied < len) {
int ret = file_write(handle, src + copied, len - copied);
if (ret < 0)
return -1;
copied += ret;
}
return 0;
}
+58
View File
@@ -0,0 +1,58 @@
#ifndef FILE_SYSTEM_INCLUDED
#define FILE_SYSTEM_INCLUDED
#if defined(MAIN_SIMULATION) || defined(MAIN_TEST)
#define QUAKEY_ENABLE_MOCKS
#endif
#include <stdint.h>
#include <quakey.h>
#include "basic.h"
typedef struct {
uint64_t data;
} Handle;
#ifdef _WIN32
typedef struct {
HANDLE handle;
WIN32_FIND_DATA find_data;
bool first;
bool done;
} DirectoryScanner;
#else
typedef struct {
DIR *d;
struct dirent *e;
bool done;
} DirectoryScanner;
#endif
bool file_exists(string path);
int file_open(string path, Handle *fd);
void file_close(Handle fd);
int file_truncate(Handle fd, size_t new_size);
int file_set_offset(Handle fd, int off);
int file_get_offset(Handle fd, int *off);
int file_lock(Handle fd);
int file_unlock(Handle fd);
int file_sync(Handle fd);
int file_read(Handle fd, char *dst, int max);
int file_write(Handle fd, char *src, int len);
int file_size(Handle fd, size_t *len);
int file_write_atomic(string path, string content);
int create_dir(string path);
int rename_file_or_dir(string oldpath, string newpath);
int remove_file_or_dir(string path);
int get_full_path(string path, char *dst);
int file_read_all(string path, string *data);
int directory_scanner_init(DirectoryScanner *scanner, string path);
int directory_scanner_next(DirectoryScanner *scanner, string *name);
void directory_scanner_free(DirectoryScanner *scanner);
int file_read_exact(Handle handle, char *dst, int len);
int file_write_exact(Handle handle, char *src, int len);
#endif // FILE_SYSTEM_INCLUDED
+1633
View File
File diff suppressed because it is too large Load Diff
+161
View File
@@ -0,0 +1,161 @@
#ifndef HTTP_PARSE_INCLUDED
#define HTTP_PARSE_INCLUDED
#include "basic.h"
#define CHTTP_MAX_HEADERS 32
typedef struct {
unsigned int data;
} CHTTP_IPv4;
typedef struct {
unsigned short data[8];
} CHTTP_IPv6;
typedef enum {
CHTTP_HOST_MODE_VOID = 0,
CHTTP_HOST_MODE_NAME,
CHTTP_HOST_MODE_IPV4,
CHTTP_HOST_MODE_IPV6,
} CHTTP_HostMode;
typedef struct {
CHTTP_HostMode mode;
string text;
union {
string name;
CHTTP_IPv4 ipv4;
CHTTP_IPv6 ipv6;
};
} CHTTP_Host;
typedef struct {
string userinfo;
CHTTP_Host host;
int port;
} CHTTP_Authority;
// ZII
typedef struct {
string scheme;
CHTTP_Authority authority;
string path;
string query;
string fragment;
} CHTTP_URL;
typedef enum {
CHTTP_METHOD_GET,
CHTTP_METHOD_HEAD,
CHTTP_METHOD_POST,
CHTTP_METHOD_PUT,
CHTTP_METHOD_DELETE,
CHTTP_METHOD_CONNECT,
CHTTP_METHOD_OPTIONS,
CHTTP_METHOD_TRACE,
CHTTP_METHOD_PATCH,
} CHTTP_Method;
typedef struct {
string name;
string value;
} CHTTP_Header;
typedef struct {
bool secure;
CHTTP_Method method;
CHTTP_URL url;
int minor;
int num_headers;
CHTTP_Header headers[CHTTP_MAX_HEADERS];
string body;
} CHTTP_Request;
typedef struct {
void* context;
int minor;
int status;
string reason;
int num_headers;
CHTTP_Header headers[CHTTP_MAX_HEADERS];
string body;
} CHTTP_Response;
int chttp_parse_ipv4(char *src, int len, CHTTP_IPv4 *ipv4);
int chttp_parse_ipv6(char *src, int len, CHTTP_IPv6 *ipv6);
int chttp_parse_url(char *src, int len, CHTTP_URL *url);
int chttp_parse_request(char *src, int len, CHTTP_Request *req);
int chttp_parse_response(char *src, int len, CHTTP_Response *res);
int chttp_find_header(CHTTP_Header *headers, int num_headers, string name);
string chttp_get_cookie(CHTTP_Request *req, string name);
string chttp_get_param(string body, string str, char *mem, int cap);
int chttp_get_param_i(string body, string str);
// Checks whether the request was meant for the host with the given
// domain an port. If port is -1, the default value of 80 is assumed.
bool chttp_match_host(CHTTP_Request *req, string domain, int port);
// Date and cookie types for Set-Cookie header parsing
typedef enum {
CHTTP_WEEKDAY_MON,
CHTTP_WEEKDAY_TUE,
CHTTP_WEEKDAY_WED,
CHTTP_WEEKDAY_THU,
CHTTP_WEEKDAY_FRI,
CHTTP_WEEKDAY_SAT,
CHTTP_WEEKDAY_SUN,
} CHTTP_WeekDay;
typedef enum {
CHTTP_MONTH_JAN,
CHTTP_MONTH_FEB,
CHTTP_MONTH_MAR,
CHTTP_MONTH_APR,
CHTTP_MONTH_MAY,
CHTTP_MONTH_JUN,
CHTTP_MONTH_JUL,
CHTTP_MONTH_AUG,
CHTTP_MONTH_SEP,
CHTTP_MONTH_OCT,
CHTTP_MONTH_NOV,
CHTTP_MONTH_DEC,
} CHTTP_Month;
typedef struct {
CHTTP_WeekDay week_day;
int day;
CHTTP_Month month;
int year;
int hour;
int minute;
int second;
} CHTTP_Date;
typedef struct {
string name;
string value;
bool secure;
bool chttp_only;
bool have_date;
CHTTP_Date date;
bool have_max_age;
uint32_t max_age;
bool have_domain;
string domain;
bool have_path;
string path;
} CHTTP_SetCookie;
// Parses a Set-Cookie header value
// Returns 0 on success, -1 on error
int chttp_parse_set_cookie(string str, CHTTP_SetCookie *out);
#endif // HTTP_PARSE_INCLUDED
+390
View File
@@ -0,0 +1,390 @@
#include <quakey.h>
#include <assert.h>
#include "http_server.h"
static string reason_phrase(int status)
{
switch(status) {
case 100: return S("Continue");
case 101: return S("Switching Protocols");
case 102: return S("Processing");
case 200: return S("OK");
case 201: return S("Created");
case 202: return S("Accepted");
case 203: return S("Non-Authoritative Information");
case 204: return S("No Content");
case 205: return S("Reset Content");
case 206: return S("Partial Content");
case 207: return S("Multi-Status");
case 208: return S("Already Reported");
case 300: return S("Multiple Choices");
case 301: return S("Moved Permanently");
case 302: return S("Found");
case 303: return S("See Other");
case 304: return S("Not Modified");
case 305: return S("Use Proxy");
case 306: return S("Switch Proxy");
case 307: return S("Temporary Redirect");
case 308: return S("Permanent Redirect");
case 400: return S("Bad Request");
case 401: return S("Unauthorized");
case 402: return S("Payment Required");
case 403: return S("Forbidden");
case 404: return S("Not Found");
case 405: return S("Method Not Allowed");
case 406: return S("Not Acceptable");
case 407: return S("Proxy Authentication Required");
case 408: return S("Request Timeout");
case 409: return S("Conflict");
case 410: return S("Gone");
case 411: return S("Length Required");
case 412: return S("Precondition Failed");
case 413: return S("Request Entity Too Large");
case 414: return S("Request-URI Too Long");
case 415: return S("Unsupported Media Type");
case 416: return S("Requested Range Not Satisfiable");
case 417: return S("Expectation Failed");
case 418: return S("I'm a teapot");
case 420: return S("Enhance your calm");
case 422: return S("Unprocessable Entity");
case 426: return S("Upgrade Required");
case 429: return S("Too many requests");
case 431: return S("Request Header Fields Too Large");
case 449: return S("Retry With");
case 451: return S("Unavailable For Legal Reasons");
case 500: return S("Internal Server Error");
case 501: return S("Not Implemented");
case 502: return S("Bad Gateway");
case 503: return S("Service Unavailable");
case 504: return S("Gateway Timeout");
case 505: return S("HTTP Version Not Supported");
case 509: return S("Bandwidth Limit Exceeded");
}
return S("???");
}
int http_server_init(HTTP_Server *server, int max_conns)
{
server->conns = malloc(max_conns * sizeof(HTTP_Conn));
if (server->conns == NULL)
return -1;
for (int i = 0; i < max_conns; i++) {
server->conns[i].state = HTTP_CONN_STATE_FREE;
server->conns[i].gen = 1;
}
server->max_conns = max_conns;
server->num_conns = 0;
int ret = tcp_init(&server->tcp, max_conns);
if (ret < 0) {
free(server->conns);
return -1;
}
return 0;
}
void http_server_free(HTTP_Server *server)
{
tcp_free(&server->tcp);
free(server->conns);
}
int http_server_listen_tcp(HTTP_Server *server, string addr, uint16_t port)
{
int ret = tcp_listen_tcp(&server->tcp, addr, port, true, 128);
if (ret < 0)
return -1;
return 0;
}
int http_server_listen_tls(HTTP_Server *server, string addr, uint16_t port,
string cert_file, string key_file)
{
#ifdef TLS_ENABLED
int ret = tls_server_init(&server->tcp.tls, cert_file, key_file);
if (ret < 0)
return -1;
ret = tcp_listen_tls(&server->tcp, addr, port, true, 128);
if (ret < 0)
return -1;
return 0;
#else
(void)server; (void)addr; (void)port; (void)cert_file; (void)key_file;
return -1;
#endif
}
int http_server_add_cert(HTTP_Server *server, string cert_file, string key_file)
{
int ret = tcp_add_cert(&server->tcp, cert_file, key_file);
if (ret < 0)
return -1;
return 0;
}
static void http_conn_init(HTTP_Conn *conn, TCP_Handle handle)
{
assert(conn->state == HTTP_CONN_STATE_FREE);
conn->state = HTTP_CONN_STATE_IDLE;
conn->ready = false;
conn->num_served = 0;
conn->request_len = 0;
conn->keep_alive = true;
conn->handle = handle;
}
static void http_conn_free(HTTP_Conn *conn)
{
assert(conn->state != HTTP_CONN_STATE_FREE);
conn->gen++;
if (conn->gen == 0)
conn->gen = 1;
tcp_close(conn->handle);
conn->state = HTTP_CONN_STATE_FREE;
}
void http_server_process_events(HTTP_Server *server,
void **ptrs, struct pollfd *arr, int num)
{
tcp_process_events(&server->tcp, ptrs, arr, num);
TCP_Event event;
while (tcp_next_event(&server->tcp, &event)) {
if (event.flags & TCP_EVENT_NEW) {
// New connection. Find an HTTP_Conn struct.
int i = 0;
while (i < server->max_conns && server->conns[i].state != HTTP_CONN_STATE_FREE)
i++;
if (i == server->max_conns) {
tcp_close(event.handle);
continue;
}
HTTP_Conn *conn = &server->conns[i];
http_conn_init(conn, event.handle);
tcp_set_user_ptr(event.handle, conn);
server->num_conns++;
}
HTTP_Conn *conn = tcp_get_user_ptr(event.handle);
assert(conn);
bool defer_close = false;
if (event.flags & TCP_EVENT_DATA) {
// Only idle connections can buffer bytes
assert(conn->state == HTTP_CONN_STATE_IDLE);
ByteView src = tcp_read_buf(event.handle);
int ret = chttp_parse_request((char*) src.ptr, src.len, &conn->request);
if (ret < 0) {
tcp_read_ack(event.handle, 0);
defer_close = true;
} else if (ret == 0) {
tcp_read_ack(event.handle, 0);
} else {
// Request was buffered
// Decide whether the connection should be closed
// after the next response
conn->keep_alive = true;
if (conn->num_served+1 >= 100)
conn->keep_alive = false;
conn->response_offset = tcp_write_off(event.handle);
conn->request_len = ret;
conn->state = HTTP_CONN_STATE_STATUS;
conn->ready = true;
}
}
if (event.flags & TCP_EVENT_HUP) {
defer_close = true;
}
if (defer_close) {
http_conn_free(conn);
server->num_conns--;
}
}
}
int http_server_register_events(HTTP_Server *server,
void **ptrs, struct pollfd *arr, int cap)
{
return tcp_register_events(&server->tcp, ptrs, arr, cap);
}
bool http_server_next_request(HTTP_Server *server,
HTTP_Request **request, HTTP_ResponseBuilder *builder)
{
for (int i = 0; i < server->max_conns; i++) {
if (server->conns[i].state == HTTP_CONN_STATE_FREE)
continue;
if (server->conns[i].ready) {
*request = &server->conns[i].request;
*builder = (HTTP_ResponseBuilder) {
.server = server,
.idx = i,
.gen = server->conns[i].gen,
};
return true;
}
}
return false;
}
static HTTP_Conn*
builder_to_conn(HTTP_ResponseBuilder builder)
{
if (builder.server == NULL)
return NULL;
HTTP_Server *server = builder.server;
if (builder.idx < 0 || builder.idx >= server->max_conns)
return NULL;
HTTP_Conn *conn = &server->conns[builder.idx];
if (conn->state == HTTP_CONN_STATE_FREE || conn->gen != builder.gen)
return NULL;
return conn;
}
void http_response_builder_status(HTTP_ResponseBuilder builder, int status)
{
HTTP_Conn *conn = builder_to_conn(builder);
if (conn == NULL)
return;
if (conn->state == HTTP_CONN_STATE_IDLE)
return;
if (conn->state != HTTP_CONN_STATE_STATUS) {
tcp_clear_from_offset(conn->handle, conn->response_offset);
conn->state = HTTP_CONN_STATE_STATUS;
}
assert(status > 99);
assert(status < 1000);
char tmp[3] = {
'0' + (status % 1000) / 100,
'0' + (status % 100) / 10,
'0' + (status % 10) / 1,
};
tcp_write(conn->handle, S("HTTP/1.1 "));
tcp_write(conn->handle, (string) { tmp, 3 });
tcp_write(conn->handle, S(" "));
tcp_write(conn->handle, reason_phrase(status));
tcp_write(conn->handle, S("\r\n"));
conn->state = HTTP_CONN_STATE_HEADER;
}
void http_response_builder_header(HTTP_ResponseBuilder builder, string header)
{
HTTP_Conn *conn = builder_to_conn(builder);
if (conn == NULL)
return;
if (conn->state != HTTP_CONN_STATE_HEADER)
return;
tcp_write(conn->handle, header);
tcp_write(conn->handle, S("\r\n"));
}
static void append_special_headers(HTTP_Conn *conn)
{
if (conn->keep_alive) {
tcp_write(conn->handle, S("Connection: Keep-Alive\r\n"));
} else {
tcp_write(conn->handle, S("Connection: Close\r\n"));
}
tcp_write(conn->handle, S("Content-Length: "));
conn->content_length_header_offset = tcp_write_off(conn->handle);
tcp_write(conn->handle, S(" "));
tcp_write(conn->handle, S("\r\n"));
tcp_write(conn->handle, S("\r\n"));
conn->content_offset = tcp_write_off(conn->handle);
}
void http_response_builder_content(HTTP_ResponseBuilder builder, string content)
{
HTTP_Conn *conn = builder_to_conn(builder);
if (conn == NULL)
return;
if (conn->state == HTTP_CONN_STATE_STATUS)
return;
if (conn->state != HTTP_CONN_STATE_CONTENT) {
append_special_headers(conn);
conn->state = HTTP_CONN_STATE_CONTENT;
}
tcp_write(conn->handle, content);
}
void http_response_builder_submit(HTTP_ResponseBuilder builder)
{
HTTP_Conn *conn = builder_to_conn(builder);
if (conn == NULL)
return;
if (conn->state == HTTP_CONN_STATE_STATUS)
return;
if (conn->state != HTTP_CONN_STATE_CONTENT) {
append_special_headers(conn);
conn->state = HTTP_CONN_STATE_CONTENT;
}
TCP_Offset current_offset = tcp_write_off(conn->handle);
int content_length = current_offset - conn->content_offset;
char buf[11];
int ret = snprintf(buf, sizeof(buf), "%d", content_length);
assert(ret > 0);
assert(ret < (int) sizeof(buf));
tcp_patch(conn->handle, conn->content_length_header_offset, buf, ret);
conn->num_served++;
tcp_read_ack(conn->handle, conn->request_len);
conn->request_len = 0;
if (conn->keep_alive) {
tcp_mark_ready(conn->handle);
conn->ready = false;
conn->state = HTTP_CONN_STATE_IDLE;
} else {
http_conn_free(conn);
builder.server->num_conns--;
}
}
+68
View File
@@ -0,0 +1,68 @@
#ifndef HTTP_SERVER_INCLUDED
#define HTTP_SERVER_INCLUDED
#include "tcp.h"
#include "http_parse.h"
typedef CHTTP_Request HTTP_Request;
typedef enum {
HTTP_CONN_STATE_FREE,
HTTP_CONN_STATE_IDLE,
HTTP_CONN_STATE_STATUS,
HTTP_CONN_STATE_HEADER,
HTTP_CONN_STATE_CONTENT,
} HTTP_ConnState;
typedef struct {
HTTP_ConnState state;
uint16_t gen;
bool ready;
int num_served;
int request_len;
bool keep_alive;
TCP_Handle handle;
HTTP_Request request;
TCP_Offset content_offset;
TCP_Offset content_length_header_offset;
TCP_Offset response_offset;
} HTTP_Conn;
typedef struct {
TCP tcp;
int num_conns;
int max_conns;
HTTP_Conn *conns;
} HTTP_Server;
int http_server_init(HTTP_Server *server, int max_conns);
void http_server_free(HTTP_Server *server);
int http_server_listen_tcp(HTTP_Server *server, string addr, uint16_t port);
int http_server_listen_tls(HTTP_Server *server, string addr, uint16_t port,
string cert_file, string key_file);
int http_server_add_cert(HTTP_Server *server, string cert_file, string key_file);
void http_server_process_events(HTTP_Server *server,
void **ptrs, struct pollfd *arr, int cap);
int http_server_register_events(HTTP_Server *server,
void **ptrs, struct pollfd *arr, int cap);
typedef struct {
HTTP_Server *server;
int idx;
int gen;
} HTTP_ResponseBuilder;
bool http_server_next_request(HTTP_Server *server,
HTTP_Request **request, HTTP_ResponseBuilder *builder);
void http_response_builder_status(HTTP_ResponseBuilder builder, int status);
void http_response_builder_header(HTTP_ResponseBuilder builder, string header);
void http_response_builder_content(HTTP_ResponseBuilder builder, string content);
void http_response_builder_submit(HTTP_ResponseBuilder builder);
#endif // HTTP_SERVER_INCLUDED
+289
View File
@@ -0,0 +1,289 @@
#include <quakey.h>
#include <assert.h>
#include "message.h"
#define MESSAGE_SYSTEM_VERSION 1
int message_system_init(MessageSystem *msys,
Address *addrs, int num_addrs)
{
if (num_addrs > MESSAGE_SYSTEM_NODE_LIMIT)
return -1;
for (int i = 0; i < num_addrs; i++)
msys->addrs[i] = addrs[i];
msys->num_addrs = num_addrs;
int max_conns = 2 * num_addrs + 1;
msys->conns = malloc(max_conns * sizeof(ConnMetadata));
if (msys->conns == NULL)
return -1;
msys->max_conns = max_conns;
for (int i = 0; i < max_conns; i++) {
msys->conns[i].used = false;
msys->conns[i].gen = 0;
}
int ret = tcp_init(&msys->tcp, max_conns);
if (ret < 0) {
free(msys->conns);
return -1;
}
return 0;
}
int message_system_free(MessageSystem *msys)
{
tcp_free(&msys->tcp);
free(msys->conns);
return 0;
}
int message_system_listen_tcp(MessageSystem *msys, Address addr)
{
int ret = tcp_listen_tcp(&msys->tcp, S(""), addr.port, true, 128);
if (ret < 0)
return -1;
return 0;
}
int message_system_listen_tls(MessageSystem *msys, Address addr)
{
int ret = tcp_listen_tls(&msys->tcp, S(""), addr.port, true, 128);
if (ret < 0)
return -1;
return 0;
}
void message_system_process_events(MessageSystem *msys,
void **ptrs, struct pollfd *arr, int num)
{
tcp_process_events(&msys->tcp, ptrs, arr, num);
}
int message_system_register_events(MessageSystem *msys,
void **ptrs, struct pollfd *arr, int cap)
{
return tcp_register_events(&msys->tcp, ptrs, arr, cap);
}
void *get_next_message(MessageSystem *msys)
{
TCP_Event event;
while (tcp_next_event(&msys->tcp, &event)) {
if (event.flags & TCP_EVENT_NEW) {
// Skip if already set up by ensure_conn (outbound connection)
if (tcp_get_user_ptr(event.handle) == NULL) {
int i = 0;
while (i < msys->max_conns && msys->conns[i].used)
i++;
if (i == msys->max_conns) {
tcp_close(event.handle);
continue;
}
ConnMetadata *meta = &msys->conns[i];
meta->used = true;
meta->num_senders = 0;
meta->message = NULL;
tcp_set_user_ptr(event.handle, meta);
}
}
ConnMetadata *meta = tcp_get_user_ptr(event.handle);
assert(meta);
if (event.flags & TCP_EVENT_HUP) {
meta->used = false;
tcp_close(event.handle);
continue;
}
if (!(event.flags & TCP_EVENT_DATA))
continue;
if (meta->message)
continue; // Already processing message
ByteView buf = tcp_read_buf(event.handle);
Message message;
if (buf.len < sizeof(message)) {
tcp_read_ack(event.handle, 0);
continue;
}
memcpy(&message, buf.ptr, sizeof(message));
// TODO: endianess?
if (message.version != MESSAGE_SYSTEM_VERSION) {
assert(0); // TODO
}
if (message.length > (uint64_t)buf.len) {
tcp_read_ack(event.handle, 0);
continue; // Still buffering
}
// Associate this sender with the TCP connection
if (message.sender < msys->num_addrs) {
bool found = false;
for (int i = 0; i < meta->num_senders; i++) {
if (meta->senders[i] == message.sender) {
found = true;
break;
}
}
if (!found) {
meta->senders[meta->num_senders++] = message.sender;
}
}
meta->message = buf.ptr;
return buf.ptr;
}
return NULL;
}
static TCP_Handle find_conn_by_message(MessageSystem *msys, void *message)
{
for (int i = 0; i < msys->max_conns; i++) {
ConnMetadata *meta = &msys->conns[i];
if (meta->message == message)
return (TCP_Handle) { &msys->tcp, meta->gen, i };
}
return (TCP_Handle) {0};
}
static TCP_Handle find_conn_by_target(MessageSystem *msys, int target)
{
for (int i = 0; i < msys->max_conns; i++) {
ConnMetadata *meta = &msys->conns[i];
for (int j = 0; j < meta->num_senders; j++) {
if (meta->senders[j] == target) {
return (TCP_Handle) { &msys->tcp, meta->gen, i };
}
}
}
return (TCP_Handle) {0};
}
static TCP_Handle ensure_conn(MessageSystem *msys, int target)
{
TCP_Handle handle = find_conn_by_target(msys, target);
if (handle.tcp)
return handle;
if (target < 0 || target >= msys->num_addrs)
return (TCP_Handle) {0};
int ret = tcp_connect(&msys->tcp, false, &msys->addrs[target], 1);
if (ret < 0)
return (TCP_Handle) {0};
// Find the newly created connection slot and pre-associate with target
for (int i = 0; i < msys->max_conns; i++) {
TCP_Conn *conn = &msys->tcp.conns[i];
if (conn->state == TCP_CONN_STATE_FREE)
continue;
if (conn->user_ptr != NULL)
continue;
if (conn->state != TCP_CONN_STATE_CONNECTING
&& conn->state != TCP_CONN_STATE_ESTABLISHED)
continue;
ConnMetadata *meta = &msys->conns[i];
meta->used = true;
meta->gen = conn->gen;
meta->num_senders = 1;
meta->senders[0] = target;
meta->message = NULL;
TCP_Handle h = { &msys->tcp, conn->gen, i };
tcp_set_user_ptr(h, meta);
return h;
}
return (TCP_Handle) {0};
}
void consume_message(MessageSystem *msys, void *ptr)
{
int i = 0;
while (i < msys->max_conns && msys->conns[i].message != ptr)
i++;
if (i == msys->max_conns)
return; // Not found
Message message;
memcpy(&message, ptr, sizeof(message));
TCP_Handle handle = { &msys->tcp, msys->conns[i].gen, i };
tcp_read_ack(handle, message.length);
tcp_mark_ready(handle);
msys->conns[i].message = NULL;
}
void send_message(MessageSystem *msys, int target, Message *message)
{
TCP_Handle handle = ensure_conn(msys, target);
if (!handle.tcp) return;
tcp_write(handle, (string) { (char*) message, message->length });
}
void send_message_ex(MessageSystem *msys, int target,
Message *header, void *extra, int extra_len)
{
TCP_Handle handle = ensure_conn(msys, target);
if (!handle.tcp) return;
int header_size = header->length - extra_len;
tcp_write(handle, (string) { (char*) header, header_size });
tcp_write(handle, (string) { (char*) extra, extra_len });
}
void reply_to_message(MessageSystem *msys, void *incoming_message,
Message *outgoing_message)
{
TCP_Handle handle = find_conn_by_message(msys, incoming_message);
if (!handle.tcp) return;
tcp_write(handle, (string) { (char*) outgoing_message, outgoing_message->length });
}
void reply_to_message_ex(MessageSystem *msys, void *incoming_message,
Message *outgoing_message, void *extra, int extra_len)
{
TCP_Handle handle = find_conn_by_message(msys, incoming_message);
if (!handle.tcp) return;
int header_size = outgoing_message->length - extra_len;
tcp_write(handle, (string) { (char*) outgoing_message, header_size });
tcp_write(handle, (string) { (char*) extra, extra_len });
}
void broadcast_message(MessageSystem *msys, int self_idx, Message *message)
{
for (int i = 0; i < msys->num_addrs; i++) {
if (i != self_idx)
send_message(msys, i, message);
}
}
void broadcast_message_ex(MessageSystem *msys, int self_idx,
Message *header, void *extra, int extra_len)
{
for (int i = 0; i < msys->num_addrs; i++) {
if (i != self_idx)
send_message_ex(msys, i, header, extra, extra_len);
}
}
+68
View File
@@ -0,0 +1,68 @@
#ifndef MESSAGE_INCLUDED
#define MESSAGE_INCLUDED
#include "tcp.h"
#define MESSAGE_SYSTEM_NODE_LIMIT 8
typedef struct {
bool used;
uint16_t gen;
int senders[MESSAGE_SYSTEM_NODE_LIMIT];
int num_senders;
void* message;
} ConnMetadata;
typedef struct {
TCP tcp;
Address addrs[MESSAGE_SYSTEM_NODE_LIMIT];
int num_addrs;
int max_conns;
ConnMetadata *conns;
} MessageSystem;
typedef struct {
uint16_t version;
uint16_t type;
uint16_t sender;
uint64_t length;
} Message;
int message_system_init(MessageSystem *msys,
Address *addrs, int num_addrs);
int message_system_free(MessageSystem *msys);
int message_system_listen_tcp(MessageSystem *msys, Address addr);
int message_system_listen_tls(MessageSystem *msys, Address addr);
void message_system_process_events(MessageSystem *msys,
void **ptrs, struct pollfd *arr, int num);
int message_system_register_events(MessageSystem *msys,
void **ptrs, struct pollfd *arr, int cap);
void *get_next_message(MessageSystem *msys);
void consume_message(MessageSystem *msys, void *ptr);
void send_message(MessageSystem *msys, int target, Message *message);
void send_message_ex(MessageSystem *msys, int target,
Message *header, void *extra, int extra_len);
void reply_to_message(MessageSystem *msys, void *incoming_message,
Message *outgoing_message);
void reply_to_message_ex(MessageSystem *msys, void *incoming_message,
Message *outgoing_message, void *extra, int extra_len);
void broadcast_message(MessageSystem *msys, int self_idx, Message *message);
void broadcast_message_ex(MessageSystem *msys, int self_idx,
Message *header, void *extra, int extra_len);
#endif // MESSAGE_INCLUDED
+917
View File
@@ -0,0 +1,917 @@
#if defined(MAIN_SIMULATION) || defined(MAIN_TEST)
#define QUAKEY_ENABLE_MOCKS
#endif
#include <quakey.h>
#include <assert.h>
#include "tcp.h"
#define MIN_RECV 4096
#ifdef _WIN32
typedef SOCKET NATIVE_SOCKET;
#else
typedef int NATIVE_SOCKET;
#endif
static void tcp_conn_free(TCP_Conn *conn);
static bool tcp_conn_free_maybe(TCP_Conn *conn);
static int set_socket_blocking(NATIVE_SOCKET sock, bool value)
{
#ifdef _WIN32
u_long mode = !value;
if (ioctlsocket(sock, FIONBIO, &mode) == SOCKET_ERROR)
return -1;
return 0;
#else
int flags = fcntl(sock, F_GETFL, 0);
if (flags < 0)
return -1;
if (value)
flags &= ~O_NONBLOCK;
else
flags |= O_NONBLOCK;
if (fcntl(sock, F_SETFL, flags) < 0)
return -1;
return 0;
#endif
}
static int create_listen_socket(string addr, uint16_t port,
bool reuse_addr, int backlog)
{
#ifdef _WIN32
// TODO: Only do this if socket creation fails due to
// winsock not being initialized, then try again
// with the socket
WSADATA wsa;
WSAStartup(MAKEWORD(2, 2), &wsa);
#endif
int fd = socket(AF_INET, SOCK_STREAM, 0);
if (fd == -1)
return -1;
if (set_socket_blocking(fd, false) < 0) {
close(fd);
return -1;
}
if (reuse_addr) {
int one = 1;
setsockopt(fd, SOL_SOCKET, SO_REUSEADDR, (void*) &one, sizeof(one));
}
struct in_addr addr_buf;
if (addr.len == 0)
addr_buf.s_addr = htonl(INADDR_ANY);
else {
char copy[100];
if (addr.len >= (int) sizeof(copy)) {
close(fd);
return -1;
}
memcpy(copy, addr.ptr, addr.len);
copy[addr.len] = '\0';
if (inet_pton(AF_INET, copy, &addr_buf) < 0) {
close(fd);
return -1;
}
}
struct sockaddr_in bind_buf;
bind_buf.sin_family = AF_INET;
bind_buf.sin_addr = addr_buf;
bind_buf.sin_port = htons(port);
if (bind(fd, (struct sockaddr*) &bind_buf, sizeof(bind_buf)) < 0) {
close(fd);
return -1;
}
if (listen(fd, backlog) < 0) {
close(fd);
return -1;
}
return fd;
}
int tcp_init(TCP *tcp, int max_conns)
{
TCP_Conn *conns = malloc(max_conns * sizeof(TCP_Conn));
if (conns == NULL)
return -1;
tcp->tls_listen_fd = -1;
tcp->tcp_listen_fd = -1;
tcp->num_conns = 0;
tcp->max_conns = max_conns;
tcp->conns = conns;
for (int i = 0; i < tcp->max_conns; i++) {
tcp->conns[i].state = TCP_CONN_STATE_FREE;
tcp->conns[i].gen = 0;
}
return 0;
}
void tcp_free(TCP *tcp)
{
for (int i = 0; i < tcp->max_conns; i++) {
if (tcp->conns[i].state != TCP_CONN_STATE_FREE)
tcp_conn_free(&tcp->conns[i]);
}
free(tcp->conns);
if (tcp->tcp_listen_fd != -1)
close(tcp->tcp_listen_fd);
#ifdef TLS_ENABLED
if (tcp->tls_listen_fd != -1) {
close(tcp->tls_listen_fd);
tls_server_free(&tcp->tls);
}
#endif
}
int tcp_listen_tcp(TCP *tcp, string addr, uint16_t port, bool reuse_addr, int backlog)
{
if (tcp->tcp_listen_fd != -1)
return -1;
int fd = create_listen_socket(addr, port, reuse_addr, backlog);
if (fd == -1)
return -1;
tcp->tcp_listen_fd = fd;
return 0;
}
int tcp_listen_tls(TCP *tcp, string addr, uint16_t port, bool reuse_addr, int backlog)
{
#ifdef TLS_ENABLED
if (tcp->tls_listen_fd != -1)
return -1;
int fd = create_listen_socket(addr, port, reuse_addr, backlog);
if (fd == -1)
return -1;
tcp->tls_listen_fd = fd;
return 0;
#else
(void)tcp; (void)addr; (void)port; (void)reuse_addr; (void)backlog;
return -1;
#endif
}
int tcp_add_cert(TCP *tcp, string cert_file, string key_file)
{
#ifdef TLS_ENABLED
return tls_server_add_cert(&tcp->tls, S(""), cert_file, key_file);
#else
(void)tcp; (void)cert_file; (void)key_file;
return -1;
#endif
}
static void tcp_conn_init(TCP *tcp, TCP_Conn *conn, bool secure, TCP_ConnState state, int fd)
{
conn->state = state;
conn->flags = 0;
conn->events = 0;
conn->handled = false;
conn->closing = false;
conn->fd = fd;
conn->num_addrs = 0;
conn->addr_idx = 0;
conn->user_ptr = NULL;
byte_queue_init(&conn->input, 1<<20);
byte_queue_init(&conn->output, 1<<20);
#ifdef TLS_ENABLED
if (secure) {
conn->flags |= TCP_CONN_FLAG_SECURE;
tls_conn_init(&conn->tls, &tcp->tls);
}
#else
(void)tcp;
(void)secure;
#endif
}
static void tcp_conn_free(TCP_Conn *conn)
{
if (conn->fd >= 0) {
close(conn->fd);
conn->fd = -1;
}
byte_queue_free(&conn->input);
byte_queue_free(&conn->output);
#ifdef TLS_ENABLED
if (conn->flags & TCP_CONN_FLAG_SECURE) {
tls_conn_free(&conn->tls);
}
#endif
conn->state = TCP_CONN_STATE_FREE;
}
static void tcp_conn_set_addrs(TCP_Conn *conn,
Address *addrs, int num_addrs)
{
assert(num_addrs <= TCP_CONNECT_ADDR_LIMIT);
for (int i = 0; i < num_addrs; i++)
conn->addrs[i] = addrs[i];
conn->num_addrs = num_addrs;
}
static ByteView tcp_conn_write_buf(TCP_Conn *conn)
{
#ifdef TLS_ENABLED
if (conn->flags & TCP_CONN_FLAG_SECURE) {
int cap;
char *ptr = tls_conn_net_write_buf(&conn->tls, &cap);
if (ptr == NULL)
return (ByteView) {0};
return (ByteView) { (uint8_t*) ptr, cap };
}
#endif
byte_queue_write_setmincap(&conn->input, MIN_RECV);
return byte_queue_write_buf(&conn->input);
}
static int tcp_conn_write_ack(TCP_Conn *conn, int num)
{
#ifdef TLS_ENABLED
if (conn->flags & TCP_CONN_FLAG_SECURE) {
int ret = 0;
tls_conn_net_write_ack(&conn->tls, num);
for (bool done = false; !done; ) {
byte_queue_write_setmincap(&conn->input, MIN_RECV);
ByteView buf = byte_queue_write_buf(&conn->input);
int n = tls_conn_app_read(&conn->tls, (char*) buf.ptr, buf.len);
if (n <= 0) {
if (n < 0) {
ret = -1;
n = 0;
}
done = true;
}
byte_queue_write_ack(&conn->input, n);
}
return ret;
}
#endif
byte_queue_write_ack(&conn->input, num);
return 0;
}
#ifdef TLS_ENABLED
// Encrypt plaintext from the output queue through SSL_write into the BIO.
static void tcp_conn_tls_encrypt_output(TCP_Conn *conn)
{
while (!byte_queue_empty(&conn->output)) {
ByteView src = byte_queue_read_buf(&conn->output);
if (!src.ptr || src.len == 0) {
byte_queue_read_ack(&conn->output, 0);
break;
}
int n = tls_conn_app_write(&conn->tls, (char*) src.ptr, src.len);
if (n <= 0) {
byte_queue_read_ack(&conn->output, 0);
break;
}
byte_queue_read_ack(&conn->output, n);
}
}
#endif
static ByteView tcp_conn_read_buf(TCP_Conn *conn)
{
#ifdef TLS_ENABLED
if (conn->flags & TCP_CONN_FLAG_SECURE) {
tcp_conn_tls_encrypt_output(conn);
int n;
char *ptr = tls_conn_net_read_buf(&conn->tls, &n);
if (ptr == NULL)
return (ByteView) {0};
return (ByteView) { (uint8_t*) ptr, n };
}
#endif
return byte_queue_read_buf(&conn->output);
}
static void tcp_conn_read_ack(TCP_Conn *conn, int num)
{
#ifdef TLS_ENABLED
if (conn->flags & TCP_CONN_FLAG_SECURE) {
tls_conn_net_read_ack(&conn->tls, num);
return;
}
#endif
byte_queue_read_ack(&conn->output, num);
}
static bool tcp_conn_needs_flushing(TCP_Conn *conn)
{
#ifdef TLS_ENABLED
if (conn->flags & TCP_CONN_FLAG_SECURE) {
return !byte_queue_empty(&conn->output)
|| tls_conn_needs_flushing(&conn->tls);
}
#endif
return !byte_queue_empty(&conn->output);
}
static bool tcp_conn_is_buffering(TCP_Conn *conn)
{
if (conn->closing)
return false;
if (conn->state == TCP_CONN_STATE_HANDSHAKE ||
conn->state == TCP_CONN_STATE_ACCEPTING)
return true;
return !byte_queue_reading(&conn->input);
}
static bool tcp_conn_free_maybe(TCP_Conn *conn)
{
if (!conn->handled && conn->fd < 0) {
tcp_conn_free(conn);
return true;
} else {
return false;
}
}
static void tcp_conn_invalidate_handles(TCP_Conn *conn)
{
conn->gen++;
if (conn->gen == 0)
conn->gen = 1;
}
static int build_sockaddr(Address *addr, struct sockaddr_in *out)
{
memset(out, 0, sizeof(*out));
if (!addr->is_ipv4)
return -1; // Only IPv4 supported for now
out->sin_family = AF_INET;
out->sin_port = htons(addr->port);
memcpy(&out->sin_addr, &addr->ipv4, sizeof(addr->ipv4));
return 0;
}
int tcp_connect(TCP *tcp, bool secure, Address *addrs, int num_addrs)
{
if (tcp->num_conns == tcp->max_conns)
return -1;
int i = 0;
while (i < tcp->max_conns && tcp->conns[i].state != TCP_CONN_STATE_FREE)
i++;
assert(i < tcp->max_conns);
int fd = socket(AF_INET, SOCK_STREAM, 0);
if (fd < 0) {
return -1;
}
if (set_socket_blocking(fd, false) < 0) {
close(fd);
return -1;
}
struct sockaddr_in sa;
if (build_sockaddr(&addrs[0], &sa) < 0) {
close(fd);
return -1;
}
TCP_ConnState state;
int ret = connect(fd, (struct sockaddr*) &sa, sizeof(sa));
if (ret == 0) {
if (secure) {
state = TCP_CONN_STATE_HANDSHAKE;
} else {
state = TCP_CONN_STATE_ESTABLISHED;
}
} else {
assert(ret < 0);
if (errno == EINPROGRESS) {
state = TCP_CONN_STATE_CONNECTING;
} else {
close(fd);
return -1;
}
}
tcp_conn_init(tcp, &tcp->conns[i], secure, state, fd);
tcp_conn_set_addrs(&tcp->conns[i], addrs, num_addrs);
tcp->num_conns++;
return 0;
}
static int restart_connect(TCP_Conn *conn)
{
close(conn->fd);
conn->fd = -1;
conn->addr_idx++;
if (conn->addr_idx == conn->num_addrs) {
return -1;
}
int fd = socket(AF_INET, SOCK_STREAM, 0);
if (fd < 0) {
return -1;
}
if (set_socket_blocking(fd, false) < 0) {
close(fd);
return -1;
}
struct sockaddr_in sa;
if (build_sockaddr(&conn->addrs[conn->addr_idx], &sa) < 0) {
close(fd);
return -1;
}
TCP_ConnState state;
int ret = connect(fd, (struct sockaddr*) &sa, sizeof(sa));
if (ret == 0) {
if (conn->flags & TCP_CONN_FLAG_SECURE) {
state = TCP_CONN_STATE_HANDSHAKE;
} else {
state = TCP_CONN_STATE_ESTABLISHED;
}
} else {
assert(ret < 0);
if (errno == EINPROGRESS) {
state = TCP_CONN_STATE_CONNECTING;
} else {
close(fd);
return -1;
}
}
conn->fd = fd;
conn->state = state;
return 0;
}
void tcp_process_events(TCP *tcp, void **ptrs, struct pollfd *arr, int num)
{
for (int i = 0; i < num; i++) {
if (arr[i].fd == tcp->tcp_listen_fd ||
arr[i].fd == tcp->tls_listen_fd) {
assert(ptrs[i] == NULL);
if (arr[i].revents & POLLIN) {
if (tcp->num_conns == tcp->max_conns)
continue;
bool is_tls = false;
if (arr[i].fd == tcp->tls_listen_fd)
is_tls = true;
int new_fd = accept(arr[i].fd, NULL, NULL);
if (new_fd == -1)
continue;
if (set_socket_blocking(new_fd, false) < 0) {
close(new_fd);
continue;
}
// Find a free connection slot
int slot = 0;
while (slot < tcp->max_conns && tcp->conns[slot].state != TCP_CONN_STATE_FREE)
slot++;
if (slot == tcp->max_conns) {
close(new_fd);
continue;
}
TCP_ConnState state;
if (is_tls) {
state = TCP_CONN_STATE_ACCEPTING;
} else {
state = TCP_CONN_STATE_ESTABLISHED;
}
TCP_Conn *conn = &tcp->conns[slot];
tcp_conn_init(tcp, conn, is_tls, state, new_fd);
if (!is_tls)
conn->events |= TCP_EVENT_NEW;
tcp->num_conns++;
}
} else {
TCP_Conn *conn = ptrs[i];
bool defer_ready = false;
bool defer_close = false;
bool defer_connect = false;
switch (conn->state) {
case TCP_CONN_STATE_CONNECTING:
{
int err = 0;
socklen_t len = sizeof(err);
if (getsockopt(conn->fd, SOL_SOCKET, SO_ERROR, (void*) &err, &len) < 0) {
defer_connect = true;
break;
}
if (err) {
defer_connect = true;
break;
}
if (conn->flags & TCP_CONN_FLAG_SECURE) {
conn->state = TCP_CONN_STATE_HANDSHAKE;
} else {
conn->state = TCP_CONN_STATE_ESTABLISHED;
conn->events |= TCP_EVENT_NEW;
}
}
break;
case TCP_CONN_STATE_HANDSHAKE:
case TCP_CONN_STATE_ACCEPTING:
{
#ifdef TLS_ENABLED
if (arr[i].revents & POLLIN) {
int cap;
char *buf = tls_conn_net_write_buf(&conn->tls, &cap);
if (buf) {
int n = recv(conn->fd, buf, cap, 0);
if (n == 0) { defer_close = true; break; }
if (n < 0) {
if (errno != EINTR && errno != EWOULDBLOCK && errno != EAGAIN)
{ defer_close = true; break; }
n = 0;
}
tls_conn_net_write_ack(&conn->tls, n);
}
}
int ret = tls_conn_handshake(&conn->tls);
if (ret == -1) {
defer_close = true;
break;
}
if (arr[i].revents & POLLOUT) {
int num;
char *buf = tls_conn_net_read_buf(&conn->tls, &num);
if (buf) {
int n = send(conn->fd, buf, num, 0);
if (n < 0) {
if (errno != EINTR && errno != EWOULDBLOCK && errno != EAGAIN)
{ defer_close = true; break; }
n = 0;
}
tls_conn_net_read_ack(&conn->tls, n);
}
}
if (ret == 1) {
conn->state = TCP_CONN_STATE_ESTABLISHED;
conn->events |= TCP_EVENT_NEW;
// Decrypt any application data already in the BIO
for (;;) {
byte_queue_write_setmincap(&conn->input, MIN_RECV);
ByteView buf = byte_queue_write_buf(&conn->input);
if (!buf.ptr) break;
int n = tls_conn_app_read(&conn->tls, (char*) buf.ptr, buf.len);
if (n <= 0) { byte_queue_write_ack(&conn->input, 0); break; }
byte_queue_write_ack(&conn->input, n);
conn->events |= TCP_EVENT_DATA;
}
}
#else
defer_close = true;
#endif
}
break;
case TCP_CONN_STATE_ESTABLISHED:
{
if (arr[i].revents & POLLIN) {
ByteView buf = tcp_conn_write_buf(conn);
int n = recv(conn->fd, (char*) buf.ptr, buf.len, 0);
if (n == 0) {
defer_close = true;
} else {
if (n < 0) {
if (errno != EINTR && errno != EWOULDBLOCK && errno != EAGAIN)
defer_close = true;
n = 0;
}
}
int ret = tcp_conn_write_ack(conn, n);
if (ret < 0)
defer_close = true;
defer_ready = true;
}
if (arr[i].revents & POLLOUT) {
ByteView buf = tcp_conn_read_buf(conn);
int n = send(conn->fd, (char*) buf.ptr, buf.len, 0);
if (n < 0) {
if (errno != EINTR && errno != EWOULDBLOCK && errno != EAGAIN)
defer_close = true;
n = 0;
}
tcp_conn_read_ack(conn, n);
if (conn->closing && !tcp_conn_needs_flushing(conn))
defer_close = true;
}
}
break;
case TCP_CONN_STATE_SHUTDOWN:
{
// TLS shutdown — just close for now
defer_close = true;
}
break;
default:
break;
}
if (defer_connect) {
int ret = restart_connect(conn);
if (ret < 0) {
defer_close = true;
}
}
if (defer_ready) {
conn->events |= TCP_EVENT_DATA;
}
if (defer_close) {
close(conn->fd);
conn->fd = -1;
conn->events |= TCP_EVENT_HUP;
if (tcp_conn_free_maybe(conn)) {
tcp->num_conns--;
}
}
}
}
}
int tcp_register_events(TCP *tcp, void **ptrs, struct pollfd *arr, int cap)
{
if (cap < tcp->num_conns+2)
return -1;
int ret = 0;
if (tcp->tcp_listen_fd > -1) {
if (tcp->num_conns < tcp->max_conns) {
arr[ret].fd = tcp->tcp_listen_fd;
arr[ret].events = POLLIN;
arr[ret].revents = 0;
ptrs[ret] = NULL;
ret++;
}
}
if (tcp->tls_listen_fd > -1) {
if (tcp->num_conns < tcp->max_conns) {
arr[ret].fd = tcp->tls_listen_fd;
arr[ret].events = POLLIN;
arr[ret].revents = 0;
ptrs[ret] = NULL;
ret++;
}
}
for (int i=0, j=0; j < tcp->num_conns; i++) {
TCP_Conn *conn = &tcp->conns[i];
if (conn->state == TCP_CONN_STATE_FREE)
continue;
j++;
int events = 0;
if (conn->state == TCP_CONN_STATE_CONNECTING)
events |= POLLOUT;
if (tcp_conn_is_buffering(conn))
events |= POLLIN;
if (tcp_conn_needs_flushing(conn))
events |= POLLOUT;
if (events) {
arr[ret].fd = conn->fd;
arr[ret].events = events;
arr[ret].revents = 0;
ptrs[ret] = conn;
ret++;
}
}
return ret;
}
static TCP_Handle conn_to_handle(TCP *tcp, TCP_Conn *conn)
{
TCP_Handle handle = {
.tcp=tcp,
.gen=conn->gen,
.idx=conn - tcp->conns,
};
return handle;
}
static TCP_Conn *handle_to_conn(TCP_Handle handle)
{
if (handle.tcp == NULL)
return NULL;
TCP *tcp = handle.tcp;
if (handle.idx < 0 || handle.idx >= tcp->max_conns)
return NULL;
TCP_Conn *conn = &tcp->conns[handle.idx];
if (conn->state == TCP_CONN_STATE_FREE || conn->gen != handle.gen)
return NULL;
return conn;
}
static bool conn_to_event(TCP *tcp, TCP_Conn *conn, TCP_Event *event)
{
if (!conn->events)
return false;
*event = (TCP_Event) {
.flags = conn->events,
.handle = conn_to_handle(tcp, conn),
};
conn->events = 0;
return true;
}
bool tcp_next_event(TCP *tcp, TCP_Event *event)
{
for (int i = 0, j = 0; j < tcp->num_conns; i++) {
TCP_Conn *conn = &tcp->conns[i];
if (conn->state == TCP_CONN_STATE_FREE)
continue;
j++;
if (conn->flags & TCP_CONN_FLAG_CLOSED)
continue; // User isn't interested in this connection anymore
if (conn_to_event(tcp, conn, event))
return true;
}
return false;
}
ByteView tcp_read_buf(TCP_Handle handle)
{
TCP_Conn *conn = handle_to_conn(handle);
if (conn == NULL)
return (ByteView) {0};
return byte_queue_read_buf(&conn->input);
}
void tcp_read_ack(TCP_Handle handle, int num)
{
TCP_Conn *conn = handle_to_conn(handle);
if (conn == NULL)
return;
byte_queue_read_ack(&conn->input, num);
}
ByteView tcp_write_buf(TCP_Handle handle)
{
TCP_Conn *conn = handle_to_conn(handle);
if (conn == NULL)
return (ByteView) {0};
return byte_queue_write_buf(&conn->output);
}
void tcp_write_ack(TCP_Handle handle, int num)
{
TCP_Conn *conn = handle_to_conn(handle);
if (conn == NULL)
return;
byte_queue_write_ack(&conn->output, num);
}
TCP_Offset tcp_write_off(TCP_Handle handle)
{
TCP_Conn *conn = handle_to_conn(handle);
if (conn == NULL)
return 0;
return byte_queue_offset(&conn->output);
}
void tcp_write(TCP_Handle handle, string str)
{
while (str.len > 0) {
byte_queue_write_setmincap(&handle_to_conn(handle)->output, str.len);
ByteView buf = tcp_write_buf(handle);
int num = MIN(buf.len, str.len);
memcpy(buf.ptr, str.ptr, num);
tcp_write_ack(handle, num);
str.ptr += num;
str.len -= num;
}
}
void tcp_patch(TCP_Handle handle, TCP_Offset offset, void *src, int len)
{
TCP_Conn *conn = handle_to_conn(handle);
if (conn == NULL)
return;
byte_queue_patch(&conn->output, offset, src, len);
}
void tcp_clear_from_offset(TCP_Handle handle, TCP_Offset offset)
{
TCP_Conn *conn = handle_to_conn(handle);
if (conn == NULL)
return;
byte_queue_remove_from_offset(&conn->output, offset);
}
void tcp_close(TCP_Handle handle)
{
TCP *tcp = handle.tcp;
TCP_Conn *conn = handle_to_conn(handle);
if (conn == NULL)
return;
conn->flags |= TCP_CONN_FLAG_CLOSED;
conn->handled = false;
tcp_conn_invalidate_handles(conn);
if (tcp_conn_free_maybe(conn)) {
tcp->num_conns--;
}
}
void tcp_set_user_ptr(TCP_Handle handle, void *ptr)
{
TCP_Conn *conn = handle_to_conn(handle);
if (conn == NULL)
return;
conn->user_ptr = ptr;
}
void *tcp_get_user_ptr(TCP_Handle handle)
{
TCP_Conn *conn = handle_to_conn(handle);
if (conn == NULL)
return NULL;
return conn->user_ptr;
}
void tcp_mark_ready(TCP_Handle handle)
{
TCP_Conn *conn = handle_to_conn(handle);
if (conn == NULL)
return;
conn->events |= TCP_EVENT_DATA;
}
+99
View File
@@ -0,0 +1,99 @@
#ifndef TCP_INCLUDED
#define TCP_INCLUDED
#include "basic.h"
#include "byte_queue.h"
#include "tls.h"
typedef struct TCP TCP;
typedef struct {
TCP* tcp;
uint16_t gen;
int idx;
} TCP_Handle;
typedef enum {
TCP_CONN_STATE_FREE,
TCP_CONN_STATE_HANDSHAKE,
TCP_CONN_STATE_ESTABLISHED,
TCP_CONN_STATE_CONNECTING,
TCP_CONN_STATE_ACCEPTING,
TCP_CONN_STATE_SHUTDOWN,
} TCP_ConnState;
#define TCP_CONNECT_ADDR_LIMIT 8
enum {
TCP_EVENT_NEW = 1<<0,
TCP_EVENT_HUP = 1<<1,
TCP_EVENT_DATA = 1<<2,
};
enum {
TCP_CONN_FLAG_CLOSED = 1<<0,
TCP_CONN_FLAG_SECURE = 1<<1,
};
typedef struct {
TCP_ConnState state;
int flags;
int events;
uint16_t gen;
int fd;
bool handled;
bool closing;
void *user_ptr;
ByteQueue input;
ByteQueue output;
#ifdef TLS_ENABLED
TLS_Conn tls;
#endif
Address addrs[TCP_CONNECT_ADDR_LIMIT];
int num_addrs;
int addr_idx;
} TCP_Conn;
struct TCP {
int tls_listen_fd;
int tcp_listen_fd;
int num_conns;
int max_conns;
TCP_Conn *conns;
#ifdef TLS_ENABLED
TLS_Server tls;
#endif
};
typedef struct {
int flags;
TCP_Handle handle;
} TCP_Event;
typedef ByteQueueOffset TCP_Offset;
struct pollfd;
int tcp_init(TCP *tcp, int max_conns);
void tcp_free(TCP *tcp);
int tcp_listen_tcp(TCP *tcp, string addr, uint16_t port, bool reuse_addr, int backlog);
int tcp_listen_tls(TCP *tcp, string addr, uint16_t port, bool reuse_addr, int backlog);
int tcp_add_cert(TCP *tcp, string cert_file, string key_file);
int tcp_connect(TCP *tcp, bool secure, Address *addrs, int num_addrs);
void tcp_process_events(TCP *tcp, void **ptrs, struct pollfd *arr, int num);
int tcp_register_events(TCP *tcp, void **ptrs, struct pollfd *arr, int cap);
bool tcp_next_event(TCP *tcp, TCP_Event *event);
ByteView tcp_read_buf(TCP_Handle handle);
void tcp_read_ack(TCP_Handle handle, int num);
ByteView tcp_write_buf(TCP_Handle handle);
void tcp_write_ack(TCP_Handle handle, int num);
TCP_Offset tcp_write_off(TCP_Handle handle);
void tcp_write(TCP_Handle handle, string str);
void tcp_patch(TCP_Handle handle, TCP_Offset offset, void *src, int len);
void tcp_clear_from_offset(TCP_Handle handle, TCP_Offset offset);
void tcp_close(TCP_Handle handle);
void tcp_set_user_ptr(TCP_Handle handle, void *ptr);
void *tcp_get_user_ptr(TCP_Handle handle);
void tcp_mark_ready(TCP_Handle handle);
#endif // TCP_INCLUDED
+92
View File
@@ -0,0 +1,92 @@
#ifdef TLS_ENABLED
#ifndef TLS_INCLUDED
#define TLS_INCLUDED
#include <stdbool.h>
#include "basic.h"
#define TLS_CERT_LIMIT 8
#ifdef TLS_SCHANNEL
#define SECURITY_WIN32
#include <windows.h>
#include <security.h>
#include <schannel.h>
#include "byte_queue.h"
typedef struct {
char domain[128];
CredHandle cred;
} TLS_Cert;
typedef struct {
CredHandle cred;
PCCERT_CONTEXT cert_ctx;
int num_certs;
TLS_Cert certs[TLS_CERT_LIMIT];
} TLS_Server;
typedef struct {
CtxtHandle ctx;
bool ctx_valid;
bool handshake;
CredHandle *cred;
ByteQueue in_buf;
ByteQueue out_buf;
SecPkgContext_StreamSizes sizes;
char *pending;
int pending_off;
int pending_len;
} TLS_Conn;
#else // OpenSSL
typedef struct ssl_ctx_st SSL_CTX;
typedef struct ssl_st SSL;
typedef struct bio_st BIO;
typedef struct {
char domain[128];
SSL_CTX *ctx;
} TLS_Cert;
typedef struct {
SSL_CTX *ctx;
int num_certs;
TLS_Cert certs[TLS_CERT_LIMIT];
} TLS_Server;
typedef struct {
SSL *ssl;
BIO *network_bio;
bool handshake;
} TLS_Conn;
#endif // TLS_SCHANNEL
void tls_global_init(void);
void tls_global_free(void);
int tls_server_init(TLS_Server *server, string cert_file, string key_file);
void tls_server_free(TLS_Server *server);
int tls_server_add_cert(TLS_Server *server, string domain, string cert_file, string key_file);
int tls_conn_init(TLS_Conn *conn, TLS_Server *server);
void tls_conn_free(TLS_Conn *conn);
int tls_conn_handshake(TLS_Conn *conn);
char *tls_conn_net_write_buf(TLS_Conn *conn, int *cap);
void tls_conn_net_write_ack(TLS_Conn *conn, int num);
char *tls_conn_net_read_buf(TLS_Conn *conn, int *num);
void tls_conn_net_read_ack(TLS_Conn *conn, int num);
int tls_conn_app_write(TLS_Conn *conn, char *dst, int num);
int tls_conn_app_read(TLS_Conn *conn, char *src, int cap);
int tls_conn_needs_flushing(TLS_Conn *conn);
#endif // TLS_INCLUDED
#endif // TLS_ENABLED
+293
View File
@@ -0,0 +1,293 @@
#ifdef TLS_ENABLED
#ifdef TLS_OPENSSL
#if defined(MAIN_SIMULATION) || defined(MAIN_TEST)
#define QUAKEY_ENABLE_MOCKS
#endif
#include <quakey.h>
#include <assert.h>
// Avoid name collision between basic.h's SHA256 typedef
// and OpenSSL's SHA256 function
#define SHA256 openssl_SHA256
#include <openssl/ssl.h>
#include <openssl/err.h>
#undef SHA256
#include "tls.h"
void tls_global_init(void)
{
SSL_library_init();
SSL_load_error_strings();
OpenSSL_add_all_algorithms();
}
void tls_global_free(void)
{
EVP_cleanup();
}
static int servername_callback(SSL *ssl, int *ad, void *arg)
{
TLS_Server *server = arg;
// The 'ad' parameter is used to set the alert description when returning
// SSL_TLSEXT_ERR_ALERT_FATAL. Since we only return OK or NOACK, it's unused.
(void) ad;
const char *servername = SSL_get_servername(ssl, TLSEXT_NAMETYPE_host_name);
if (servername == NULL)
return SSL_TLSEXT_ERR_NOACK;
for (int i = 0; i < server->num_certs; i++) {
TLS_Cert *cert = &server->certs[i];
if (!strcmp(cert->domain, servername)) {
SSL_set_SSL_CTX(ssl, cert->ctx);
return SSL_TLSEXT_ERR_OK;
}
}
return SSL_TLSEXT_ERR_NOACK;
}
int tls_server_init(TLS_Server *server, string cert_file, string key_file)
{
SSL_CTX *ctx = SSL_CTX_new(TLS_server_method());
if (ctx == NULL)
return -1;
SSL_CTX_set_min_proto_version(ctx, TLS1_2_VERSION);
char cert_buf[1024];
if (cert_file.len >= (int) sizeof(cert_buf)) {
SSL_CTX_free(ctx);
return -1;
}
memcpy(cert_buf, cert_file.ptr, cert_file.len);
cert_buf[cert_file.len] = '\0';
char key_buf[1024];
if (key_file.len >= (int) sizeof(key_buf)) {
SSL_CTX_free(ctx);
return -1;
}
memcpy(key_buf, key_file.ptr, key_file.len);
key_buf[key_file.len] = '\0';
// Load certificate and private key
if (SSL_CTX_use_certificate_chain_file(ctx, cert_buf) != 1) {
SSL_CTX_free(ctx);
return -1;
}
if (SSL_CTX_use_PrivateKey_file(ctx, key_buf, SSL_FILETYPE_PEM) != 1) {
SSL_CTX_free(ctx);
return -1;
}
if (SSL_CTX_check_private_key(ctx) != 1) {
SSL_CTX_free(ctx);
return -1;
}
SSL_CTX_set_tlsext_servername_callback(ctx, servername_callback);
SSL_CTX_set_tlsext_servername_arg(ctx, server);
server->ctx = ctx;
server->num_certs = 0;
return 0;
}
void tls_server_free(TLS_Server *server)
{
for (int i = 0; i < server->num_certs; i++)
SSL_CTX_free(server->certs[i].ctx);
SSL_CTX_free(server->ctx);
}
// TODO: Can the domain be inferred from the cert?
int tls_server_add_cert(TLS_Server *server, string domain, string cert_file, string key_file)
{
if (server->num_certs == TLS_CERT_LIMIT)
return -1;
SSL_CTX *ctx = SSL_CTX_new(TLS_server_method());
if (!ctx)
return -1;
SSL_CTX_set_min_proto_version(ctx, TLS1_2_VERSION);
char cert_buf[1024];
if (cert_file.len >= (int) sizeof(cert_buf)) {
SSL_CTX_free(ctx);
return -1;
}
memcpy(cert_buf, cert_file.ptr, cert_file.len);
cert_buf[cert_file.len] = '\0';
char key_buf[1024];
if (key_file.len >= (int) sizeof(key_buf)) {
SSL_CTX_free(ctx);
return -1;
}
memcpy(key_buf, key_file.ptr, key_file.len);
key_buf[key_file.len] = '\0';
if (SSL_CTX_use_certificate_chain_file(ctx, cert_buf) != 1) {
SSL_CTX_free(ctx);
return -1;
}
if (SSL_CTX_use_PrivateKey_file(ctx, key_buf, SSL_FILETYPE_PEM) != 1) {
SSL_CTX_free(ctx);
return -1;
}
if (SSL_CTX_check_private_key(ctx) != 1) {
SSL_CTX_free(ctx);
return -1;
}
TLS_Cert *cert = &server->certs[server->num_certs];
if (domain.len >= (int) sizeof(cert->domain)) {
SSL_CTX_free(ctx);
return -1;
}
memcpy(cert->domain, domain.ptr, domain.len);
cert->domain[domain.len] = '\0';
cert->ctx = ctx;
server->num_certs++;
return 0;
}
int tls_conn_init(TLS_Conn *conn, TLS_Server *server)
{
SSL *ssl = SSL_new(server->ctx);
if (ssl == NULL)
return -1;
// Create a BIO pair:
// internal_bio — attached to the SSL object (OpenSSL uses it internally)
// network_bio — you read/write encrypted data from/to this
BIO *internal_bio = NULL;
BIO *network_bio = NULL;
if (!BIO_new_bio_pair(&internal_bio, 0, &network_bio, 0)) {
SSL_free(ssl);
return -1;
}
// Bind the internal side to the SSL object
SSL_set_bio(ssl, internal_bio, internal_bio);
// We're the server side
SSL_set_accept_state(ssl);
conn->ssl = ssl;
conn->network_bio = network_bio;
conn->handshake = true;
return 0;
}
void tls_conn_free(TLS_Conn *conn)
{
SSL_free(conn->ssl);
BIO_free(conn->network_bio);
}
// Write ciphertext from the connection object to the network
char *tls_conn_net_write_buf(TLS_Conn *conn, int *cap)
{
char *buf;
int ret = BIO_nwrite0(conn->network_bio, &buf);
if (ret <= 0)
return NULL;
*cap = ret;
return buf;
}
// Complete the write from the connection object
void tls_conn_net_write_ack(TLS_Conn *conn, int num)
{
char *dummy;
BIO_nwrite(conn->network_bio, &dummy, num);
}
// Read ciphertext from the network into the connection object
char *tls_conn_net_read_buf(TLS_Conn *conn, int *num)
{
char *buf;
int ret = BIO_nread0(conn->network_bio, &buf);
if (ret <= 0)
return NULL;
*num = ret;
return buf;
}
// Complete the read from the network
void tls_conn_net_read_ack(TLS_Conn *conn, int num)
{
char *dummy;
BIO_nread(conn->network_bio, &dummy, num);
}
// Write plaintext from the application to the connection object
int tls_conn_app_write(TLS_Conn *conn, char *dst, int num)
{
assert(!conn->handshake);
int n = SSL_write(conn->ssl, dst, num);
if (n > 0)
return n;
int err = SSL_get_error(conn->ssl, n);
if (err == SSL_ERROR_WANT_READ ||
err == SSL_ERROR_WANT_WRITE)
return 0;
return -1;
}
// Read plaintext from the connection object into the application
int tls_conn_app_read(TLS_Conn *conn, char *src, int cap)
{
assert(!conn->handshake);
int n = SSL_read(conn->ssl, src, cap);
if (n > 0)
return n;
int err = SSL_get_error(conn->ssl, n);
if (err == SSL_ERROR_WANT_READ ||
err == SSL_ERROR_WANT_WRITE)
return 0;
return -1;
}
int tls_conn_handshake(TLS_Conn *conn)
{
assert(conn->handshake);
int n = SSL_do_handshake(conn->ssl);
if (n == 1) {
conn->handshake = false;
return 1;
}
int err = SSL_get_error(conn->ssl, n);
if (err == SSL_ERROR_WANT_READ ||
err == SSL_ERROR_WANT_WRITE)
return 0;
return -1;
}
int tls_conn_needs_flushing(TLS_Conn *conn)
{
return BIO_ctrl_pending(conn->network_bio) > 0;
}
#endif // TLS_OPENSSL
#endif // TLS_ENABLED
+695
View File
@@ -0,0 +1,695 @@
#ifdef TLS_ENABLED
#ifdef TLS_SCHANNEL
#include <assert.h>
#include <stdio.h>
#include <string.h>
#define SECURITY_WIN32
#include <windows.h>
#include <wincrypt.h>
#include <security.h>
#include <schannel.h>
#include "tls.h"
#pragma comment(lib, "secur32.lib")
#pragma comment(lib, "crypt32.lib")
#pragma comment(lib, "ncrypt.lib")
#define TLS_BUF_LIMIT (32 * 1024)
// ============================================================
// Certificate loading
// ============================================================
// Read entire file into malloc'd buffer. Caller must free.
static char *read_file(const char *path, int *out_len)
{
FILE *f = fopen(path, "rb");
if (!f) return NULL;
fseek(f, 0, SEEK_END);
long len = ftell(f);
fseek(f, 0, SEEK_SET);
if (len <= 0 || len > 1024 * 1024) {
fclose(f);
return NULL;
}
char *buf = malloc((size_t) len + 1);
if (!buf) {
fclose(f);
return NULL;
}
size_t nread = fread(buf, 1, (size_t) len, f);
fclose(f);
buf[nread] = '\0';
*out_len = (int) nread;
return buf;
}
// Decode PEM base64 content to DER binary. Returns malloc'd buffer.
static BYTE *pem_to_der(const char *pem, int pem_len, DWORD *out_len)
{
DWORD len = 0;
if (!CryptStringToBinaryA(pem, pem_len, CRYPT_STRING_BASE64HEADER,
NULL, &len, NULL, NULL))
return NULL;
BYTE *der = malloc(len);
if (!der) return NULL;
if (!CryptStringToBinaryA(pem, pem_len, CRYPT_STRING_BASE64HEADER,
der, &len, NULL, NULL)) {
free(der);
return NULL;
}
*out_len = len;
return der;
}
// Import an in-memory PFX blob, acquire SChannel credential.
static int load_pfx_blob(const BYTE *pfx_data, DWORD pfx_len,
CredHandle *cred_out, PCCERT_CONTEXT *cert_ctx_out)
{
CRYPT_DATA_BLOB blob;
blob.pbData = (BYTE *) pfx_data;
blob.cbData = pfx_len;
HCERTSTORE store = PFXImportCertStore(&blob, L"", CRYPT_EXPORTABLE);
if (!store) return -1;
// Find first cert with a private key
PCCERT_CONTEXT cert_ctx = NULL;
while ((cert_ctx = CertEnumCertificatesInStore(store, cert_ctx)) != NULL) {
HCRYPTPROV_OR_NCRYPT_KEY_HANDLE hkey = 0;
DWORD key_spec = 0;
BOOL caller_free = FALSE;
if (CryptAcquireCertificatePrivateKey(cert_ctx,
CRYPT_ACQUIRE_ALLOW_NCRYPT_KEY_FLAG | CRYPT_ACQUIRE_SILENT_FLAG,
NULL, &hkey, &key_spec, &caller_free)) {
if (caller_free && hkey) {
if (key_spec == CERT_NCRYPT_KEY_SPEC)
NCryptFreeObject(hkey);
else
CryptReleaseContext(hkey, 0);
}
break;
}
}
if (!cert_ctx) {
CertCloseStore(store, 0);
return -1;
}
PCCERT_CONTEXT dup_ctx = CertDuplicateCertificateContext(cert_ctx);
SCHANNEL_CRED sc_cred = {0};
sc_cred.dwVersion = SCHANNEL_CRED_VERSION;
sc_cred.cCreds = 1;
sc_cred.paCred = &dup_ctx;
sc_cred.grbitEnabledProtocols = SP_PROT_TLS1_2_SERVER;
TimeStamp expiry;
SECURITY_STATUS ss = AcquireCredentialsHandleA(
NULL, UNISP_NAME_A, SECPKG_CRED_INBOUND,
NULL, &sc_cred, NULL, NULL,
cred_out, &expiry);
if (ss != SEC_E_OK) {
CertFreeCertificateContext(dup_ctx);
CertCloseStore(store, 0);
return -1;
}
*cert_ctx_out = dup_ctx;
// Keep store open — SChannel needs access to the cert+key
return 0;
}
// Load PEM cert+key using pure CryptoAPI (no external tools).
// Decodes PEM, imports key into a temporary CAPI keyset, exports as
// in-memory PFX, then imports via PFXImportCertStore.
static int load_pem_credential(string cert_file, string key_file,
CredHandle *cred_out, PCCERT_CONTEXT *cert_ctx_out)
{
int result = -1;
char cert_z[1024], key_z[1024], container[64];
char *cert_pem = NULL, *key_pem = NULL;
BYTE *cert_der = NULL, *key_der = NULL;
BYTE *rsa_blob = NULL, *pkcs8_buf = NULL;
PCCERT_CONTEXT cert_ctx = NULL;
HCRYPTPROV hprov = 0;
HCRYPTKEY hkey = 0;
HCERTSTORE mem_store = NULL;
CRYPT_DATA_BLOB pfx = {0};
int cert_pem_len, key_pem_len;
DWORD cert_der_len, key_der_len, rsa_blob_len, pkcs8_len;
snprintf(container, sizeof(container), "tls_tmp_%lu",
(unsigned long) GetCurrentProcessId());
// Null-terminate file paths
if (cert_file.len >= (int) sizeof(cert_z)) goto done;
memcpy(cert_z, cert_file.ptr, cert_file.len);
cert_z[cert_file.len] = '\0';
if (key_file.len >= (int) sizeof(key_z)) goto done;
memcpy(key_z, key_file.ptr, key_file.len);
key_z[key_file.len] = '\0';
// --- Certificate: PEM -> DER -> CERT_CONTEXT ---
cert_pem = read_file(cert_z, &cert_pem_len);
if (!cert_pem) goto done;
cert_der = pem_to_der(cert_pem, cert_pem_len, &cert_der_len);
if (!cert_der) goto done;
cert_ctx = CertCreateCertificateContext(
X509_ASN_ENCODING | PKCS_7_ASN_ENCODING, cert_der, cert_der_len);
if (!cert_ctx) goto done;
// --- Private key: PEM -> DER -> CAPI RSA blob ---
key_pem = read_file(key_z, &key_pem_len);
if (!key_pem) goto done;
key_der = pem_to_der(key_pem, key_pem_len, &key_der_len);
if (!key_der) goto done;
if (strstr(key_pem, "-----BEGIN PRIVATE KEY-----")) {
// PKCS#8: unwrap to get inner RSA key
if (!CryptDecodeObjectEx(X509_ASN_ENCODING, PKCS_PRIVATE_KEY_INFO,
key_der, key_der_len, CRYPT_DECODE_ALLOC_FLAG,
NULL, &pkcs8_buf, &pkcs8_len))
goto done;
CRYPT_PRIVATE_KEY_INFO *pki = (CRYPT_PRIVATE_KEY_INFO *) pkcs8_buf;
if (!CryptDecodeObjectEx(X509_ASN_ENCODING, PKCS_RSA_PRIVATE_KEY,
pki->PrivateKey.pbData, pki->PrivateKey.cbData,
CRYPT_DECODE_ALLOC_FLAG, NULL, &rsa_blob, &rsa_blob_len))
goto done;
} else {
// Traditional RSA (BEGIN RSA PRIVATE KEY)
if (!CryptDecodeObjectEx(X509_ASN_ENCODING, PKCS_RSA_PRIVATE_KEY,
key_der, key_der_len, CRYPT_DECODE_ALLOC_FLAG,
NULL, &rsa_blob, &rsa_blob_len))
goto done;
}
// --- Import key into temporary CAPI keyset ---
CryptAcquireContextA(&hprov, container, MS_ENH_RSA_AES_PROV_A,
PROV_RSA_AES, CRYPT_DELETEKEYSET);
hprov = 0;
if (!CryptAcquireContextA(&hprov, container, MS_ENH_RSA_AES_PROV_A,
PROV_RSA_AES, CRYPT_NEWKEYSET))
goto done;
if (!CryptImportKey(hprov, rsa_blob, rsa_blob_len, 0, CRYPT_EXPORTABLE, &hkey))
goto done;
// --- Bind key to cert, export PFX in memory ---
{
wchar_t containerW[64];
MultiByteToWideChar(CP_ACP, 0, container, -1, containerW, 64);
CRYPT_KEY_PROV_INFO prov_info = {0};
prov_info.pwszContainerName = containerW;
prov_info.pwszProvName = (LPWSTR) L"Microsoft Enhanced RSA and AES Cryptographic Provider";
prov_info.dwProvType = PROV_RSA_AES;
prov_info.dwKeySpec = AT_KEYEXCHANGE;
if (!CertSetCertificateContextProperty(cert_ctx,
CERT_KEY_PROV_INFO_PROP_ID, 0, &prov_info))
goto done;
}
mem_store = CertOpenStore(CERT_STORE_PROV_MEMORY, 0, 0, 0, NULL);
if (!mem_store) goto done;
if (!CertAddCertificateContextToStore(mem_store, cert_ctx,
CERT_STORE_ADD_ALWAYS, NULL))
goto done;
// Two-pass PFX export: get size, then export
pfx.cbData = 0;
pfx.pbData = NULL;
if (!PFXExportCertStoreEx(mem_store, &pfx, L"", NULL, EXPORT_PRIVATE_KEYS))
goto done;
pfx.pbData = malloc(pfx.cbData);
if (!pfx.pbData) goto done;
if (!PFXExportCertStoreEx(mem_store, &pfx, L"", NULL, EXPORT_PRIVATE_KEYS))
goto done;
// --- Import PFX blob and acquire SChannel credential ---
result = load_pfx_blob(pfx.pbData, pfx.cbData, cred_out, cert_ctx_out);
done:
free(pfx.pbData);
if (mem_store) CertCloseStore(mem_store, 0);
if (hkey) CryptDestroyKey(hkey);
if (hprov) CryptReleaseContext(hprov, 0);
{ HCRYPTPROV tmp = 0;
CryptAcquireContextA(&tmp, container, MS_ENH_RSA_AES_PROV_A,
PROV_RSA_AES, CRYPT_DELETEKEYSET); }
if (pkcs8_buf) LocalFree(pkcs8_buf);
if (rsa_blob) LocalFree(rsa_blob);
free(key_der);
free(key_pem);
free(cert_der);
free(cert_pem);
if (cert_ctx) CertFreeCertificateContext(cert_ctx);
return result;
}
// ============================================================
// Global init/free (no-ops for SChannel)
// ============================================================
void tls_global_init(void)
{
}
void tls_global_free(void)
{
}
// ============================================================
// Server init/free
// ============================================================
int tls_server_init(TLS_Server *server, string cert_file, string key_file)
{
memset(server, 0, sizeof(*server));
server->num_certs = 0;
int ret = load_pem_credential(cert_file, key_file, &server->cred, &server->cert_ctx);
if (ret < 0)
return -1;
return 0;
}
void tls_server_free(TLS_Server *server)
{
FreeCredentialsHandle(&server->cred);
if (server->cert_ctx) {
HCRYPTPROV_OR_NCRYPT_KEY_HANDLE hkey = 0;
DWORD key_spec = 0;
BOOL caller_free = FALSE;
if (CryptAcquireCertificatePrivateKey(server->cert_ctx,
CRYPT_ACQUIRE_ONLY_NCRYPT_KEY_FLAG, NULL, &hkey, &key_spec, &caller_free)) {
if (caller_free && hkey)
NCryptFreeObject(hkey);
}
CertFreeCertificateContext(server->cert_ctx);
}
for (int i = 0; i < server->num_certs; i++)
FreeCredentialsHandle(&server->certs[i].cred);
}
int tls_server_add_cert(TLS_Server *server, string domain, string cert_file, string key_file)
{
if (server->num_certs >= TLS_CERT_LIMIT)
return -1;
TLS_Cert *cert = &server->certs[server->num_certs];
if (domain.len >= (int) sizeof(cert->domain))
return -1;
PCCERT_CONTEXT cert_ctx = NULL;
int ret = load_pem_credential(cert_file, key_file, &cert->cred, &cert_ctx);
if (ret < 0)
return -1;
if (cert_ctx) {
HCRYPTPROV_OR_NCRYPT_KEY_HANDLE hkey = 0;
DWORD key_spec = 0;
BOOL caller_free = FALSE;
if (CryptAcquireCertificatePrivateKey(cert_ctx,
CRYPT_ACQUIRE_ONLY_NCRYPT_KEY_FLAG, NULL, &hkey, &key_spec, &caller_free)) {
if (caller_free && hkey)
NCryptFreeObject(hkey);
}
CertFreeCertificateContext(cert_ctx);
}
memcpy(cert->domain, domain.ptr, domain.len);
cert->domain[domain.len] = '\0';
server->num_certs++;
return 0;
}
// ============================================================
// Connection init/free
// ============================================================
int tls_conn_init(TLS_Conn *conn, TLS_Server *server)
{
memset(conn, 0, sizeof(*conn));
conn->ctx_valid = false;
conn->handshake = true;
conn->cred = &server->cred;
byte_queue_init(&conn->in_buf, TLS_BUF_LIMIT);
byte_queue_init(&conn->out_buf, TLS_BUF_LIMIT);
conn->pending = NULL;
conn->pending_off = 0;
conn->pending_len = 0;
SecInvalidateHandle(&conn->ctx);
return 0;
}
void tls_conn_free(TLS_Conn *conn)
{
if (conn->ctx_valid)
DeleteSecurityContext(&conn->ctx);
byte_queue_free(&conn->in_buf);
byte_queue_free(&conn->out_buf);
free(conn->pending);
}
// ============================================================
// Handshake
// ============================================================
int tls_conn_handshake(TLS_Conn *conn)
{
assert(conn->handshake);
// Read available ciphertext from in_buf
ByteView in = byte_queue_read_buf(&conn->in_buf);
if (!in.ptr || in.len == 0) {
byte_queue_read_ack(&conn->in_buf, 0);
return 0;
}
int in_avail = (int) in.len;
// Input buffers
SecBuffer in_bufs[2];
in_bufs[0].BufferType = SECBUFFER_TOKEN;
in_bufs[0].pvBuffer = in.ptr;
in_bufs[0].cbBuffer = (unsigned long) in_avail;
in_bufs[1].BufferType = SECBUFFER_EMPTY;
in_bufs[1].pvBuffer = NULL;
in_bufs[1].cbBuffer = 0;
SecBufferDesc in_desc;
in_desc.ulVersion = SECBUFFER_VERSION;
in_desc.cBuffers = 2;
in_desc.pBuffers = in_bufs;
// Output buffers
SecBuffer out_bufs[1];
out_bufs[0].BufferType = SECBUFFER_TOKEN;
out_bufs[0].pvBuffer = NULL;
out_bufs[0].cbBuffer = 0;
SecBufferDesc out_desc;
out_desc.ulVersion = SECBUFFER_VERSION;
out_desc.cBuffers = 1;
out_desc.pBuffers = out_bufs;
DWORD flags = ASC_REQ_STREAM
| ASC_REQ_SEQUENCE_DETECT
| ASC_REQ_REPLAY_DETECT
| ASC_REQ_CONFIDENTIALITY
| ASC_REQ_ALLOCATE_MEMORY;
DWORD out_flags = 0;
TimeStamp expiry;
SECURITY_STATUS ss = AcceptSecurityContext(
conn->cred,
conn->ctx_valid ? &conn->ctx : NULL,
&in_desc,
flags,
0,
conn->ctx_valid ? NULL : &conn->ctx,
&out_desc,
&out_flags,
&expiry);
if (ss == SEC_E_OK || ss == SEC_I_CONTINUE_NEEDED)
conn->ctx_valid = true;
// Copy output token to out_buf
if (out_bufs[0].pvBuffer && out_bufs[0].cbBuffer > 0) {
byte_queue_write_setmincap(&conn->out_buf, out_bufs[0].cbBuffer);
byte_queue_write(&conn->out_buf, out_bufs[0].pvBuffer, out_bufs[0].cbBuffer);
FreeContextBuffer(out_bufs[0].pvBuffer);
}
// Calculate how much input was consumed
int consumed = in_avail;
if (in_bufs[1].BufferType == SECBUFFER_EXTRA && in_bufs[1].cbBuffer > 0)
consumed = in_avail - (int) in_bufs[1].cbBuffer;
if (ss == SEC_E_INCOMPLETE_MESSAGE) {
// SChannel didn't consume anything
byte_queue_read_ack(&conn->in_buf, 0);
return 0;
}
byte_queue_read_ack(&conn->in_buf, consumed);
if (ss == SEC_I_CONTINUE_NEEDED)
return 0;
if (ss == SEC_E_OK) {
conn->handshake = false;
ss = QueryContextAttributes(&conn->ctx, SECPKG_ATTR_STREAM_SIZES, &conn->sizes);
if (ss != SEC_E_OK)
return -1;
return 1;
}
return -1;
}
// ============================================================
// Network I/O (ciphertext ↔ socket)
// ============================================================
char *tls_conn_net_write_buf(TLS_Conn *conn, int *cap)
{
byte_queue_write_setmincap(&conn->in_buf, 4096);
ByteView bv = byte_queue_write_buf(&conn->in_buf);
if (!bv.ptr || bv.len == 0) {
byte_queue_write_ack(&conn->in_buf, 0);
return NULL;
}
*cap = (int) bv.len;
return (char *) bv.ptr;
}
void tls_conn_net_write_ack(TLS_Conn *conn, int num)
{
byte_queue_write_ack(&conn->in_buf, num);
}
char *tls_conn_net_read_buf(TLS_Conn *conn, int *num)
{
ByteView bv = byte_queue_read_buf(&conn->out_buf);
if (!bv.ptr || bv.len == 0) {
byte_queue_read_ack(&conn->out_buf, 0);
return NULL;
}
*num = (int) bv.len;
return (char *) bv.ptr;
}
void tls_conn_net_read_ack(TLS_Conn *conn, int num)
{
byte_queue_read_ack(&conn->out_buf, num);
}
// ============================================================
// Application I/O (encrypt/decrypt)
// ============================================================
int tls_conn_app_write(TLS_Conn *conn, char *src, int num)
{
assert(!conn->handshake);
if (num <= 0) return 0;
int max_msg = (int) conn->sizes.cbMaximumMessage;
if (num > max_msg)
num = max_msg;
int header_size = (int) conn->sizes.cbHeader;
int trailer_size = (int) conn->sizes.cbTrailer;
int total = header_size + num + trailer_size;
// Ensure output buffer has enough space
byte_queue_write_setmincap(&conn->out_buf, total);
ByteView bv = byte_queue_write_buf(&conn->out_buf);
if (!bv.ptr || (int) bv.len < total) {
// Try with less data
if (!bv.ptr || (int) bv.len < header_size + trailer_size + 1) {
byte_queue_write_ack(&conn->out_buf, 0);
return 0;
}
num = (int) bv.len - header_size - trailer_size;
total = header_size + num + trailer_size;
}
char *out_ptr = (char *) bv.ptr;
// Copy plaintext into the data portion
memcpy(out_ptr + header_size, src, num);
// Set up SecBuffers for in-place encryption
SecBuffer bufs[4];
bufs[0].BufferType = SECBUFFER_STREAM_HEADER;
bufs[0].pvBuffer = out_ptr;
bufs[0].cbBuffer = (unsigned long) header_size;
bufs[1].BufferType = SECBUFFER_DATA;
bufs[1].pvBuffer = out_ptr + header_size;
bufs[1].cbBuffer = (unsigned long) num;
bufs[2].BufferType = SECBUFFER_STREAM_TRAILER;
bufs[2].pvBuffer = out_ptr + header_size + num;
bufs[2].cbBuffer = (unsigned long) trailer_size;
bufs[3].BufferType = SECBUFFER_EMPTY;
bufs[3].pvBuffer = NULL;
bufs[3].cbBuffer = 0;
SecBufferDesc desc;
desc.ulVersion = SECBUFFER_VERSION;
desc.cBuffers = 4;
desc.pBuffers = bufs;
SECURITY_STATUS ss = EncryptMessage(&conn->ctx, 0, &desc, 0);
if (ss != SEC_E_OK) {
byte_queue_write_ack(&conn->out_buf, 0);
return -1;
}
int written = (int)(bufs[0].cbBuffer + bufs[1].cbBuffer + bufs[2].cbBuffer);
byte_queue_write_ack(&conn->out_buf, written);
return num;
}
int tls_conn_app_read(TLS_Conn *conn, char *dst, int cap)
{
assert(!conn->handshake);
// Drain any pending plaintext from a previous partial read
if (conn->pending_len > 0) {
int n = conn->pending_len;
if (n > cap) n = cap;
memcpy(dst, conn->pending + conn->pending_off, n);
conn->pending_off += n;
conn->pending_len -= n;
if (conn->pending_len == 0) {
free(conn->pending);
conn->pending = NULL;
conn->pending_off = 0;
}
return n;
}
ByteView in = byte_queue_read_buf(&conn->in_buf);
if (!in.ptr || in.len == 0) {
byte_queue_read_ack(&conn->in_buf, 0);
return 0;
}
int in_avail = (int) in.len;
// DecryptMessage operates in-place
SecBuffer bufs[4];
bufs[0].BufferType = SECBUFFER_DATA;
bufs[0].pvBuffer = in.ptr;
bufs[0].cbBuffer = (unsigned long) in_avail;
bufs[1].BufferType = SECBUFFER_EMPTY;
bufs[1].pvBuffer = NULL;
bufs[1].cbBuffer = 0;
bufs[2].BufferType = SECBUFFER_EMPTY;
bufs[2].pvBuffer = NULL;
bufs[2].cbBuffer = 0;
bufs[3].BufferType = SECBUFFER_EMPTY;
bufs[3].pvBuffer = NULL;
bufs[3].cbBuffer = 0;
SecBufferDesc desc;
desc.ulVersion = SECBUFFER_VERSION;
desc.cBuffers = 4;
desc.pBuffers = bufs;
SECURITY_STATUS ss = DecryptMessage(&conn->ctx, &desc, 0, NULL);
if (ss == SEC_E_INCOMPLETE_MESSAGE) {
byte_queue_read_ack(&conn->in_buf, 0);
return 0;
}
if (ss != SEC_E_OK && ss != SEC_I_RENEGOTIATE) {
byte_queue_read_ack(&conn->in_buf, 0);
return -1;
}
// Find decrypted data and extra ciphertext buffers
SecBuffer *data_buf = NULL;
SecBuffer *extra_buf = NULL;
for (int i = 0; i < 4; i++) {
if (bufs[i].BufferType == SECBUFFER_DATA)
data_buf = &bufs[i];
else if (bufs[i].BufferType == SECBUFFER_EXTRA)
extra_buf = &bufs[i];
}
int result = 0;
if (data_buf && data_buf->cbBuffer > 0) {
int total = (int) data_buf->cbBuffer;
int n = total;
if (n > cap) n = cap;
memcpy(dst, data_buf->pvBuffer, n);
result = n;
// Save excess plaintext for next call
int leftover = total - n;
if (leftover > 0) {
conn->pending = malloc(leftover);
if (conn->pending) {
memcpy(conn->pending, (char *) data_buf->pvBuffer + n, leftover);
conn->pending_off = 0;
conn->pending_len = leftover;
}
}
}
// Consume processed input, keeping any extra ciphertext
int consumed = in_avail;
if (extra_buf && extra_buf->cbBuffer > 0)
consumed = in_avail - (int) extra_buf->cbBuffer;
byte_queue_read_ack(&conn->in_buf, consumed);
if (ss == SEC_I_RENEGOTIATE)
return result > 0 ? result : 0;
return result;
}
int tls_conn_needs_flushing(TLS_Conn *conn)
{
return !byte_queue_empty(&conn->out_buf);
}
#endif // TLS_SCHANNEL
#endif // TLS_ENABLED