Megalodon/internal/manager/instance.go

package manager

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"git.veillax.com/Veillax/megalodon/internal/ark"
	"git.veillax.com/Veillax/megalodon/internal/config"
	"git.veillax.com/Veillax/megalodon/internal/rcon"
)

const (
	sampleEvery = time.Minute
	sampleKeep  = 1440 // a day
)

// Sample is one point of the player-count history.
type Sample struct {
	T  int64 `json:"t"` // unix seconds
	N  int   `json:"n"` // players online
	Up bool  `json:"up"`
}

// PlayerView is a connected player as the UI sees them.
type PlayerView struct {
	ark.Player
	JoinedAt time.Time `json:"joined_at"`
	Watch    bool      `json:"watch,omitempty"`
	Notes    string    `json:"notes,omitempty"`
	Banned   bool      `json:"banned,omitempty"`
}

// RestartPlan is a restart or shutdown that is counting down.
type RestartPlan struct {
	At     time.Time `json:"at"`
	By     string    `json:"by"`
	Reason string    `json:"reason,omitempty"`
	Stop   bool      `json:"stop"` // shut down and stay down, instead of restarting
}

// State is what is currently known about a server.
type State struct {
	Online       bool         `json:"online"`
	Since        *time.Time   `json:"since,omitempty"` // when it came up, as far as Megalodon saw
	OfflineSince *time.Time   `json:"offline_since,omitempty"`
	LastSeen     *time.Time   `json:"last_seen,omitempty"`
	LastError    string       `json:"last_error,omitempty"`
	RTTMS        int64        `json:"rtt_ms"`
	Players      []PlayerView `json:"players"`
	Info         *ark.Info    `json:"info,omitempty"`
	Held         bool         `json:"held"` // stopped on purpose: auto-start stays off
	Restart      *RestartPlan `json:"restart,omitempty"`
	Peak         int          `json:"peak"` // most players seen at once since Megalodon started
}

// View is a server plus its state, as sent to the UI.
type View struct {
	ID        string `json:"id"`
	Name      string `json:"name"`
	Host      string `json:"host"`
	RCONPort  int    `json:"rcon_port"`
	QueryPort int    `json:"query_port"`
	CanStart  bool   `json:"can_start"`
	AutoStart bool   `json:"auto_start"`
	State
}

type strike struct {
	n    int
	last time.Time
}

// Instance watches and controls one server.
type Instance struct {
	cfg config.Server
	m   *Manager
	rc  *rcon.Client

	ctx    context.Context
	cancel context.CancelFunc
	done   chan struct{}

	mu          sync.RWMutex
	st          State
	known       bool // a poll has finished since Megalodon started, so changes are real events
	fails       int
	lastPoll    time.Time
	lastSample  time.Time
	samples     []Sample
	planCancel  context.CancelFunc
	autoStarted time.Time
	strikes     map[string]strike
}

func newInstance(m *Manager, cfg config.Server) *Instance {
	i := &Instance{
		cfg: cfg, m: m, done: make(chan struct{}), strikes: map[string]strike{},
		rc: rcon.New(fmt.Sprintf("%s:%d", cfg.Host, cfg.RCONPort), cfg.Password),
	}
	i.rc.Timeout = 6 * time.Second
	i.loadSamples()
	return i
}

func (i *Instance) start(parent context.Context) {
	i.ctx, i.cancel = context.WithCancel(parent)
	go func() {
		defer close(i.done)
		i.run()
	}()
}

func (i *Instance) stop() {
	i.cancel()
	<-i.done
	i.rc.Close()
}

func (i *Instance) event(typ, actor, msg string) Event {
	return i.m.hub.Publish(Event{Server: i.cfg.ID, Type: typ, Actor: actor, Message: msg})
}

// View returns the server's current state.
func (i *Instance) View() View {
	i.mu.RLock()
	defer i.mu.RUnlock()
	st := i.st
	st.Players = append([]PlayerView(nil), i.st.Players...)
	return View{
		ID: i.cfg.ID, Name: i.cfg.Name, Host: i.cfg.Host, RCONPort: i.cfg.RCONPort, QueryPort: i.cfg.QueryPort,
		CanStart: i.cfg.StartCommand != "", AutoStart: i.cfg.AutoStart, State: st,
	}
}

// Samples returns the player-count history, oldest first.
func (i *Instance) Samples() []Sample {
	i.mu.RLock()
	defer i.mu.RUnlock()
	return append([]Sample(nil), i.samples...)
}

// Exec sends one console command.
func (i *Instance) Exec(cmd string) (string, error) { return i.rc.Exec(cmd) }

func (i *Instance) run() {
	cfg := i.m.cfg.Get()
	poll := time.NewTicker(time.Duration(cfg.PollSeconds) * time.Second)
	chat := time.NewTicker(time.Duration(cfg.ChatSeconds) * time.Second)
	defer poll.Stop()
	defer chat.Stop()
	i.poll()
	for {
		select {
		case <-i.ctx.Done():
			return
		case <-poll.C:
			i.poll()
		case <-chat.C:
			i.pollChat()
		}
	}
}

// ---- polling ----

func (i *Instance) poll() {
	start := time.Now()
	out, err := i.rc.Exec(ark.ListPlayers)
	rtt := time.Since(start)
	if err != nil {
		i.markFailed(err)
		return
	}
	var info *ark.Info
	if i.cfg.QueryPort > 0 {
		info, _ = ark.Query(fmt.Sprintf("%s:%d", i.cfg.Host, i.cfg.QueryPort), 2*time.Second)
	}
	i.apply(ark.ParsePlayers(out), info, rtt)
}

func (i *Instance) markFailed(err error) {
	now := time.Now()
	i.mu.Lock()
	i.fails++
	i.st.LastError = err.Error()
	wentDown := i.st.Online && (i.fails >= 2 || errors.Is(err, rcon.ErrAuth))
	if wentDown {
		i.st.Online = false
		i.st.Players, i.st.Info, i.st.Since = nil, nil, nil
	}
	if !i.st.Online && i.st.OfflineSince == nil {
		i.st.OfflineSince = &now
	}
	i.known = true
	i.recordSample(now, 0, false)
	i.mu.Unlock()
	if wentDown {
		i.event("offline", "", "Server went offline: "+err.Error())
	}
	i.maybeAutoStart()
}

func (i *Instance) apply(players []ark.Player, info *ark.Info, rtt time.Duration) {
	now := time.Now()
	i.mu.Lock()
	wasOnline, known := i.st.Online, i.known
	prev := map[string]PlayerView{}
	for _, p := range i.st.Players {
		prev[p.SteamID] = p
	}
	dt := time.Duration(0)
	if wasOnline && !i.lastPoll.IsZero() {
		dt = now.Sub(i.lastPoll)
		if limit := 3 * time.Duration(i.m.cfg.Get().PollSeconds) * time.Second; dt > limit {
			dt = limit
		}
	}
	i.lastPoll, i.fails, i.known = now, 0, true
	i.st.Online, i.st.LastSeen, i.st.LastError, i.st.RTTMS, i.st.Info = true, &now, "", rtt.Milliseconds(), info
	if !wasOnline {
		i.st.Since, i.st.OfflineSince = &now, nil
	}
	views := make([]PlayerView, 0, len(players))
	var joined, left []PlayerView
	cur := map[string]bool{}
	for _, p := range players {
		cur[p.SteamID] = true
		v, was := prev[p.SteamID]
		if !was {
			v = PlayerView{Player: p, JoinedAt: now}
			joined = append(joined, v)
		}
		v.Name = p.Name
		views = append(views, v)
	}
	for id, v := range prev {
		if !cur[id] {
			left = append(left, v)
		}
	}
	if len(views) > i.st.Peak {
		i.st.Peak = len(views)
	}
	i.recordSample(now, len(views), true)
	i.mu.Unlock()

	// Who is on: remember them, and decorate with notes and ban status for the UI.
	isNew := joinedSet(joined)
	for k, v := range views {
		_, joinedNow := isNew[v.SteamID]
		rec := i.m.players.Seen(v.SteamID, v.Name, i.cfg.ID, dt, joinedNow && known)
		_, banned := i.m.bans.Active(v.SteamID, i.cfg.ID)
		views[k].Watch, views[k].Notes, views[k].Banned = rec.Watch, rec.Notes, banned
	}
	i.mu.Lock()
	i.st.Players = views
	i.mu.Unlock()

	if !wasOnline && known {
		i.event("online", "", "Server is online")
	}
	for _, v := range left {
		i.event("leave", "", fmt.Sprintf("%s left after %s", v.Name, ark.Duration(now.Sub(v.JoinedAt))))
	}
	for _, v := range views {
		if b, banned := i.m.bans.Active(v.SteamID, i.cfg.ID); banned {
			i.enforceBan(v, b)
			continue
		}
		if _, fresh := isNew[v.SteamID]; fresh && known {
			i.onJoin(v, len(views))
		}
	}
}

func joinedSet(l []PlayerView) map[string]struct{} {
	m := make(map[string]struct{}, len(l))
	for _, v := range l {
		m[v.SteamID] = struct{}{}
	}
	return m
}

func (i *Instance) onJoin(v PlayerView, online int) {
	i.m.hub.Publish(Event{Server: i.cfg.ID, Type: "join", Message: v.Name + " joined", SteamID: v.SteamID})
	if v.Watch {
		note := ""
		if v.Notes != "" {
			note = " (" + v.Notes + ")"
		}
		i.event("watch", "", "Watched player "+v.Name+" joined"+note)
	}
	if w := i.m.cfg.Get().Welcome; w.Enabled && w.Message != "" {
		msg := strings.NewReplacer("{name}", v.Name, "{server}", i.cfg.Name, "{players}", strconv.Itoa(online)).Replace(w.Message)
		if cmd, err := ark.ServerChatTo(v.SteamID, msg); err == nil {
			// Give the player a moment to finish loading in, or they never see it.
			go func() {
				select {
				case <-time.After(15 * time.Second):
					i.rc.Exec(cmd)
				case <-i.ctx.Done():
				}
			}()
		}
	}
}

func (i *Instance) enforceBan(v PlayerView, b Ban) {
	cmd, err := ark.ForPlayer("BanPlayer", v.SteamID)
	if err != nil {
		return
	}
	if _, err := i.rc.Exec(cmd); err == nil {
		why := b.Reason
		if why == "" {
			why = "banned"
		}
		i.event("kick", "megalodon", fmt.Sprintf("Removed banned player %s (%s)", v.Name, why))
	}
}

// recordSample must be called with i.mu held.
func (i *Instance) recordSample(now time.Time, n int, up bool) {
	if now.Sub(i.lastSample) < sampleEvery {
		return
	}
	i.lastSample = now
	s := Sample{T: now.Unix(), N: n, Up: up}
	i.samples = append(i.samples, s)
	if len(i.samples) > sampleKeep {
		i.samples = i.samples[len(i.samples)-sampleKeep:]
	}
	if f, err := os.OpenFile(i.samplePath(), os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o600); err == nil {
		b, _ := json.Marshal(s)
		f.Write(append(b, '\n'))
		f.Close()
	}
}

func (i *Instance) samplePath() string {
	return filepath.Join(i.m.dir, "samples-"+i.cfg.ID+".jsonl")
}

func (i *Instance) loadSamples() {
	f, err := os.Open(i.samplePath())
	if err != nil {
		return
	}
	defer f.Close()
	cutoff := time.Now().Add(-24 * time.Hour).Unix()
	sc := bufio.NewScanner(f)
	for sc.Scan() {
		var s Sample
		if json.Unmarshal(sc.Bytes(), &s) == nil && s.T >= cutoff {
			i.samples = append(i.samples, s)
		}
	}
	// Rewrite the file so it only ever holds about a day.
	var sb strings.Builder
	for _, s := range i.samples {
		b, _ := json.Marshal(s)
		sb.Write(b)
		sb.WriteByte('\n')
	}
	os.WriteFile(i.samplePath(), []byte(sb.String()), 0o600)
}

// ---- chat ----

func (i *Instance) pollChat() {
	i.mu.RLock()
	online := i.st.Online
	i.mu.RUnlock()
	if !online {
		return
	}
	out, err := i.rc.Exec(ark.GetChat)
	if err != nil {
		return
	}
	for _, cl := range ark.ParseChat(out) {
		who := cl.Sender
		if cl.Server && who == "" {
			who = "SERVER"
		}
		i.m.hub.Publish(Event{Server: i.cfg.ID, Type: "chat", Actor: who, Message: cl.Message, Detail: cl.Character})
		if !cl.Server {
			i.filter(cl)
		}
	}
}

// filter applies the word filter to one line of player chat.
func (i *Instance) filter(cl ark.ChatLine) {
	cf := i.m.cfg.Get().ChatFilter
	if !cf.Enabled {
		return
	}
	low := strings.ToLower(cl.Message)
	hit := false
	for _, w := range cf.Words {
		if w = strings.ToLower(strings.TrimSpace(w)); w != "" && strings.Contains(low, w) {
			hit = true
			break
		}
	}
	if !hit {
		return
	}
	i.mu.Lock()
	var who *PlayerView
	for k := range i.st.Players {
		if strings.EqualFold(i.st.Players[k].Name, cl.Sender) {
			who = &i.st.Players[k]
			break
		}
	}
	if who == nil {
		i.mu.Unlock()
		return
	}
	p := *who
	s := i.strikes[p.SteamID]
	if time.Since(s.last) > time.Hour {
		s.n = 0
	}
	s.n++
	s.last = time.Now()
	i.strikes[p.SteamID] = s
	over := s.n >= cf.Strikes
	if over {
		delete(i.strikes, p.SteamID)
	}
	i.mu.Unlock()

	switch {
	case over && cf.Action == "kick":
		i.run1("KickPlayer", p.SteamID)
		i.event("filter", "megalodon", fmt.Sprintf("Kicked %s for language (%d strikes)", p.Name, cf.Strikes))
	case over && cf.Action == "ban":
		i.m.Ban("megalodon", p.SteamID, p.Name, "chat filter", time.Duration(cf.BanMinutes)*time.Minute, []string{i.cfg.ID})
		i.event("filter", "megalodon", fmt.Sprintf("Banned %s for %d minutes for language", p.Name, cf.BanMinutes))
	default:
		msg := fmt.Sprintf("Please keep chat friendly (warning %d of %d).", s.n, cf.Strikes)
		if over {
			msg = "Please keep chat friendly."
		}
		if cmd, err := ark.ServerChatTo(p.SteamID, msg); err == nil {
			i.rc.Exec(cmd)
		}
		i.event("filter", "megalodon", fmt.Sprintf("Warned %s (%d of %d)", p.Name, s.n, cf.Strikes))
	}
}

func (i *Instance) run1(cmd, steamID string) {
	if c, err := ark.ForPlayer(cmd, steamID); err == nil {
		i.rc.Exec(c)
	}
}

// ---- actions ----

// Perform runs one action on this server and returns the server's reply.
func (i *Instance) Perform(actor string, a Action) (string, error) {
	var cmd string
	var err error
	switch a.Name {
	case "save":
		cmd = ark.SaveWorld
	case "broadcast":
		cmd, err = ark.Broadcast(a.Message)
	case "say":
		cmd, err = ark.ServerChat(a.Message)
	case "whisper":
		cmd, err = ark.ServerChatTo(a.SteamID, a.Message)
	case "kick":
		cmd, err = ark.ForPlayer("KickPlayer", a.SteamID)
	case "allow":
		cmd, err = ark.ForPlayer("AllowPlayerToJoinNoCheck", a.SteamID)
	case "disallow":
		cmd, err = ark.ForPlayer("DisallowPlayerToJoinNoCheck", a.SteamID)
	case "destroy_wild_dinos":
		cmd = ark.DestroyWildDinos
	case "set_time":
		cmd, err = ark.SetTimeOfDay(a.Time)
	case "motd":
		cmd, err = ark.SetMOTD(a.Message)
	case "exec":
		if cmd = strings.TrimSpace(a.Command); cmd == "" {
			err = errors.New("no command given")
		}
	case "restart", "stop":
		err = i.ScheduleRestart(actor, time.Duration(a.Minutes*float64(time.Minute)), a.Reason, a.Name == "stop")
		if err == nil {
			return "scheduled", nil
		}
	case "cancel_restart":
		err = i.CancelRestart(actor)
		if err == nil {
			return "cancelled", nil
		}
	case "start":
		return i.Start(actor)
	default:
		err = fmt.Errorf("unknown action %q", a.Name)
	}
	if err != nil {
		return "", err
	}
	out, err := i.rc.Exec(cmd)
	if err != nil {
		return "", err
	}
	i.audit(actor, a)
	if ark.IsEmptyReply(out) {
		return "", nil
	}
	return strings.TrimSpace(out), nil
}

func (i *Instance) audit(actor string, a Action) {
	detail := a.Name
	switch a.Name {
	case "exec":
		detail = "console: " + a.Command
	case "broadcast", "say", "motd":
		detail = a.Name + ": " + a.Message
	case "whisper":
		detail = "whisper to " + a.SteamID + ": " + a.Message
	case "kick", "allow", "disallow":
		detail = a.Name + " " + a.SteamID
	case "set_time":
		detail = "set time of day to " + a.Time
	}
	i.event("audit", actor, detail)
}

// ScheduleRestart announces a restart (or, with stop, a shutdown) and carries it out after delay.
// Players are warned along the way; the world is saved first.
func (i *Instance) ScheduleRestart(by string, delay time.Duration, reason string, stop bool) error {
	if delay < 0 || delay > 24*time.Hour {
		return errors.New("the delay must be between 0 and 24 hours")
	}
	i.mu.Lock()
	if i.st.Restart != nil {
		i.mu.Unlock()
		return errors.New("a restart is already scheduled; cancel it first")
	}
	at := time.Now().Add(delay)
	plan := &RestartPlan{At: at, By: by, Reason: ark.Clean(reason), Stop: stop}
	ctx, cancel := context.WithCancel(i.ctx)
	i.st.Restart, i.planCancel = plan, cancel
	i.mu.Unlock()
	verb := "restart"
	if stop {
		verb = "shutdown"
	}
	i.event("restart", by, fmt.Sprintf("Scheduled %s in %s%s", verb, ark.Duration(delay), withReason(plan.Reason)))
	go i.runPlan(ctx, plan)
	return nil
}

func withReason(r string) string {
	if r == "" {
		return ""
	}
	return " (" + r + ")"
}

// warnings are the moments before a restart at which players are told.
var warnings = []time.Duration{30 * time.Minute, 15 * time.Minute, 10 * time.Minute, 5 * time.Minute, 2 * time.Minute, time.Minute, 30 * time.Second, 10 * time.Second}

func (i *Instance) runPlan(ctx context.Context, plan *RestartPlan) {
	finish := func() {
		i.mu.Lock()
		if i.st.Restart == plan {
			i.st.Restart, i.planCancel = nil, nil
		}
		i.mu.Unlock()
	}
	defer finish()
	verb, ing := "restarting", "restart"
	if plan.Stop {
		verb, ing = "shutting down", "shutdown"
	}
	say := func(until time.Duration) {
		msg := fmt.Sprintf("Server %s in %s.", verb, ark.Duration(until))
		if plan.Reason != "" {
			msg += " " + plan.Reason
		}
		if cmd, err := ark.Broadcast(msg); err == nil {
			i.rc.Exec(cmd)
		}
	}
	wait := func(until time.Time) bool {
		t := time.NewTimer(time.Until(until))
		defer t.Stop()
		select {
		case <-t.C:
			return true
		case <-ctx.Done():
			return false
		}
	}
	total := time.Until(plan.At)
	if total >= 15*time.Second {
		say(total) // tell everyone straight away, not at the next fixed warning
	}
	for _, w := range warnings {
		if w >= total-5*time.Second { // already announced above, or too late to matter
			continue
		}
		if !wait(plan.At.Add(-w)) {
			return
		}
		say(w)
	}
	if !wait(plan.At) {
		return
	}
	i.event("restart", plan.By, "Saving and "+ing+"ing now")
	i.rc.Exec(ark.SaveWorld)
	time.Sleep(2 * time.Second)
	i.rc.Exec(ark.DoExit) // the connection usually drops before a reply
	if plan.Stop {
		i.mu.Lock()
		i.st.Held = true
		i.mu.Unlock()
		finish()
		return
	}
	if i.cfg.StartCommand == "" {
		return // a supervisor (docker, systemd) is expected to bring it back
	}
	// Wait for the server to be seen going down, then start it.
	deadline := time.Now().Add(3 * time.Minute)
	for time.Now().Before(deadline) {
		select {
		case <-ctx.Done():
			return
		case <-time.After(time.Second):
		}
		i.mu.RLock()
		online := i.st.Online
		i.mu.RUnlock()
		if !online {
			break
		}
	}
	select {
	case <-ctx.Done():
	case <-time.After(5 * time.Second):
		i.Start(plan.By)
	}
}

// CancelRestart stops a countdown.
func (i *Instance) CancelRestart(by string) error {
	i.mu.Lock()
	defer i.mu.Unlock()
	if i.st.Restart == nil {
		return errors.New("no restart is scheduled")
	}
	i.planCancel()
	i.st.Restart, i.planCancel = nil, nil
	go func() {
		if cmd, err := ark.Broadcast("The scheduled restart was cancelled."); err == nil {
			i.rc.Exec(cmd)
		}
	}()
	i.m.hub.Publish(Event{Server: i.cfg.ID, Type: "restart", Actor: by, Message: "Scheduled restart cancelled"})
	return nil
}

// Start runs the server's start command.
func (i *Instance) Start(by string) (string, error) {
	if i.cfg.StartCommand == "" {
		return "", errors.New("this server has no start command configured")
	}
	i.mu.Lock()
	i.st.Held = false
	i.autoStarted = time.Now()
	i.mu.Unlock()
	out, err := runShell(i.cfg.StartCommand)
	if err != nil {
		i.event("error", by, "Start command failed: "+err.Error()+" "+out)
		return out, err
	}
	i.event("start", by, "Ran the start command")
	return out, nil
}

func (i *Instance) maybeAutoStart() {
	if !i.cfg.AutoStart || i.cfg.StartCommand == "" {
		return
	}
	after := time.Duration(i.cfg.AutoStartSec) * time.Second
	if after <= 0 {
		after = 2 * time.Minute
	}
	i.mu.RLock()
	due := !i.st.Online && !i.st.Held && i.st.OfflineSince != nil && time.Since(*i.st.OfflineSince) >= after && time.Since(i.autoStarted) >= 5*time.Minute
	i.mu.RUnlock()
	if due {
		i.Start("auto-start")
	}
}

func runShell(cmd string) (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
	defer cancel()
	var c *exec.Cmd
	if runtime.GOOS == "windows" {
		c = exec.CommandContext(ctx, "cmd", "/C", cmd)
	} else {
		c = exec.CommandContext(ctx, "sh", "-c", cmd)
	}
	out, err := c.CombinedOutput()
	s := strings.TrimSpace(string(out))
	if len(s) > 500 {
		s = s[:500]
	}
	return s, err
}