elmprodvpn/selective-vpn-api/app/trafficappmarks/store.go

package trafficappmarks

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strings"
	"time"
)

type State struct {
	Version   int    `json:"version"`
	UpdatedAt string `json:"updated_at"`
	Items     []Item `json:"items,omitempty"`
}

type Item struct {
	ID        uint64 `json:"id"`
	Target    string `json:"target"` // vpn|direct
	Cgroup    string `json:"cgroup"` // absolute path ("/user.slice/..."), informational
	CgroupRel string `json:"cgroup_rel"`
	Level     int    `json:"level"`
	Unit      string `json:"unit,omitempty"`
	Command   string `json:"command,omitempty"`
	AppKey    string `json:"app_key,omitempty"`
	AddedAt   string `json:"added_at"`
	ExpiresAt string `json:"expires_at"`
}

func LoadState(statePath string, canonicalizeAppKey func(appKey, command string) string) State {
	st := State{Version: 1}
	data, err := os.ReadFile(statePath)
	if err != nil {
		return st
	}
	if err := json.Unmarshal(data, &st); err != nil {
		return State{Version: 1}
	}
	if st.Version == 0 {
		st.Version = 1
	}

	changed := false
	for i := range st.Items {
		st.Items[i].Target = strings.ToLower(strings.TrimSpace(st.Items[i].Target))
		if canonicalizeAppKey != nil {
			canon := canonicalizeAppKey(st.Items[i].AppKey, st.Items[i].Command)
			if canon != "" && strings.TrimSpace(st.Items[i].AppKey) != canon {
				st.Items[i].AppKey = canon
				changed = true
			}
		}
	}
	if deduped, dedupChanged := DedupeItems(st.Items, canonicalizeAppKey); dedupChanged {
		st.Items = deduped
		changed = true
	}
	if changed {
		_ = SaveState(statePath, st)
	}
	return st
}

func SaveState(statePath string, st State) error {
	st.Version = 1
	st.UpdatedAt = time.Now().UTC().Format(time.RFC3339)

	data, err := json.MarshalIndent(st, "", "  ")
	if err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(statePath), 0o755); err != nil {
		return err
	}
	tmp := statePath + ".tmp"
	if err := os.WriteFile(tmp, data, 0o644); err != nil {
		return err
	}
	return os.Rename(tmp, statePath)
}

func DedupeItems(in []Item, canonicalizeAppKey func(appKey, command string) string) ([]Item, bool) {
	if len(in) <= 1 {
		return in, false
	}
	out := make([]Item, 0, len(in))
	byTargetID := map[string]int{}
	byTargetApp := map[string]int{}
	changed := false

	for _, raw := range in {
		it := raw
		it.Target = strings.ToLower(strings.TrimSpace(it.Target))
		if it.Target != "vpn" && it.Target != "direct" {
			changed = true
			continue
		}
		if canonicalizeAppKey != nil {
			it.AppKey = canonicalizeAppKey(it.AppKey, it.Command)
		}

		if it.ID > 0 {
			idKey := fmt.Sprintf("%s:%d", it.Target, it.ID)
			if idx, ok := byTargetID[idKey]; ok {
				if preferItem(it, out[idx]) {
					out[idx] = it
				}
				changed = true
				continue
			}
			byTargetID[idKey] = len(out)
		}

		if it.AppKey != "" {
			appKey := it.Target + "|" + it.AppKey
			if idx, ok := byTargetApp[appKey]; ok {
				if preferItem(it, out[idx]) {
					out[idx] = it
				}
				changed = true
				continue
			}
			byTargetApp[appKey] = len(out)
		}

		out = append(out, it)
	}
	return out, changed
}

func UpsertItem(items []Item, next Item) []Item {
	out := items[:0]
	for _, it := range items {
		if strings.ToLower(strings.TrimSpace(it.Target)) == strings.ToLower(strings.TrimSpace(next.Target)) && it.ID == next.ID {
			continue
		}
		out = append(out, it)
	}
	out = append(out, next)
	return out
}

func IsAllDigits(s string) bool {
	s = strings.TrimSpace(s)
	if s == "" {
		return false
	}
	for i := 0; i < len(s); i++ {
		ch := s[i]
		if ch < '0' || ch > '9' {
			return false
		}
	}
	return true
}

func PruneExpired(st *State, now time.Time, deleteRule func(target string, id uint64)) (changed bool) {
	if st == nil {
		return false
	}
	kept := st.Items[:0]
	for _, it := range st.Items {
		expRaw := strings.TrimSpace(it.ExpiresAt)
		if expRaw == "" {
			kept = append(kept, it)
			continue
		}
		exp, err := time.Parse(time.RFC3339, expRaw)
		if err != nil {
			it.ExpiresAt = ""
			kept = append(kept, it)
			changed = true
			continue
		}
		if !exp.After(now) {
			if deleteRule != nil {
				deleteRule(strings.ToLower(strings.TrimSpace(it.Target)), it.ID)
			}
			changed = true
			continue
		}
		kept = append(kept, it)
	}
	st.Items = kept
	return changed
}

func preferItem(cand, cur Item) bool {
	ca := strings.TrimSpace(cand.AddedAt)
	oa := strings.TrimSpace(cur.AddedAt)
	if ca != oa {
		if ca == "" {
			return false
		}
		if oa == "" {
			return true
		}
		return ca > oa
	}
	if strings.TrimSpace(cand.Command) != "" && strings.TrimSpace(cur.Command) == "" {
		return true
	}
	return false
}