#!/usr/bin/env bash
set -euo pipefail
umask 077

[[ "$#" -eq 1 || ("$#" -eq 2 && "$2" = --render) ]] || {
  echo 'usage: apply-wireguard-mss.sh MTU [--render]' >&2
  exit 2
}
mtu="$1"
[[ "$mtu" =~ ^[1-9][0-9]{3}$ && "$mtu" -ge 1280 && "$mtu" -le 1420 ]] || exit 2
mss4=$((mtu - 40))
mss6=$((mtu - 60))
render() {
  printf 'table inet wireguard_mss {\n  chain forward {\n    type filter hook forward priority mangle; policy accept;\n'
  for direction in iifname oifname; do
    printf '    %s "wg0" meta nfproto ipv4 tcp flags & (syn | rst) == syn tcp option maxseg size > %s tcp option maxseg size set %s\n' "$direction" "$mss4" "$mss4"
    printf '    %s "wg0" meta nfproto ipv6 tcp flags & (syn | rst) == syn tcp option maxseg size > %s tcp option maxseg size set %s\n' "$direction" "$mss6" "$mss6"
  done
  printf '  }\n}\n'
}
if [[ "${2:-}" = --render ]]; then
  render
  exit 0
fi

config=/etc/sysconfig/nftables.conf
fragment=/etc/nftables/wireguard-mss.nft
if [[ "${WG_MSS_TEST:-}" = 1 ]]; then
  config="${WG_MSS_CONFIG:?}"
  fragment="${WG_MSS_FRAGMENT:?}"
fi
[[ -f "$config" && ! -L "$config" && ! -L "$fragment" ]] || exit 1
include="include \"$fragment\""
[[ "$(grep -Fxc "$include" "$config" || true)" -le 1 ]] || exit 1
work="$(mktemp -d /var/tmp/wireguard-mss.XXXXXX)"
cp -a "$config" "$work/config.before"
fragment_existed=false
if [[ -e "$fragment" ]]; then
  cp -a "$fragment" "$work/fragment.before"
  fragment_existed=true
fi
# Fail on inspection errors rather than confusing them with an absent owned table.
nft list tables >"$work/tables"
table_existed=false
if grep -Fxq 'table inet wireguard_mss' "$work/tables"; then
  nft list table inet wireguard_mss >"$work/table.before"
  table_existed=true
fi
render >"$work/fragment"
cp -a "$config" "$work/config"
if ! grep -Fxq "$include" "$work/config"; then printf '\n%s\n' "$include" >>"$work/config"; fi
# Replace only our exact include in the validation copy, never the live ruleset.
awk -v include_line="$include" -v candidate="$work/fragment" '
  $0 == include_line { print "include \"" candidate "\""; next }
  { print }
' "$work/config" >"$work/config.check"
nft --check --file "$work/config.check"
: >"$work/live.batch"
if "$table_existed"; then printf 'delete table inet wireguard_mss\n' >>"$work/live.batch"; fi
cat "$work/fragment" >>"$work/live.batch"
nft --check --file "$work/live.batch"

persistence_changed=false
cleanup() {
  local status=$?
  if [[ "$status" -ne 0 && "$persistence_changed" = true ]]; then
    cp -a "$work/config.before" "$config" || echo 'ERROR: nft persistence rollback failed' >&2
    if "$fragment_existed"; then
      cp -a "$work/fragment.before" "$fragment" || echo 'ERROR: MSS fragment rollback failed' >&2
    else
      rm -f "$fragment" || echo 'ERROR: MSS fragment removal failed' >&2
    fi
  fi
  printf 'MSS backup evidence: %s\n' "$work"
  exit "$status"
}
trap cleanup EXIT
install -d -m 0755 "$(dirname "$fragment")"
persistence_changed=true
install -m 0644 "$work/fragment" "$fragment"
cp -a "$work/config" "$config"
# Last fallible mutation: nft submits this batch atomically; no full flush/restart.
nft --file "$work/live.batch"
printf 'WireGuard MSS caps: IPv4=%s IPv6=%s; no firewall exposure changes\n' "$mss4" "$mss6"