lpra/internal/server/dbpg.go

package server

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"strconv"
	"strings"
	"time"

	_ "github.com/jackc/pgx/v5/stdlib"
)

// The panel keeps its data in SQLite by default, or in PostgreSQL when the
// admins choose so in Настройки. The choice lives in a small file beside the
// SQLite database (<db>.backend.json, readable by the server only) and
// takes effect on the next start; the first start on an empty PostgreSQL
// copies everything over from the SQLite file, which stays as it was.

type PGConf struct {
	Host     string `json:"host"`
	Port     int    `json:"port"`
	Name     string `json:"name"`
	User     string `json:"user"`
	Password string `json:"password"`
	SSLMode  string `json:"sslmode"`
}

type backendConf struct {
	Type string  `json:"type"` // sqlite or postgres
	PG   *PGConf `json:"pg,omitempty"`
}

func (c PGConf) DSN() string {
	port := c.Port
	if port == 0 {
		port = 5432
	}
	ssl := c.SSLMode
	if ssl == "" {
		ssl = "prefer"
	}
	u := url.URL{Scheme: "postgres", User: url.UserPassword(c.User, c.Password), Host: c.Host + ":" + strconv.Itoa(port), Path: "/" + c.Name}
	u.RawQuery = url.Values{"sslmode": {ssl}, "connect_timeout": {"10"}}.Encode()
	return u.String()
}

func backendPath(dbPath string) string { return dbPath + ".backend.json" }

func readBackend(dbPath string) backendConf {
	var c backendConf
	if b, err := os.ReadFile(backendPath(dbPath)); err == nil {
		_ = json.Unmarshal(b, &c)
	}
	if c.Type != "postgres" || c.PG == nil {
		c = backendConf{Type: "sqlite"}
	}
	return c
}

func writeBackend(dbPath string, c backendConf) error {
	b, _ := json.MarshalIndent(c, "", "  ")
	return os.WriteFile(backendPath(dbPath), b, 0o600)
}

// OpenStore opens the database the admins chose: SQLite at path, or the
// PostgreSQL named beside it.
func OpenStore(path string) (*DB, error) {
	c := readBackend(path)
	if c.Type == "postgres" {
		return OpenPostgres(*c.PG, path)
	}
	return OpenDB(path)
}

// pingPG checks that PostgreSQL c answers.
func pingPG(c PGConf) error {
	d, err := sql.Open("pgx", c.DSN())
	if err != nil {
		return err
	}
	defer d.Close()
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()
	return d.PingContext(ctx)
}

// OpenPostgres opens PostgreSQL c, brings its tables up to date and, when it
// has no users yet, copies over the SQLite database at sqlitePath.
func OpenPostgres(c PGConf, sqlitePath string) (*DB, error) {
	d, err := sql.Open("pgx", c.DSN())
	if err != nil {
		return nil, err
	}
	d.SetMaxOpenConns(16)
	d.SetMaxIdleConns(4)
	d.SetConnMaxIdleTime(5 * time.Minute)
	if err := d.Ping(); err != nil {
		return nil, fmt.Errorf("PostgreSQL: %w", err)
	}
	db := &DB{sql: &dbx{DB: d, pg: true}, path: sqlitePath}
	if err := db.migrate(); err != nil {
		return nil, fmt.Errorf("PostgreSQL: %w", err)
	}
	var n int
	_ = db.sql.QueryRow(`SELECT COUNT(*) FROM admins`).Scan(&n)
	if _, err := os.Stat(sqlitePath); n == 0 && err == nil {
		src, err := OpenDB(sqlitePath)
		if err != nil {
			return nil, err
		}
		defer src.Close()
		if err := copyDB(src, db); err != nil {
			return nil, fmt.Errorf("перенос из SQLite: %w", err)
		}
		log.Printf("db: data copied from %s to PostgreSQL", sqlitePath)
	}
	return db, nil
}

// copyTables are copied parents first, for the foreign keys.
var copyTables = []string{"config", "admins", "agents", "admin_sessions", "audit", "ufw_rules", "admin_keys",
	"forward_rules", "admin_agents", "ssh_hosts", "ai_messages"}

// copyDB copies every row of src into dst (an empty PostgreSQL), in one
// transaction, then moves the id counters past the copied ids.
func copyDB(src, dst *DB) error {
	tx, err := dst.sql.DB.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()
	for _, t := range copyTables {
		rows, err := src.sql.Query(`SELECT * FROM ` + t)
		if err != nil {
			return err
		}
		cols, _ := rows.Columns()
		ph := make([]string, len(cols))
		for i := range ph {
			ph[i] = "$" + strconv.Itoa(i+1)
		}
		ins := `INSERT INTO ` + t + ` (` + strings.Join(cols, ", ") + `) VALUES (` + strings.Join(ph, ", ") + `)`
		for rows.Next() {
			vals := make([]any, len(cols))
			ptrs := make([]any, len(cols))
			for i := range vals {
				ptrs[i] = &vals[i]
			}
			if err := rows.Scan(ptrs...); err != nil {
				rows.Close()
				return err
			}
			if _, err := tx.Exec(ins, vals...); err != nil {
				rows.Close()
				return fmt.Errorf("%s: %w", t, err)
			}
		}
		rows.Close()
		if err := rows.Err(); err != nil {
			return err
		}
	}
	for _, t := range copyTables {
		var seq sql.NullString
		_ = tx.QueryRow(`SELECT pg_get_serial_sequence($1, 'id')`, t).Scan(&seq)
		if seq.Valid {
			if _, err := tx.Exec(`SELECT setval($1, COALESCE((SELECT MAX(id) FROM `+t+`), 0) + 1, false)`, seq.String); err != nil {
				return err
			}
		}
	}
	return tx.Commit()
}

// Backend tells which database is in use and which is chosen for the next
// start.
func (db *DB) Backend() (now string, chosen backendConf) {
	now = "sqlite"
	if db.sql.pg {
		now = "postgres"
	}
	return now, readBackend(db.path)
}

var errNoPGHost = errors.New("укажите хост, базу и пользователя PostgreSQL")

// POST /settings/db - the database for the next start; PostgreSQL is tried
// first and kept only if it answers.
func (s *Server) handleSetDB(w http.ResponseWriter, r *http.Request) {
	if !checkCSRF(r) {
		http.Error(w, "bad csrf", http.StatusForbidden)
		return
	}
	back := func(q string) { http.Redirect(w, r, "/settings?"+q+"#db", http.StatusSeeOther) }
	if r.FormValue("db_type") != "postgres" {
		if err := writeBackend(s.db.path, backendConf{Type: "sqlite"}); err != nil {
			back("err=" + url.QueryEscape(err.Error()))
			return
		}
		back("db=saved")
		return
	}
	old := readBackend(s.db.path).PG
	c := PGConf{Host: strings.TrimSpace(r.FormValue("pg_host")), Name: strings.TrimSpace(r.FormValue("pg_name")),
		User: strings.TrimSpace(r.FormValue("pg_user")), Password: r.FormValue("pg_password"), SSLMode: r.FormValue("pg_sslmode")}
	c.Port, _ = strconv.Atoi(r.FormValue("pg_port"))
	if c.Password == "" && old != nil {
		c.Password = old.Password
	}
	if c.Host == "" || c.Name == "" || c.User == "" {
		back("err=" + url.QueryEscape(errNoPGHost.Error()))
		return
	}
	if err := pingPG(c); err != nil {
		back("err=" + url.QueryEscape("PostgreSQL не отвечает: "+err.Error()))
		return
	}
	if err := writeBackend(s.db.path, backendConf{Type: "postgres", PG: &c}); err != nil {
		back("err=" + url.QueryEscape(err.Error()))
		return
	}
	back("db=saved")
}

// POST /settings/db/restart - the server ends; its service (systemd,
// Restart=always) starts it again on the database chosen.
func (s *Server) handleRestart(w http.ResponseWriter, r *http.Request) {
	if !checkCSRF(r) {
		http.Error(w, "bad csrf", http.StatusForbidden)
		return
	}
	log.Printf("restart asked from the panel")
	http.Redirect(w, r, "/settings?restarted=1#db", http.StatusSeeOther)
	go func() {
		time.Sleep(700 * time.Millisecond)
		os.Exit(0)
	}()
}