/core/src/roon/connection.zig at f3fa46a40da6b38662281273307383a9a41af9e6

// Copyright 2025 Shota FUJI
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// SPDX-License-Identifier: Apache-2.0

const std = @import("std");

const moo = @import("moo");
const websocket = @import("websocket");

const Response = struct {
    wrote: *std.Thread.ResetEvent,

    /// `null` on OOM.
    data: ?[]const u8,
};

const ResponsesStore = std.AutoHashMap(i64, *Response);

const RequestHandler = *const fn (
    response_writer: anytype,
    meta: moo.Metadata,
    header_ctx: moo.HeaderParsingContext,
    message: []const u8,
) anyerror!bool;

pub const Connection = struct {
    allocator: std.mem.Allocator,
    thread_safe_allocator: *std.heap.ThreadSafeAllocator,
    ws: websocket.Client,
    addr: []const u8,
    rng: std.Random.Xoshiro256,
    thread: ?std.Thread = null,
    responses: ResponsesStore,
    responses_mutex: *std.Thread.Mutex,

    pub const InitError = error{
        WebSocketClientCreationError,
        WebSocketHandshakeError,
        PRNGSeedGenerationFailure,
    } || std.mem.Allocator.Error;

    pub fn init(child_allocator: std.mem.Allocator, address: std.net.Address) InitError!@This() {
        var tsa = try child_allocator.create(std.heap.ThreadSafeAllocator);
        tsa.* = std.heap.ThreadSafeAllocator{ .child_allocator = child_allocator };
        const allocator = tsa.allocator();

        var addr = std.ArrayList(u8).init(allocator);
        defer addr.deinit();

        try addr.writer().print("{}", .{address});

        var addr_string = try addr.toOwnedSlice();
        errdefer allocator.free(addr_string);

        // Zig std always prints "<addr>:<port>" for IPv4 and IPv6
        const port_start = std.mem.lastIndexOfScalar(u8, addr_string, ':') orelse {
            unreachable;
        };

        var client = websocket.Client.init(allocator, .{
            .port = address.getPort(),
            .host = addr_string[0..port_start],
        }) catch return InitError.WebSocketClientCreationError;
        errdefer client.deinit();

        std.log.debug("Performing WebSocket handshake...", .{});
        client.handshake("/api", .{
            .timeout_ms = 1_000,
        }) catch return InitError.WebSocketHandshakeError;

        const responses = ResponsesStore.init(allocator);

        const responses_mutex = try allocator.create(std.Thread.Mutex);
        responses_mutex.* = std.Thread.Mutex{};

        return .{
            .allocator = allocator,
            .ws = client,
            .addr = addr_string,
            .rng = std.Random.DefaultPrng.init(seed: {
                var seed: u64 = undefined;
                std.posix.getrandom(std.mem.asBytes(&seed)) catch {
                    return InitError.PRNGSeedGenerationFailure;
                };
                break :seed seed;
            }),
            .responses = responses,
            .responses_mutex = responses_mutex,
            .thread_safe_allocator = tsa,
        };
    }

    pub fn listen(self: *Connection, on_request: RequestHandler) std.Thread.SpawnError!void {
        std.log.debug("Spawning WebSocket request handler thread...", .{});
        self.thread = try std.Thread.spawn(.{}, readLoop, .{ self, on_request });
    }

    pub fn deinit(self: *Connection) void {
        std.log.debug("Closing WebSocket connection...", .{});
        self.ws.close(.{}) catch |err| {
            std.log.warn("Failed to close WebSocket connection, proceeding: {s}", .{@errorName(err)});
        };
        if (self.thread) |thread| {
            std.log.debug("Waiting WebSocket thread to terminate...", .{});
            // Wait for read thread to terminate.
            thread.join();
            self.thread = null;
        }
        self.ws.deinit();

        self.responses_mutex.lock();
        // TODO: Release response bytes.
        self.responses.deinit();
        self.responses_mutex.unlock();

        self.allocator.destroy(self.responses_mutex);

        self.allocator.free(self.addr);
        self.thread_safe_allocator.child_allocator.destroy(self.thread_safe_allocator);
        self.thread_safe_allocator = undefined;
    }

    pub fn newRequestId(self: *Connection) i64 {
        return self.rng.random().int(i64);
    }

    pub fn request(self: *Connection, request_id: i64, message: []u8) ![]const u8 {
        var wrote = std.Thread.ResetEvent{};
        var response = Response{
            .wrote = &wrote,
            .data = null,
        };

        {
            self.responses_mutex.lock();
            defer self.responses_mutex.unlock();
            try self.responses.put(request_id, &response);
        }

        defer {
            self.responses_mutex.lock();
            _ = self.responses.remove(request_id);
            self.responses_mutex.unlock();
        }

        try self.ws.writeBin(message);

        wrote.wait();

        return response.data orelse return std.mem.Allocator.Error.OutOfMemory;
    }
};

// TODO: Better logging
fn readLoop(conn: *Connection, on_request: RequestHandler) void {
    conn.ws.readTimeout(1_000) catch |err| {
        std.log.err("Unable to set WebSocket read timeout: {s}", .{@errorName(err)});
        return;
    };

    while (true) {
        const msg = conn.ws.read() catch return orelse continue;
        defer conn.ws.done(msg);

        std.log.debug("Received WebSocket message: type={s}", .{@tagName(msg.type)});

        switch (msg.type) {
            // NOTE: roon-node-api does not check whether message is binaryType.
            .text, .binary => {
                const meta, const header_ctx = moo.Metadata.parse(msg.data) catch |err| {
                    std.log.warn("Failed to parse MOO metadata: {s}", .{@errorName(err)});
                    continue;
                };

                // We just want to know Request-Id header.
                const header, _ = moo.NoBodyHeaders.parse(msg.data, header_ctx) catch |err| {
                    std.log.warn("Failed to parse MOO headers: {s}", .{@errorName(err)});
                    continue;
                };

                {
                    conn.responses_mutex.lock();
                    defer conn.responses_mutex.unlock();

                    if (conn.responses.get(header.request_id)) |store| {
                        if (store.wrote.isSet()) {
                            std.log.warn(
                                "Received more than one message having same Request-Id({d})",
                                .{header.request_id},
                            );
                            continue;
                        }

                        defer store.wrote.set();

                        const bytes = conn.allocator.dupe(u8, msg.data) catch |err| {
                            std.log.err("Unable to release incoming WS message: {s}", .{@errorName(err)});
                            return;
                        };

                        store.data = bytes;
                        continue;
                    }
                }

                var buffer = std.ArrayList(u8).init(conn.allocator);
                defer buffer.deinit();

                const wrote = on_request(buffer.writer(), meta, header_ctx, msg.data) catch |err| {
                    std.log.warn("Service server handler returned an error: {s}", .{@errorName(err)});
                    continue;
                };

                if (!wrote) {
                    std.log.info("Unhandled incoming request: service={s}\n", .{meta.service});
                    continue;
                }

                const bytes = buffer.toOwnedSlice() catch |err| {
                    std.log.warn("Unable to prepare response bytes: {s}\n", .{@errorName(err)});
                    continue;
                };

                conn.ws.writeBin(bytes) catch |err| {
                    std.log.warn("Failed to write response message: {s}\n", .{@errorName(err)});
                    continue;
                };
            },
            .ping => conn.ws.writePong(msg.data) catch {},
            .pong => {},
            .close => {
                conn.ws.close(.{}) catch return;
                break;
            },
        }
    }
}