webui4frpc/internal/store/store.go

// Package store implements the persistence layer for webui-frpc.
package store

import (
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	_ "modernc.org/sqlite"
)

// Local is a local forward service node on the canvas.
type Local struct {
	Name     string `json:"name"`
	IP       string `json:"ip"`
	Port     int    `json:"port"`
	Protocol string `json:"protocol"` // tcp | udp | http | https

	// Advanced transport knobs (M1). All optional; empty/zero = frpc defaults.
	UseEncryption  bool              `json:"useEncryption,omitempty"`
	UseCompression bool              `json:"useCompression,omitempty"`
	BandwidthLimit string            `json:"bandwidthLimit,omitempty"` // e.g. "1MB", "2KB"
	PoolCount      int               `json:"poolCount,omitempty"`
	Metadatas      map[string]string `json:"metadatas,omitempty"`
	Annotations    map[string]string `json:"annotations,omitempty"`
}

// Remote is a remote server node on the canvas.
type Remote struct {
	Name    string `json:"name"`
	IP      string `json:"ip"`
	Port    int    `json:"port"` // port to connect to frps
	Token   string `json:"token,omitempty"`
	URL     string `json:"url,omitempty"`
	Enabled bool   `json:"enabled"`

	// Transport section (M1). Protocol: tcp | quic | kcp | websocket.
	TransportProtocol      string `json:"transportProtocol,omitempty"`
	TransportTLS           bool   `json:"transportTls,omitempty"`
	TransportPool          int    `json:"transportPool,omitempty"`
	TransportTLSServerName string `json:"transportTlsServerName,omitempty"`
}

// Link connects one local to one remote.
type Link struct {
	ID         int64  `json:"id,omitempty"`
	Local      string `json:"local"`
	Remote     string `json:"remote"`
	RemotePort int    `json:"remotePort"`
	OffsetX    int    `json:"offsetX,omitempty"`
	OffsetY    int    `json:"offsetY,omitempty"`
}

// Settings holds runtime options.
type Settings struct {
	AutoStartProfiles      bool   `json:"autoStartProfiles"`
	RestartOnExit          bool   `json:"restartOnExit"`
	RestartIntervalSeconds int    `json:"restartIntervalSeconds"`
	BinaryPath             string `json:"binaryPath,omitempty"`
}

// Forward is a rendered link row attached to a remote.
type Forward struct {
	Service    string `json:"service"`
	RemotePort int    `json:"remotePort"`
	LocalPort  int    `json:"localPort,omitempty"`
	OffsetX    int    `json:"offsetX,omitempty"`
	OffsetY    int    `json:"offsetY,omitempty"`
}

const schema = `
CREATE TABLE IF NOT EXISTS locals (
	name     TEXT PRIMARY KEY,
	ip       TEXT NOT NULL,
	port     INTEGER NOT NULL,
	protocol TEXT NOT NULL DEFAULT 'tcp'
);
CREATE TABLE IF NOT EXISTS remotes (
	name     TEXT PRIMARY KEY,
	ip       TEXT NOT NULL,
	port     INTEGER NOT NULL,
	token    TEXT NOT NULL DEFAULT '',
	url      TEXT NOT NULL DEFAULT '',
	enabled  INTEGER NOT NULL DEFAULT 1
);
CREATE TABLE IF NOT EXISTS links (
	id          INTEGER PRIMARY KEY AUTOINCREMENT,
	local       TEXT NOT NULL REFERENCES locals(name) ON DELETE CASCADE,
	remote      TEXT NOT NULL REFERENCES remotes(name) ON DELETE CASCADE,
	remote_port INTEGER NOT NULL DEFAULT 0,
	offset_x    INTEGER NOT NULL DEFAULT 0,
	offset_y    INTEGER NOT NULL DEFAULT 0
);
CREATE INDEX IF NOT EXISTS idx_links_remote ON links(remote);
CREATE TABLE IF NOT EXISTS settings (
	key   TEXT PRIMARY KEY,
	value TEXT NOT NULL
);
`

// Store is the SQLite persistence layer.
type Store struct {
	db *sql.DB
}

// New opens the database at path, creating the schema if needed.
func New(path string) (*Store, error) {
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return nil, fmt.Errorf("create dir: %w", err)
	}
	db, err := sql.Open("sqlite", path)
	if err != nil {
		return nil, fmt.Errorf("open sqlite: %w", err)
	}
	db.SetMaxOpenConns(1)
	if _, err := db.Exec(schema); err != nil {
		db.Close()
		return nil, fmt.Errorf("init schema: %w", err)
	}
	if _, err := db.Exec("PRAGMA foreign_keys = ON"); err != nil {
		db.Close()
		return nil, fmt.Errorf("enable fk: %w", err)
	}
	st := &Store{db: db}
	if err := st.migrate(); err != nil {
		db.Close()
		return nil, fmt.Errorf("migrate schema: %w", err)
	}
	if err := os.Chmod(path, 0o600); err != nil {
		db.Close()
		return nil, fmt.Errorf("secure db: %w", err)
	}
	return st, nil
}

// migrate adds columns introduced after the initial schema so old databases
// keep working (CREATE TABLE IF NOT EXISTS does not touch existing tables).
func (s *Store) migrate() error {
	tables := map[string][]string{
		"locals": {
			"use_encryption INTEGER NOT NULL DEFAULT 0",
			"use_compression INTEGER NOT NULL DEFAULT 0",
			"bandwidth_limit TEXT NOT NULL DEFAULT ''",
			"pool_count INTEGER NOT NULL DEFAULT 0",
			"metadatas TEXT NOT NULL DEFAULT ''",
			"annotations TEXT NOT NULL DEFAULT ''",
		},
		"remotes": {
			"transport_protocol TEXT NOT NULL DEFAULT ''",
			"transport_tls INTEGER NOT NULL DEFAULT 0",
			"transport_pool INTEGER NOT NULL DEFAULT 0",
			"transport_tls_server_name TEXT NOT NULL DEFAULT ''",
		},
	}
	for table, cols := range tables {
		rows, err := s.db.Query("PRAGMA table_info(" + table + ")")
		if err != nil {
			return err
		}
		have := map[string]bool{}
		for rows.Next() {
			var cid int
			var name, typ string
			var notnull, pk int
			var dflt sql.NullString
			if err := rows.Scan(&cid, &name, &typ, &notnull, &dflt, &pk); err != nil {
				rows.Close()
				return err
			}
			have[name] = true
		}
		rows.Close()
		for _, def := range cols {
			name := def[:strings.Index(def, " ")]
			if !have[name] {
				if _, err := s.db.Exec("ALTER TABLE " + table + " ADD COLUMN " + def); err != nil {
					return fmt.Errorf("migrate %s.%s: %w", table, name, err)
				}
			}
		}
	}
	return nil
}

func (s *Store) Close() error { return s.db.Close() }

// ---- Locals ----

func (s *Store) ListLocals() ([]Local, error) {
	rows, err := s.db.Query("SELECT name, ip, port, protocol, use_encryption, use_compression, bandwidth_limit, pool_count, metadatas, annotations FROM locals ORDER BY name")
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []Local
	for rows.Next() {
		var l Local
		var enc, comp int
		var metaRaw, annoRaw string
		if err := rows.Scan(&l.Name, &l.IP, &l.Port, &l.Protocol, &enc, &comp, &l.BandwidthLimit, &l.PoolCount, &metaRaw, &annoRaw); err != nil {
			return nil, err
		}
		l.UseEncryption = enc != 0
		l.UseCompression = comp != 0
		l.Metadatas = decodeMap(metaRaw)
		l.Annotations = decodeMap(annoRaw)
		out = append(out, l)
	}
	return out, rows.Err()
}

func (s *Store) GetLocal(name string) (Local, bool) {
	var l Local
	var enc, comp int
	var metaRaw, annoRaw string
	row := s.db.QueryRow("SELECT name, ip, port, protocol, use_encryption, use_compression, bandwidth_limit, pool_count, metadatas, annotations FROM locals WHERE name = ?", name)
	if err := row.Scan(&l.Name, &l.IP, &l.Port, &l.Protocol, &enc, &comp, &l.BandwidthLimit, &l.PoolCount, &metaRaw, &annoRaw); err != nil {
		return Local{}, false
	}
	l.UseEncryption = enc != 0
	l.UseCompression = comp != 0
	l.Metadatas = decodeMap(metaRaw)
	l.Annotations = decodeMap(annoRaw)
	return l, true
}

func (s *Store) UpsertLocal(l Local) error {
	_, err := s.db.Exec(
		"INSERT INTO locals(name, ip, port, protocol, use_encryption, use_compression, bandwidth_limit, pool_count, metadatas, annotations) VALUES(?,?,?,?,?,?,?,?,?,?) "+
			"ON CONFLICT(name) DO UPDATE SET ip=excluded.ip, port=excluded.port, protocol=excluded.protocol, use_encryption=excluded.use_encryption, use_compression=excluded.use_compression, bandwidth_limit=excluded.bandwidth_limit, pool_count=excluded.pool_count, metadatas=excluded.metadatas, annotations=excluded.annotations",
		l.Name, l.IP, l.Port, l.Protocol, boolToInt(l.UseEncryption), boolToInt(l.UseCompression), l.BandwidthLimit, l.PoolCount, encodeMap(l.Metadatas), encodeMap(l.Annotations),
	)
	return err
}

// encodeMap serializes an optional JSON map so nil stays an empty string.
func encodeMap(m map[string]string) string {
	if len(m) == 0 {
		return ""
	}
	b, err := json.Marshal(m)
	if err != nil {
		return ""
	}
	return string(b)
}

// decodeMap reads a JSON map column; empty/invalid becomes nil.
func decodeMap(raw string) map[string]string {
	if raw == "" {
		return nil
	}
	var m map[string]string
	if err := json.Unmarshal([]byte(raw), &m); err != nil {
		return nil
	}
	return m
}

func (s *Store) DeleteLocal(name string) error {
	_, err := s.db.Exec("DELETE FROM locals WHERE name = ?", name)
	return err
}

// ---- Remotes ----

func (s *Store) ListRemotes() ([]Remote, error) {
	rows, err := s.db.Query("SELECT name, ip, port, token, url, enabled, transport_protocol, transport_tls, transport_pool, transport_tls_server_name FROM remotes ORDER BY name")
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []Remote
	for rows.Next() {
		var r Remote
		var en, tls int
		if err := rows.Scan(&r.Name, &r.IP, &r.Port, &r.Token, &r.URL, &en, &r.TransportProtocol, &tls, &r.TransportPool, &r.TransportTLSServerName); err != nil {
			return nil, err
		}
		r.Enabled = en != 0
		r.TransportTLS = tls != 0
		out = append(out, r)
	}
	return out, rows.Err()
}

func (s *Store) GetRemote(name string) (Remote, bool) {
	var r Remote
	var en, tls int
	row := s.db.QueryRow("SELECT name, ip, port, token, url, enabled, transport_protocol, transport_tls, transport_pool, transport_tls_server_name FROM remotes WHERE name = ?", name)
	if err := row.Scan(&r.Name, &r.IP, &r.Port, &r.Token, &r.URL, &en, &r.TransportProtocol, &tls, &r.TransportPool, &r.TransportTLSServerName); err != nil {
		return Remote{}, false
	}
	r.Enabled = en != 0
	r.TransportTLS = tls != 0
	return r, true
}

func (s *Store) UpsertRemote(r Remote) error {
	_, err := s.db.Exec(
		"INSERT INTO remotes(name, ip, port, token, url, enabled, transport_protocol, transport_tls, transport_pool, transport_tls_server_name) VALUES(?,?,?,?,?,?,?,?,?,?) "+
			"ON CONFLICT(name) DO UPDATE SET ip=excluded.ip, port=excluded.port, token=excluded.token, url=excluded.url, enabled=excluded.enabled, transport_protocol=excluded.transport_protocol, transport_tls=excluded.transport_tls, transport_pool=excluded.transport_pool, transport_tls_server_name=excluded.transport_tls_server_name",
		r.Name, r.IP, r.Port, r.Token, r.URL, boolToInt(r.Enabled), r.TransportProtocol, boolToInt(r.TransportTLS), r.TransportPool, r.TransportTLSServerName,
	)
	return err
}

func (s *Store) DeleteRemote(name string) error {
	_, err := s.db.Exec("DELETE FROM remotes WHERE name = ?", name)
	return err
}

// ---- Links ----

func (s *Store) ListLinks() ([]Link, error) {
	rows, err := s.db.Query("SELECT id, local, remote, remote_port, offset_x, offset_y FROM links")
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []Link
	for rows.Next() {
		var l Link
		if err := rows.Scan(&l.ID, &l.Local, &l.Remote, &l.RemotePort, &l.OffsetX, &l.OffsetY); err != nil {
			return nil, err
		}
		out = append(out, l)
	}
	return out, rows.Err()
}

// LocalTarget describes one outgoing forward of a local service.
type LocalTarget struct {
	Remote     string `json:"remote"`
	RemotePort int    `json:"remotePort"`
}

// LinksForLocal returns the targets a local service forwards to.
func (s *Store) LinksForLocal(local string) ([]LocalTarget, error) {
	links, err := s.ListLinks()
	if err != nil {
		return nil, err
	}
	var out []LocalTarget
	for _, l := range links {
		if l.Local != local {
			continue
		}
		out = append(out, LocalTarget{Remote: l.Remote, RemotePort: l.RemotePort})
	}
	return out, nil
}

// LinksForRemote returns forwards of one remote with local port resolved.
func (s *Store) LinksForRemote(remote string) ([]Forward, error) {
	links, err := s.ListLinks()
	if err != nil {
		return nil, err
	}
	var out []Forward
	for _, l := range links {
		if l.Remote != remote {
			continue
		}
		loc, ok := s.GetLocal(l.Local)
		if !ok {
			continue
		}
		out = append(out, Forward{
			Service:    l.Local,
			RemotePort: l.RemotePort,
			LocalPort:  loc.Port,
			OffsetX:    l.OffsetX,
			OffsetY:    l.OffsetY,
		})
	}
	return out, nil
}

// ReplaceLinks clears all links and inserts the given set in one transaction.
func (s *Store) ReplaceLinks(links []Link) error {
	tx, err := s.db.Begin()
	if err != nil {
		return err
	}
	defer func() { _ = tx.Rollback() }()
	if _, err := tx.Exec("DELETE FROM links"); err != nil {
		return err
	}
	for _, l := range links {
		if _, err := tx.Exec(
			"INSERT INTO links(local, remote, remote_port, offset_x, offset_y) VALUES(?,?,?,?,?)",
			l.Local, l.Remote, l.RemotePort, l.OffsetX, l.OffsetY,
		); err != nil {
			return err
		}
	}
	return tx.Commit()
}

// ---- Settings ----

func (s *Store) Settings() (Settings, error) {
	var st Settings
	row := s.db.QueryRow("SELECT value FROM settings WHERE key = 'settings'")
	var raw string
	if err := row.Scan(&raw); err != nil {
		return Settings{AutoStartProfiles: true, RestartOnExit: true, RestartIntervalSeconds: 5}, nil
	}
	if err := json.Unmarshal([]byte(raw), &st); err != nil {
		return Settings{AutoStartProfiles: true, RestartOnExit: true, RestartIntervalSeconds: 5}, nil
	}
	return st, nil
}

func (s *Store) UpdateSettings(st Settings) error {
	raw, err := json.Marshal(st)
	if err != nil {
		return err
	}
	_, err = s.db.Exec(
		"INSERT INTO settings(key, value) VALUES('settings', ?) ON CONFLICT(key) DO UPDATE SET value=excluded.value",
		string(raw),
	)
	return err
}

var (
	ErrNotFound      = errors.New("not found")
	ErrInvalid       = errors.New("invalid argument")
	ErrAlreadyExists = errors.New("already exists")
)

func boolToInt(b bool) int {
	if b {
		return 1
	}
	return 0
}