play.wenpai.net/scripts/lib/nats-msg.sh

#!/usr/bin/env bash
# NATS 消息辅助库 — 被 source 使用，禁止 set -euo pipefail

NATS_MSG_CONTEXT="${NATS_MSG_CONTEXT:-vm-hub}"
NATS_MSG_DM_PREFIX="${NATS_MSG_DM_PREFIX:-vm.dm}"
NATS_MSG_GROUP_SUBJECT="${NATS_MSG_GROUP_SUBJECT:-vm.group}"
NATS_MSG_HMAC_SECRET="${NATS_MSG_HMAC_SECRET:-}"

if [[ -z "$NATS_MSG_HMAC_SECRET" ]]; then
    for _nats_hmac_path in \
        "${HOME}/.config/nats/hmac-secret" \
        "/mnt/shared-context/certs/hmac-secret" \
        "/Volumes/shared-context/certs/hmac-secret"
    do
        if [[ -f "$_nats_hmac_path" ]]; then
            NATS_MSG_HMAC_SECRET=$(<"$_nats_hmac_path")
            break
        fi
    done
    unset _nats_hmac_path
fi

nats_msg_available() {
    command -v nats >/dev/null 2>&1 && command -v jq >/dev/null 2>&1
}

nats_msg_subject_for_target() {
    local target="$1"
    if [[ "$target" == "all" || "$target" == "group" ]]; then
        printf '%s' "$NATS_MSG_GROUP_SUBJECT"
    else
        printf '%s.%s' "$NATS_MSG_DM_PREFIX" "$target"
    fi
}

nats_msg_sign() {
    local id="$1" from="$2" type="$3" content="$4" timestamp="$5" subject="${6:-}"
    [[ -n "$NATS_MSG_HMAC_SECRET" ]] || return 0
    local data="${id}|${from}|${type}|${subject}|${content}|${timestamp}"
    printf '%s' "$data" | openssl dgst -sha256 -hmac "$NATS_MSG_HMAC_SECRET" -binary | base64 | tr -d '\n'
}

nats_msg_json_object() {
    local raw_json='{}'
    if [[ $# -ge 1 && -n "${1:-}" ]]; then
        raw_json="$1"
    fi
    printf '%s\n' "$raw_json" | jq -ce 'if type == "object" then . else error("object required") end' 2>/dev/null
}

nats_msg_verify_payload() {
    local payload="$1"
    command -v jq >/dev/null 2>&1 || return 1
    command -v openssl >/dev/null 2>&1 || return 1

    local id from type content timestamp subject sig expected
    id=$(printf '%s\n' "$payload" | jq -r 'if type == "object" then (.id // "") else "" end' 2>/dev/null) || return 1
    from=$(printf '%s\n' "$payload" | jq -r 'if type == "object" then (.from // "") else "" end' 2>/dev/null) || return 1
    type=$(printf '%s\n' "$payload" | jq -r 'if type == "object" then (.type // "") else "" end' 2>/dev/null) || return 1
    content=$(printf '%s\n' "$payload" | jq -r 'if type == "object" then (.content // "") else "" end' 2>/dev/null) || return 1
    timestamp=$(printf '%s\n' "$payload" | jq -r 'if type == "object" then (.timestamp // "") else "" end' 2>/dev/null) || return 1
    subject=$(printf '%s\n' "$payload" | jq -r 'if type == "object" then (.subject // "") else "" end' 2>/dev/null) || return 1
    sig=$(printf '%s\n' "$payload" | jq -r 'if type == "object" then (.sig // "") else "" end' 2>/dev/null) || return 1

    [[ -n "$id" && -n "$from" && -n "$type" && -n "$timestamp" && -n "$sig" ]] || return 1

    expected=$(nats_msg_sign "$id" "$from" "$type" "$content" "$timestamp" "$subject") || return 1
    [[ -n "$expected" && "$expected" == "$sig" ]]
}

nats_msg_publish() {
    local from="$1" target="$2" type="$3" subject="$4" content="$5"
    local priority="normal"
    local extra_json='{}'

    if [[ $# -ge 6 && -n "${6:-}" ]]; then
        priority="$6"
    fi
    if [[ $# -ge 7 && -n "${7:-}" ]]; then
        extra_json="$7"
    fi

    nats_msg_publish_subject "$from" "$(nats_msg_subject_for_target "$target")" "$type" "$subject" "$content" "$priority" "$extra_json"
}

nats_msg_publish_subject() {
    local from="$1" subject_name="$2" type="$3" subject="$4" content="$5"
    local priority="normal"
    local extra_json='{}'
    local extra_compact

    if [[ $# -ge 6 && -n "${6:-}" ]]; then
        priority="$6"
    fi
    if [[ $# -ge 7 && -n "${7:-}" ]]; then
        extra_json="$7"
    fi

    nats_msg_available || return 1
    extra_compact=$(nats_msg_json_object "$extra_json") || return 1

    local msg_id timestamp sig payload
    msg_id="${from}-$(date +%s%N)"
    timestamp="$(date -Iseconds)"
    sig=$(nats_msg_sign "$msg_id" "$from" "$type" "$content" "$timestamp" "$subject")

    payload=$(jq -nc \
        --arg id "$msg_id" \
        --arg from "$from" \
        --arg type "$type" \
        --arg priority "$priority" \
        --arg subject "$subject" \
        --arg content "$content" \
        --arg timestamp "$timestamp" \
        --arg sig "${sig:-}" \
        '{
          id:$id,
          from:$from,
          type:$type,
          priority:$priority,
          subject:$subject,
          content:$content,
          timestamp:$timestamp
        }
        + (if $sig != "" then {sig:$sig} else {} end)') || return 1

    extra_compact=$(printf '%s\n' "$extra_compact" | jq -c 'del(.id, .from, .type, .priority, .subject, .content, .timestamp, .sig)') || return 1

    payload=$(jq -c -s '.[0] + .[1]' \
        <(printf '%s\n' "$payload") \
        <(printf '%s\n' "$extra_compact")) || return 1

    NATS_CONTEXT="$NATS_MSG_CONTEXT" nats pub "$subject_name" "$payload" >/dev/null 2>&1
}