llmgw/internal/firewall/follow.go

package firewall

import (
	"context"
	"fmt"
	"strings"
)

// A store member that is nothing else runs no llmgw process: Patroni, etcd
// and PostgreSQL only. What the cluster's settings say of its firewall it
// reads from its own copy of the store, where the leading control plane
// keeps it (the cluster_state table, key firewall), once a minute, by a
// root timer running FollowScript. It turns llmgw's table off and on, and
// lets in the networks the settings add; what `server init` let in stays.

// FollowScript is the script, as python3 runs it. LLMGW_FW_STATE and
// LLMGW_FW_RULES replace the store and the rules file, for its tests.
const FollowScript = `#!/usr/bin/python3
# Written by llmgw: llmgw's firewall on this store member, as the cluster's
# settings say. Run by llmgw-firewall-follow.timer.
import ipaddress, json, os, subprocess, sys, tempfile

RULES = os.environ.get("LLMGW_FW_RULES", "` + RulesFile + `")

def run(*a):
    return subprocess.run(a, capture_output=True, text=True)

raw = os.environ.get("LLMGW_FW_STATE")
if raw is None:
    r = run("runuser", "-u", "postgres", "--", "psql", "-d", "llmgw", "-tAqc",
            "SELECT value FROM cluster_state WHERE key = 'firewall'")
    # No table yet, or nothing told: a control plane older than this, or
    # none leading. Left as it is.
    if r.returncode != 0:
        sys.exit(0)
    raw = r.stdout
if not raw.strip():
    sys.exit(0)
st = json.loads(raw)
present = run("nft", "list", "table", "inet", "` + Table + `").returncode == 0

if not st.get("enabled", True):
    if present:
        run("systemctl", "disable", "--now", "` + Unit + `")
        run("nft", "delete", "table", "inet", "` + Table + `")
        print("llmgw's firewall is off here")
    sys.exit(0)

if not os.path.exists(RULES):
    sys.exit(0)
lines = open(RULES).read().split("\n")
want = []
for n in st.get("networks") or []:
    try:
        want.append(ipaddress.ip_network(n, strict=False))
    except ValueError:
        pass

def saddr(l):
    s = l.strip()
    return s.startswith("ip saddr {") or s.startswith("ip6 saddr {")

changed = False
for fam, kw in ((4, "ip saddr {"), (6, "ip6 saddr {")):
    add = [str(n) for n in want if n.version == fam]
    if not add:
        continue
    same = [i for i, l in enumerate(lines) if l.strip().startswith(kw)]
    if same:
        i = same[0]
        l = lines[i]
        a, b = l.index("{") + 1, l.index("}")
        have = [x.strip() for x in l[a:b].split(",") if x.strip()]
        new = [x for x in add if x not in have]
        if new:
            lines[i] = l[:a] + " " + ", ".join(have + new) + " " + l[b:]
            changed = True
        continue
    other = [i for i, l in enumerate(lines) if saddr(l)]
    if not other:
        continue
    l = lines[other[0]]
    indent = l[:len(l) - len(l.lstrip())]
    lines.insert(other[0] + 1, indent + kw + " " + ", ".join(add) + " }" + l[l.index("}") + 1:])
    changed = True

if changed:
    fd, tmp = tempfile.mkstemp(dir=os.path.dirname(RULES))
    with os.fdopen(fd, "w") as f:
        f.write("\n".join(lines))
    if run("nft", "-c", "-f", tmp).returncode != 0:
        os.unlink(tmp)
        print("the settings' networks make a ruleset nft refuses; left as it was", file=sys.stderr)
        sys.exit(1)
    os.chmod(tmp, 0o644)
    os.replace(tmp, RULES)
if changed or not present:
    run("systemctl", "enable", "` + Unit + `")
    if run("nft", "-f", RULES).returncode != 0:
        sys.exit(1)
    print("llmgw's firewall is on here")
`

// FollowerPath is where FollowScript is installed.
const FollowerPath = "/usr/local/lib/llmgw/firewall-follow"

// FollowerTimer is the timer that runs it.
const FollowerTimer = "llmgw-firewall-follow.timer"

const followService = `[Unit]
Description=llmgw firewall: follow the cluster's settings on this store member
After=postgresql.service patroni.service

[Service]
Type=oneshot
ExecStart=/usr/bin/python3 ` + FollowerPath + `
`

const followTimer = `[Unit]
Description=llmgw firewall: follow the cluster's settings on this store member

[Timer]
OnBootSec=1min
OnUnitActiveSec=1min

[Install]
WantedBy=timers.target
`

// InstallFollower puts FollowScript and its timer on a store member.
func InstallFollower(ctx context.Context, r Runner) error {
	body := FollowScript + "\x00" + followService + "\x00" + followTimer
	out, err := r.Run(ctx, `set -e
t=$(mktemp -d); trap 'rm -rf "$t"' EXIT
cat > "$t/all"
python3 - "$t" <<'PY'
import sys
d = sys.argv[1]
s, svc, tim = open(d + "/all", "rb").read().split(b"\x00")
open(d + "/script", "wb").write(s)
open(d + "/service", "wb").write(svc)
open(d + "/timer", "wb").write(tim)
PY
sudo install -d -m 755 /usr/local/lib/llmgw
sudo install -m 755 "$t/script" `+FollowerPath+`
sudo install -m 644 "$t/service" /etc/systemd/system/llmgw-firewall-follow.service
sudo install -m 644 "$t/timer" /etc/systemd/system/`+FollowerTimer+`
sudo systemctl daemon-reload
sudo systemctl enable --now `+FollowerTimer+` >/dev/null 2>&1`, []byte(body))
	if err != nil {
		return fmt.Errorf("installing %s: %v: %s", FollowerTimer, err, strings.TrimSpace(string(out)))
	}
	return nil
}