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"

	_ "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
}

// 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"`
}

// 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)
	}
	if err := os.Chmod(path, 0o600); err != nil {
		db.Close()
		return nil, fmt.Errorf("secure db: %w", err)
	}
	return &Store{db: db}, 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 FROM locals ORDER BY name")
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []Local
	for rows.Next() {
		var l Local
		if err := rows.Scan(&l.Name, &l.IP, &l.Port, &l.Protocol); err != nil {
			return nil, err
		}
		out = append(out, l)
	}
	return out, rows.Err()
}

func (s *Store) GetLocal(name string) (Local, bool) {
	var l Local
	row := s.db.QueryRow("SELECT name, ip, port, protocol FROM locals WHERE name = ?", name)
	if err := row.Scan(&l.Name, &l.IP, &l.Port, &l.Protocol); err != nil {
		return Local{}, false
	}
	return l, true
}

func (s *Store) UpsertLocal(l Local) error {
	_, err := s.db.Exec(
		"INSERT INTO locals(name, ip, port, protocol) VALUES(?,?,?,?) "+
			"ON CONFLICT(name) DO UPDATE SET ip=excluded.ip, port=excluded.port, protocol=excluded.protocol",
		l.Name, l.IP, l.Port, l.Protocol,
	)
	return err
}

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 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 int
		if err := rows.Scan(&r.Name, &r.IP, &r.Port, &r.Token, &r.URL, &en); err != nil {
			return nil, err
		}
		r.Enabled = en != 0
		out = append(out, r)
	}
	return out, rows.Err()
}

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

func (s *Store) UpsertRemote(r Remote) error {
	_, err := s.db.Exec(
		"INSERT INTO remotes(name, ip, port, token, url, enabled) VALUES(?,?,?,?,?,?) "+
			"ON CONFLICT(name) DO UPDATE SET ip=excluded.ip, port=excluded.port, token=excluded.token, url=excluded.url, enabled=excluded.enabled",
		r.Name, r.IP, r.Port, r.Token, r.URL, boolToInt(r.Enabled),
	)
	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
}