#!/usr/bin/perl
# truenas-plugin-broker
#
# Holds one persistent authenticated WebSocket per (host, api_key) pair to
# the TrueNAS middleware and proxies JSON-RPC for any number of local PVE
# plugin clients via a Unix socket.
#
# Goal: every PVE plugin invocation (forked qm/pvesh/pveproxy worker, every
# pvedaemon process) reuses the broker's WS instead of opening a fresh one.
# TrueNAS rate-limits auth.login_with_api_key at 20 calls / 60 s / IP /
# method, so per-process re-auth used to trip a sticky 60 s lockout
# (death-spiral D2). With the broker, the node logs in once per upstream
# session and never re-auths under normal operation.
#
# Protocol (newline-delimited JSON over Unix socket):
#   client -> broker:
#     {"scfg":{"api_host":"...","api_key":"...","api_scheme":"wss",
#              "api_insecure":0,"prefer_ipv4":1,"api_port":443},
#      "method":"system.version","params":[]}
#   broker -> client:
#     {"result":...}            on success
#     {"error":"..."}           on failure (string contains TN error body)
#
# Trust model: Unix socket is root-only (mode 0600). API keys are passed
# in each request because the broker is stateless w.r.t. credentials.

use strict;
use warnings;
use IO::Socket::UNIX;
use IO::Socket::SSL;
use IO::Socket::INET;
use IO::Select;
use Fcntl qw(:flock);
use POSIX qw(setsid strftime WNOHANG);
use Digest::SHA qw(sha1 sha1_hex);
use MIME::Base64 qw(encode_base64);
use JSON::PP;
use Socket;
use Time::HiRes qw(time);

my $SOCKET_PATH    = '/run/truenas-plugin/broker.sock';
my $PIDFILE        = '/run/truenas-plugin/broker.pid';
my $LOG_FILE       = '/var/log/truenas-plugin-broker.log';
my $IDLE_REAP_SECS = 600;   # close WS conns idle longer than this
my $WS_PATH = '/api/current';

# Cache the pool-entry liveness probe (`auth.me`) result for up to
# this many seconds. Every plugin call goes through get_or_open_ws();
# running an auth-required RPC on every call adds a per-call round
# trip whose per-call cost is small in isolation but compounds under
# multi-node load, and 'auth.me' in particular returns a large
# response (full user + role list, several KiB) that TN's middleware
# serialises around user-record locks -- enough to trip 596 API
# timeouts in PVE under a busy 3-node test suite. Caching the last
# successful probe timestamp per pool entry and skipping the probe
# when it is <LIVENESS_CACHE_SECS old costs essentially nothing on
# healthy connections and still detects TrueNAS's 30-day silent
# session expiry (AA_LEVEL1.max_session_age, issue #98) within one
# cache window. Round-trip cost drops from "every call" to "at most
# one per LIVENESS_CACHE_SECS per pool entry".
my $LIVENESS_CACHE_SECS = 60;

# Prefork worker count. Each worker holds one upstream WS (auth'd once) and
# accepts on the shared listen fd; the kernel load-balances accept() across
# them. Under 3-node cluster test load a single-threaded broker serialized
# every RPC through one WS, so 4 concurrent client calls (extent.query x2,
# targetextent.query, pool.dataset.query) took 9.8 s wall vs 0.14+1.6+3.2+
# 9.8 ≈ 15 s serial-sum -- close to fully serialized. With N workers each
# client call parallelizes up to N-way. TN's rate-limit on
# auth.login_with_api_key is 20 / 60 s / IP; N=4 auths at broker startup
# stay well under that (and each worker only re-auths on WS disconnect).
# Bump this if a node has many pvedaemon workers; keep it small enough
# that startup burst + reconnect churn stays inside TN's 20/60s ceiling
# with headroom for other operational traffic.
my $PREFORK_N = 4;

my $foreground = grep { $_ eq '--foreground' || $_ eq '-f' } @ARGV;
my $debug      = grep { $_ eq '--debug' || $_ eq '-d' } @ARGV;

# ---------------- logging ----------------
sub _log_line {
    my ($lvl, $msg) = @_;
    my $ts = strftime('%Y-%m-%dT%H:%M:%S', gmtime());
    my $line = "$ts $lvl $msg\n";
    if ($foreground) {
        print STDERR $line;
    } else {
        if (open(my $fh, '>>', $LOG_FILE)) {
            print $fh $line;
            close $fh;
        }
    }
}
sub linfo  { _log_line('INFO',  $_[0]); }
sub lwarn  { _log_line('WARN',  $_[0]); }
sub lerr   { _log_line('ERROR', $_[0]); }
sub ldebug { _log_line('DEBUG', $_[0]) if $debug; }

# ---------------- setup ----------------
sub ensure_rundir {
    my $dir = '/run/truenas-plugin';
    mkdir($dir, 0700) unless -d $dir;
    chmod 0700, $dir;
}

sub daemonize {
    return if $foreground;
    defined(my $pid = fork()) or die "fork failed: $!";
    exit 0 if $pid;
    setsid() or die "setsid failed: $!";
    chdir('/') or die "chdir / failed: $!";
    open(STDIN,  '<', '/dev/null');
    open(STDOUT, '>>', '/dev/null');
    open(STDERR, '>>', $LOG_FILE) or open(STDERR, '>>', '/dev/null');
}

sub write_pidfile {
    open(my $fh, '>', $PIDFILE) or die "cannot write pidfile $PIDFILE: $!";
    print $fh "$$\n";
    close $fh;
}

# ---------------- WebSocket helpers (subset of plugin code) ----------------

# Resolve $host to an IPv4 literal so IO::Socket::INET/SSL cannot pick the
# AAAA record on a dual-stack TrueNAS when the caller asked for IPv4. Mirrors
# TrueNASPlugin.pm::_host_ipv4. Returns $host unchanged when it is already an
# IPv4 literal, when the lookup fails, or when the host is an IPv6 literal.
sub _host_ipv4 {
    my ($host) = @_;
    return $host if $host =~ /^\d+\.\d+\.\d+\.\d+$/;
    my @ent = Socket::gethostbyname($host);
    if (@ent && defined $ent[4]) {
        my $ip = inet_ntoa($ent[4]);
        return $ip if $ip;
    }
    return $host;
}

sub ws_open {
    my ($scfg) = @_;
    # Normalize scheme the same way TrueNASPlugin.pm::_ws_defaults does.
    # Without this, an operator who set tn_api_scheme=https (rejected by the
    # plugin's own validator only against the direct-WS path) would slip
    # through the broker envelope raw, and the eq 'wss' check below would
    # miss — the broker would open plain TCP and the handshake would 400.
    my $scheme = $scfg->{api_scheme} || 'wss';
    if    ($scheme =~ /^https$/i) { $scheme = 'wss'; }
    elsif ($scheme =~ /^http$/i)  { $scheme = 'ws';  }
    my $host   = $scfg->{api_host} or die "ws_open: no api_host";
    my $port   = $scfg->{api_port} // (($scheme eq 'wss') ? 443 : 80);
    # Honor prefer_ipv4 (default 1). The wire schema always listed this field
    # but the earlier broker ignored it, so the OS resolver could pick AAAA
    # on a dual-stack TN even though the plugin had resolved A already.
    my $peer   = (($scfg->{prefer_ipv4} // 1) ? _host_ipv4($host) : $host);

    my $sock;
    if ($scheme eq 'wss') {
        $sock = IO::Socket::SSL->new(
            PeerHost => $peer, PeerPort => $port,
            SSL_verify_mode => $scfg->{api_insecure} ? 0x00 : 0x02,
            SSL_hostname    => $host,
            Timeout => 15,
        ) or die "TLS connect to $host:$port failed: $SSL_ERROR";
    } else {
        $sock = IO::Socket::INET->new(
            PeerHost => $peer, PeerPort => $port, Proto => 'tcp', Timeout => 15,
        ) or die "TCP connect to $host:$port failed: $!";
    }

    my $key_raw = join '', map { chr(int(rand(256))) } 1..16;
    my $key_b64 = encode_base64($key_raw, '');
    my $req =
      "GET $WS_PATH HTTP/1.1\r\n".
      "Host: $host:$port\r\n".
      "Upgrade: websocket\r\n".
      "Connection: Upgrade\r\n".
      "Sec-WebSocket-Key: $key_b64\r\n".
      "Sec-WebSocket-Version: 13\r\n".
      "\r\n";
    print $sock $req;
    my $resp = '';
    while ($sock->sysread(my $buf, 1024)) {
        $resp .= $buf;
        last if $resp =~ /\r\n\r\n/s;
    }
    die "WS handshake failed: no HTTP 101 response" if $resp !~ m#^HTTP/1\.([01]) 101#;
    my ($accept) = $resp =~ /Sec-WebSocket-Accept:\s*(\S+)/i;
    my $expect = encode_base64(sha1($key_b64 . '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'), '');
    die "WS handshake invalid accept key" if ($accept // '') ne $expect;

    my $conn = { sock => $sock, next_id => 1 };
    my $auth = ws_rpc($conn, {
        jsonrpc => "2.0", id => $conn->{next_id}++,
        method  => "auth.login_with_api_key",
        params  => [ $scfg->{api_key} ],
    });
    die "TrueNAS authentication failed" unless $auth;
    return $conn;
}

sub xor_mask {
    my ($data, $mask) = @_;
    my $out = $data;
    my $len = length($data);
    for (my $i = 0; $i < $len; $i++) {
        substr($out, $i, 1) = chr(ord(substr($data, $i, 1)) ^ ord(substr($mask, $i % 4, 1)));
    }
    return $out;
}

sub ws_read_exact {
    my ($sock, $bufref, $len) = @_;
    $$bufref = '';
    while (length($$bufref) < $len) {
        my $got = $sock->sysread(my $tmp, $len - length($$bufref));
        return 0 unless defined $got && $got > 0;
        $$bufref .= $tmp;
    }
    return 1;
}

sub ws_send_text {
    my ($sock, $text) = @_;
    my $len = length($text);
    my $mask_key = join '', map { chr(int(rand(256))) } 1..4;
    my $hdr = pack('C', 0x81);   # FIN=1, opcode=0x1 (text)
    my $masked_flag = 0x80;
    if    ($len <= 125)       { $hdr .= pack('C',   $masked_flag | $len); }
    elsif ($len <= 0xFFFF)    { $hdr .= pack('C n', $masked_flag | 126, $len); }
    else                      { $hdr .= pack('C Q>',$masked_flag | 127, $len); }
    $sock->syswrite($hdr . $mask_key . xor_mask($text, $mask_key));
}

sub ws_recv_text {
    my ($sock) = @_;
    my $message = '';
    while (1) {
        my $hdr;
        return undef unless ws_read_exact($sock, \$hdr, 2);
        my ($b1, $b2) = unpack('CC', $hdr);
        my $fin    = $b1 & 0x80;
        my $opcode = $b1 & 0x0F;
        my $masked = $b2 & 0x80;
        my $len    = $b2 & 0x7F;
        if ($len == 126) {
            my $ext; return undef unless ws_read_exact($sock, \$ext, 2);
            $len = unpack('n', $ext);
        } elsif ($len == 127) {
            my $ext; return undef unless ws_read_exact($sock, \$ext, 8);
            $len = unpack('Q>', $ext);
        }
        my $mask_key = '';
        if ($masked) {
            return undef unless ws_read_exact($sock, \$mask_key, 4);
        }
        my $payload = '';
        if ($len > 0) {
            return undef unless ws_read_exact($sock, \$payload, $len);
            $payload = xor_mask($payload, $mask_key) if $masked;
        }
        if    ($opcode == 0x01) { $message  = $payload; return $message if $fin; }
        elsif ($opcode == 0x00) { $message .= $payload; return $message if $fin; }
        elsif ($opcode == 0x08) { return undef; }            # close
        elsif ($opcode == 0x09) {                            # ping
            my $pong = pack('C', 0x8A);
            if ($len <= 125)       { $pong .= pack('C', $len); }
            elsif ($len <= 0xFFFF) { $pong .= pack('C n', 126, $len); }
            else                   { $pong .= pack('C Q>', 127, $len); }
            $sock->syswrite($pong . $payload);
        }
        elsif ($opcode == 0x0A) { }                          # pong
        else { return undef; }
    }
}

sub ws_rpc {
    my ($conn, $obj) = @_;
    my $req_id = $obj->{id};
    my $text = encode_json($obj);
    ws_send_text($conn->{sock}, $text);
    my $max_skip = 10;
    for (my $skipped = 0; $skipped <= $max_skip; $skipped++) {
        my $resp = ws_recv_text($conn->{sock});
        return undef unless defined $resp;
        my $decoded = eval { decode_json($resp) };
        next if $@ || !$decoded;
        if (defined($req_id) && ref($decoded) eq 'HASH'
            && defined($decoded->{id}) && "$decoded->{id}" eq "$req_id")
        {
            return $decoded;
        }
    }
    return undef;
}

# ---------------- per-(host,key) connection pool ----------------
my %POOL;       # scfg_key -> { conn => $ws_conn, last_used => $epoch }

sub scfg_key {
    my ($scfg) = @_;
    my $host = $scfg->{api_host} // '';
    my $key  = $scfg->{api_key}  // '';
    # Don't put raw api_key in memory key; hash it
    return $host . ':' . sha1_hex($key);
}

sub get_or_open_ws {
    my ($scfg) = @_;
    my $sk = scfg_key($scfg);
    my $entry = $POOL{$sk};
    if ($entry && $entry->{conn} && $entry->{conn}->{sock}) {
        # Liveness check: call an AUTHENTICATED method to detect
        # TrueNAS's silent 30-day session expiry (AA_LEVEL1.max_session_age;
        # 'core.ping' is no_auth_required and can't). 'auth.me' fits the
        # semantic but its response is large and TN serialises around
        # user-record locks; running it on every call caused 596 PVE-API
        # timeouts under multi-node load. Cache the last successful probe
        # per pool entry and skip the probe if it's still fresh -- a
        # 30-day expiry is still detected within LIVENESS_CACHE_SECS of
        # the next call, and the round-trip cost drops to near-zero on
        # healthy connections. See issue #98 for the full write-up.
        my $now = time();
        my $last_ok = $entry->{last_liveness_ok} // 0;
        if ($now - $last_ok < $LIVENESS_CACHE_SECS) {
            $entry->{last_used} = $now;
            return $entry->{conn};
        }
        my $ok = eval {
            ws_rpc($entry->{conn}, {
                jsonrpc => "2.0", id => 999999, method => "auth.me", params => [],
            });
        };
        if ($ok) {
            $entry->{last_used} = $now;
            $entry->{last_liveness_ok} = $now;
            return $entry->{conn};
        }
        # Dead or unauthenticated; drop and reopen (which re-auths).
        eval { $entry->{conn}->{sock}->close(); };
        delete $POOL{$sk};
    }
    my $conn = ws_open($scfg);
    my $now = time();
    # ws_open() ran auth.login_with_api_key against TN successfully to
    # get here, so the session is definitionally authenticated as of
    # `now`. Stamp last_liveness_ok so the first N calls (within
    # LIVENESS_CACHE_SECS) skip the auth.me probe entirely.
    $POOL{$sk} = { conn => $conn, last_used => $now, last_liveness_ok => $now };
    linfo("opened WS for $sk (pool size " . scalar(keys %POOL) . ")");
    return $conn;
}

sub reap_idle {
    my $now = time();
    for my $sk (keys %POOL) {
        my $e = $POOL{$sk};
        next unless $e && $e->{conn};
        if ($now - $e->{last_used} > $IDLE_REAP_SECS) {
            linfo("reaping idle WS for $sk");
            eval { $e->{conn}->{sock}->close(); };
            delete $POOL{$sk};
        }
    }
}

# ---------------- Unix-socket server ----------------
sub make_listener {
    unlink $SOCKET_PATH if -e $SOCKET_PATH;
    my $srv = IO::Socket::UNIX->new(
        Local  => $SOCKET_PATH,
        Listen => 32,
        Type   => SOCK_STREAM,
    ) or die "cannot bind $SOCKET_PATH: $!";
    chmod 0600, $SOCKET_PATH;
    return $srv;
}

# Read one line (ending in \n) from a non-blocking-ish client. Returns
# undef on EOF/error, string (without trailing \n) on success.
sub read_line_blocking {
    my ($fh) = @_;
    my $buf = '';
    while (1) {
        my $got = $fh->sysread(my $tmp, 4096);
        return undef unless defined $got && $got > 0;
        $buf .= $tmp;
        if ((my $idx = index($buf, "\n")) >= 0) {
            my $line = substr($buf, 0, $idx);
            return $line;
        }
    }
}

sub handle_request {
    my ($client) = @_;
    my $line = read_line_blocking($client);
    return unless defined $line;
    my $req = eval { decode_json($line) };
    if ($@ || !$req) {
        lwarn("bad request JSON: $@");
        my $err = encode_json({ error => "broker: bad request JSON: $@" });
        $client->syswrite("$err\n");
        return;
    }
    my $scfg = $req->{scfg};
    if (!$scfg || !$scfg->{api_host} || !$scfg->{api_key}) {
        my $h    = ref($scfg) eq 'HASH' && defined $scfg->{api_host} ? $scfg->{api_host} : '<undef>';
        my $kraw = ref($scfg) eq 'HASH' ? $scfg->{api_key} : undef;
        my $k    = !defined $kraw       ? '<undef>'
                 : length($kraw) == 0   ? '<empty>'
                                        : '<set len=' . length($kraw) . '>';
        my $m    = $req->{method} // '<no method>';
        lwarn("scfg missing api_host/api_key in request (api_host=$h api_key=$k method=$m)");
        if (ref($scfg) eq 'HASH') {
            lwarn("  scfg keys: " . join(",", sort keys %$scfg));
        } else {
            lwarn("  scfg is not a HASH (ref=" . ref($scfg) . ")");
        }
        # Dump the raw line so we can tell whether the client really
        # sent a malformed envelope or whether our decode is at fault.
        # Redact any api_key= occurrences before logging (defense in
        # depth; the key shouldn't be in this code path but we still
        # don't want it in journal). Truncate to 600 bytes.
        my $raw = $line // '<undef>';
        $raw =~ s/("api_key"\s*:\s*)"[^"]*"/$1"<redacted>"/g;
        $raw = substr($raw, 0, 600) . (length($raw) > 600 ? "...[truncated]" : "");
        lwarn("  raw request: $raw");
        # Top-level keys of the decoded request envelope.
        if (ref($req) eq 'HASH') {
            lwarn("  request top-level keys: " . join(",", sort keys %$req));
        }
        my $err = encode_json({ error => "broker: scfg missing api_host/api_key" });
        $client->syswrite("$err\n");
        return;
    }
    my $method = $req->{method};
    my $params = $req->{params} // [];
    if (!$method) {
        lwarn("request missing method (host=$scfg->{api_host})");
        my $err = encode_json({ error => "broker: no method" });
        $client->syswrite("$err\n");
        return;
    }

    my $sk = scfg_key($scfg);
    my $result_obj;
    eval {
        my $conn = get_or_open_ws($scfg);
        my $rpc = {
            jsonrpc => "2.0", id => $conn->{next_id}++,
            method => $method, params => $params,
        };
        $result_obj = ws_rpc($conn, $rpc);
    };
    if (my $err = $@) {
        # Drop the cached conn — likely dead — and surface the failure in
        # the journal so the operator can tell the difference between
        # "broker quiet because all RPCs succeed" and "broker quiet because
        # auth keeps failing". Before this lwarn the daemon was opaque on
        # failure: only WS lifecycle events were logged.
        lwarn("upstream RPC failed (sk=$sk method=$method): $err");
        if (my $e = delete $POOL{$sk}) {
            eval { $e->{conn}->{sock}->close(); };
        }
        my $msg = encode_json({ error => "broker: $err" });
        $client->syswrite("$msg\n");
        return;
    }
    if (!$result_obj) {
        lwarn("no response from upstream (sk=$sk method=$method)");
        my $msg = encode_json({ error => "broker: no response from upstream" });
        $client->syswrite("$msg\n");
        return;
    }

    if (exists $result_obj->{error}) {
        # Upstream returned a structured JSON-RPC error (e.g. EBUSY /
        # Rate Limit Exceeded, invalid params, method not found). Worth
        # logging so users can match a failing test to a TN-side cause.
        my $err_repr = ref($result_obj->{error}) ? encode_json($result_obj->{error}) : $result_obj->{error};
        # Downgrade the well-known "already gone" case to INFO. A DELETE
        # against a resource that TN reports as InstanceNotFound / ENOENT
        # is a benign race the callers already handle idempotently
        # (see e.g. free_image, _teardown_snapshot_device). Leaving it
        # at WARN level drowns real errors under vzdump-* snapshot-clone
        # teardown noise: Max Carrara's test_run7 2026-08-25 3-node
        # cluster-01 logged ~2000 of these per node per subrun.
        my $is_benign_delete_enoent =
            ($method =~ /\.delete$/) &&
            ($err_repr =~ /InstanceNotFound|\bENOENT\b|does not exist/i);
        if ($is_benign_delete_enoent) {
            linfo("upstream ENOENT on $method (already gone; sk=$sk)");
        } else {
            lwarn("upstream error (sk=$sk method=$method): $err_repr");
        }
        my $msg = encode_json({ error => encode_json($result_obj->{error}) });
        $client->syswrite("$msg\n");
    } else {
        my $msg = encode_json({ result => $result_obj->{result} });
        $client->syswrite("$msg\n");
    }
}

# ---------------- main ----------------
ensure_rundir();
daemonize();
write_pidfile();
linfo("truenas-plugin-broker starting (pid=$$, socket=$SOCKET_PATH)");

$SIG{PIPE} = 'IGNORE';
# NOTE: parent SIG{TERM}/INT is set AFTER the prefork below so it can
# also propagate to the worker pool. Workers install their own TERM
# handler in spawn_worker so they close their local %POOL WS connections
# and exit cleanly without touching parent's socket / pidfile.

my $srv = make_listener();

# Prefork the worker pool. Parent creates listen fd, forks N children; each
# child re-runs the accept loop against the shared fd. Kernel serializes
# accept() among children so no race. Parent then waits and re-spawns any
# child that dies.
my %WORKER_PIDS;   # pid => 1

sub spawn_worker {
    my $pid = fork();
    if (!defined $pid) {
        lerr("fork worker failed: $!");
        return;
    }
    if ($pid == 0) {
        # child: close its own WS pool cleanly on TERM; do NOT touch
        # parent's socket file or pidfile.
        $SIG{TERM} = $SIG{INT} = sub {
            for my $sk (keys %POOL) {
                my $e = $POOL{$sk};
                eval { $e->{conn}->{sock}->close() if $e && $e->{conn}; };
            }
            exit 0;
        };
        $SIG{CHLD} = 'DEFAULT';
        worker_loop($srv);
        exit 0;
    }
    $WORKER_PIDS{$pid} = 1;
    ldebug("spawned worker pid=$pid");
}

sub worker_loop {
    my ($srv) = @_;
    my $sel = IO::Select->new($srv);
    my $last_reap = time();
    while (1) {
        my @ready = $sel->can_read(60);
        for my $fh (@ready) {
            if ($fh == $srv) {
                my $c = $srv->accept();
                next unless $c;
                eval { handle_request($c); };
                if ($@) { lerr("handle_request died: $@"); }
                eval { $c->close(); };
            }
        }
        if (time() - $last_reap > 60) {
            reap_idle();
            $last_reap = time();
        }
    }
}

# Parent: spawn N workers, then supervise
for (1..$PREFORK_N) { spawn_worker(); }
linfo("prefork pool up (N=$PREFORK_N)");

# Parent supervision loop: reap dead workers, respawn to keep pool full.
# Set the parent's TERM/INT handler HERE (after workers are up) so it
# only affects the parent process -- workers already installed their own
# in spawn_worker.
$SIG{TERM} = $SIG{INT} = sub {
    linfo("parent shutting down, signalling workers");
    kill 'TERM', keys %WORKER_PIDS;
    # Give workers a moment to close their WS cleanly.
    for (1..10) {
        last unless keys %WORKER_PIDS;
        my $dead = waitpid(-1, WNOHANG);
        if ($dead > 0) { delete $WORKER_PIDS{$dead}; }
        else { select(undef, undef, undef, 0.1); }
    }
    unlink $SOCKET_PATH;
    unlink $PIDFILE;
    exit 0;
};

while (1) {
    my $dead = waitpid(-1, 0);
    if ($dead > 0) {
        delete $WORKER_PIDS{$dead};
        lwarn("worker pid=$dead exited (status=$?), respawning");
        spawn_worker();
    }
}
