DragonCoreSSH-NewWEB/xray_clients.go

package main

import (
	"context"
	"database/sql"
	"fmt"
	"log"
	"strings"
	"time"
)

// XrayClientMeta holds metadata stored in PostgreSQL for an Xray client.
// Xray's own config only stores uuid/email/level; expiry, display name,
// reseller owner, and connection policy live here.
type XrayClientMeta struct {
	UUID               string
	Name               string
	Email              string
	InboundTag         string
	OwnerUsername      string
	ExpiresAt          *time.Time
	MaxConns           int
	CreatedAt          time.Time
	TotalUplinkBytes   int64
	TotalDownlinkBytes int64
	LastActive         *time.Time
	ActiveConnections  int
}

func (s *Store) EnsureXrayClientsSchema(ctx context.Context) error {
	stmts := []string{
		`CREATE TABLE IF NOT EXISTS xray_clients (
			uuid                  TEXT PRIMARY KEY,
			name                  TEXT NOT NULL DEFAULT '',
			email                 TEXT NOT NULL DEFAULT '',
			inbound_tag           TEXT NOT NULL DEFAULT '',
			owner_username        TEXT NOT NULL DEFAULT '',
			expires_at            TIMESTAMPTZ,
			max_conns             INT NOT NULL DEFAULT 0,
			created_at            TIMESTAMPTZ NOT NULL DEFAULT NOW(),
			total_uplink_bytes    BIGINT NOT NULL DEFAULT 0,
			total_downlink_bytes  BIGINT NOT NULL DEFAULT 0,
			last_active           TIMESTAMPTZ,
			active_connections    INT NOT NULL DEFAULT 0
		)`,
		`ALTER TABLE xray_clients ADD COLUMN IF NOT EXISTS owner_username TEXT NOT NULL DEFAULT ''`,
		`ALTER TABLE xray_clients ADD COLUMN IF NOT EXISTS total_uplink_bytes BIGINT NOT NULL DEFAULT 0`,
		`ALTER TABLE xray_clients ADD COLUMN IF NOT EXISTS total_downlink_bytes BIGINT NOT NULL DEFAULT 0`,
		`ALTER TABLE xray_clients ADD COLUMN IF NOT EXISTS last_active TIMESTAMPTZ`,
		`ALTER TABLE xray_clients ADD COLUMN IF NOT EXISTS active_connections INT NOT NULL DEFAULT 0`,
		// Keep legacy reseller-owned Xray accounts aligned with weighted quota
		// accounting. A reseller account always consumes at least one slot.
		`UPDATE xray_clients SET max_conns = 1
			WHERE owner_username <> '' AND max_conns < 1`,
	}
	for _, stmt := range stmts {
		if _, err := s.db.ExecContext(ctx, stmt); err != nil {
			return err
		}
	}
	return nil
}

func (s *Store) UpsertXrayClientMeta(ctx context.Context, m XrayClientMeta) error {
	var expiresAt interface{}
	if m.ExpiresAt != nil {
		expiresAt = *m.ExpiresAt
	}
	_, err := s.db.ExecContext(ctx, `
		INSERT INTO xray_clients (uuid, name, email, inbound_tag, owner_username, expires_at, max_conns)
		VALUES ($1, $2, $3, $4, $5, $6, $7)
		ON CONFLICT (uuid) DO UPDATE SET
			name           = EXCLUDED.name,
			email          = EXCLUDED.email,
			inbound_tag    = CASE WHEN EXCLUDED.inbound_tag <> '' THEN EXCLUDED.inbound_tag ELSE xray_clients.inbound_tag END,
			owner_username = CASE WHEN EXCLUDED.owner_username <> '' THEN EXCLUDED.owner_username ELSE xray_clients.owner_username END,
			expires_at     = EXCLUDED.expires_at,
			max_conns      = EXCLUDED.max_conns`,
		m.UUID, m.Name, m.Email, m.InboundTag, m.OwnerUsername, expiresAt, m.MaxConns)
	return err
}

func (s *Store) GetXrayClientMeta(ctx context.Context, uuid string) (*XrayClientMeta, error) {
	m := &XrayClientMeta{}
	var expiresAt sql.NullTime
	var lastActive sql.NullTime
	err := s.db.QueryRowContext(ctx, `
		SELECT uuid, name, email, inbound_tag, COALESCE(owner_username, ''), expires_at, max_conns, created_at,
		       COALESCE(total_uplink_bytes, 0), COALESCE(total_downlink_bytes, 0), last_active, COALESCE(active_connections, 0)
		FROM xray_clients WHERE uuid = $1`, uuid).
		Scan(&m.UUID, &m.Name, &m.Email, &m.InboundTag, &m.OwnerUsername, &expiresAt, &m.MaxConns, &m.CreatedAt, &m.TotalUplinkBytes, &m.TotalDownlinkBytes, &lastActive, &m.ActiveConnections)
	if err != nil {
		return nil, err
	}
	if expiresAt.Valid {
		m.ExpiresAt = &expiresAt.Time
	}
	if lastActive.Valid {
		m.LastActive = &lastActive.Time
	}
	return m, nil
}

func (s *Store) DeleteXrayClientMeta(ctx context.Context, uuid string) error {
	_, err := s.db.ExecContext(ctx, `DELETE FROM xray_clients WHERE uuid = $1`, uuid)
	return err
}

func (s *Store) ListAllXrayClients(ctx context.Context) ([]*XrayClientMeta, error) {
	rows, err := s.db.QueryContext(ctx, `
		SELECT uuid, name, email, inbound_tag, COALESCE(owner_username, ''), expires_at, max_conns, created_at,
		       COALESCE(total_uplink_bytes, 0), COALESCE(total_downlink_bytes, 0), last_active, COALESCE(active_connections, 0)
		FROM xray_clients ORDER BY created_at DESC`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanXrayClientMetaRows(rows)
}

func (s *Store) ListXrayClientsByOwner(ctx context.Context, ownerUsername string) ([]*XrayClientMeta, error) {
	rows, err := s.db.QueryContext(ctx, `
		SELECT uuid, name, email, inbound_tag, COALESCE(owner_username, ''), expires_at, max_conns, created_at,
		       COALESCE(total_uplink_bytes, 0), COALESCE(total_downlink_bytes, 0), last_active, COALESCE(active_connections, 0)
		FROM xray_clients WHERE owner_username = $1 ORDER BY created_at DESC`, ownerUsername)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanXrayClientMetaRows(rows)
}

func (s *Store) CountXrayClientsByOwner(ctx context.Context, ownerUsername string) (int, error) {
	var n int
	err := s.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM xray_clients WHERE owner_username = $1`, ownerUsername).Scan(&n)
	return n, err
}

func (s *Store) ListExpiredXrayClients(ctx context.Context) ([]*XrayClientMeta, error) {
	rows, err := s.db.QueryContext(ctx, `
		SELECT uuid, name, email, inbound_tag, COALESCE(owner_username, ''), expires_at, max_conns, created_at,
		       COALESCE(total_uplink_bytes, 0), COALESCE(total_downlink_bytes, 0), last_active, COALESCE(active_connections, 0)
		FROM xray_clients WHERE expires_at IS NOT NULL AND expires_at <= NOW()`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	return scanXrayClientMetaRows(rows)
}

func scanXrayClientMetaRows(rows *sql.Rows) ([]*XrayClientMeta, error) {
	var out []*XrayClientMeta
	for rows.Next() {
		m := &XrayClientMeta{}
		var expiresAt sql.NullTime
		var lastActive sql.NullTime
		if err := rows.Scan(&m.UUID, &m.Name, &m.Email, &m.InboundTag, &m.OwnerUsername, &expiresAt, &m.MaxConns, &m.CreatedAt, &m.TotalUplinkBytes, &m.TotalDownlinkBytes, &lastActive, &m.ActiveConnections); err != nil {
			return nil, err
		}
		if expiresAt.Valid {
			m.ExpiresAt = &expiresAt.Time
		}
		if lastActive.Valid {
			m.LastActive = &lastActive.Time
		}
		out = append(out, m)
	}
	return out, rows.Err()
}

// ResetXrayActiveConnections clears stale online counters after the panel starts.
// Native mode then increments/decrements active_connections for real live streams.
func (s *Store) ResetXrayActiveConnections(ctx context.Context) error {
	_, err := s.db.ExecContext(ctx, `UPDATE xray_clients SET active_connections = 0`)
	return err
}

// AddXrayClientTrafficBatch persists native-emulator traffic deltas. It keeps
// totals in PostgreSQL so bandwidth remains visible after panel restarts.
func (s *Store) AddXrayClientTrafficBatch(ctx context.Context, deltas map[string]xrayPendingTraffic) error {
	if len(deltas) == 0 {
		return nil
	}
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	stmt, err := tx.PrepareContext(ctx, `
		UPDATE xray_clients SET
			email                = CASE WHEN email = '' AND $2 <> '' THEN $2 ELSE email END,
			name                 = CASE WHEN name = '' AND $2 <> '' THEN $2 ELSE name END,
			total_uplink_bytes   = GREATEST(total_uplink_bytes + GREATEST($3::BIGINT, 0), 0),
			total_downlink_bytes = GREATEST(total_downlink_bytes + GREATEST($4::BIGINT, 0), 0),
			last_active          = NOW()
		WHERE uuid = $1`)
	if err != nil {
		_ = tx.Rollback()
		return err
	}
	defer stmt.Close()
	for uuid, d := range deltas {
		if uuid == "" || (d.Uplink == 0 && d.Downlink == 0) {
			continue
		}
		if _, err := stmt.ExecContext(ctx, uuid, d.Email, d.Uplink, d.Downlink); err != nil {
			_ = tx.Rollback()
			return err
		}
	}
	return tx.Commit()
}

// UpdateXrayClientActive adjusts the native online connection counter.
func (s *Store) UpdateXrayClientActive(ctx context.Context, uuid, email string, delta int) error {
	if uuid == "" || delta == 0 {
		return nil
	}
	_, err := s.db.ExecContext(ctx, `
		UPDATE xray_clients SET
			email              = CASE WHEN email = '' AND $2 <> '' THEN $2 ELSE email END,
			name               = CASE WHEN name = '' AND $2 <> '' THEN $2 ELSE name END,
			last_active        = CASE WHEN $3::INT > 0 THEN NOW() ELSE last_active END,
			active_connections = GREATEST(active_connections + $3::INT, 0)
		WHERE uuid = $1`, uuid, email, delta)
	return err
}

func countOwnedXrayClients(ctx context.Context, store *Store, ownerUsername string) int {
	if store == nil || ownerUsername == "" {
		return 0
	}
	n, err := store.CountXrayClientsByOwner(ctx, ownerUsername)
	if err != nil {
		log.Printf("count xray clients for %s: %v", ownerUsername, err)
		return 0
	}
	return n
}

func (s *Store) SumXrayClientQuotaByOwner(ctx context.Context, ownerUsername string) (int, error) {
	if s == nil || ownerUsername == "" {
		return 0, nil
	}
	var total int
	err := s.db.QueryRowContext(ctx, `
		SELECT COALESCE(SUM(GREATEST(max_conns, 1)), 0)
		FROM xray_clients WHERE owner_username=$1`, ownerUsername).Scan(&total)
	return total, err
}

func countOwnedSSHQuota(ownerUsername string) int {
	total := 0
	for _, user := range userMgr.List() {
		if user.Cfg.OwnerUsername == ownerUsername {
			total += resellerProvisionCost(user.Cfg.MaxConnections)
		}
	}
	return total
}

func countOwnedXrayQuota(ctx context.Context, store *Store, ownerUsername string) int {
	if store == nil || ownerUsername == "" {
		return 0
	}
	total, err := store.SumXrayClientQuotaByOwner(ctx, ownerUsername)
	if err != nil {
		log.Printf("sum Xray quota for %s: %v", ownerUsername, err)
		return 0
	}
	return total
}

func countOwnedQuota(ctx context.Context, store *Store, ownerUsername string) int {
	return countOwnedSSHQuota(ownerUsername) + countOwnedXrayQuota(ctx, store, ownerUsername)
}

func removeOwnerXrayClients(ctx context.Context, store *Store, ownerUsername string) {
	if store == nil || ownerUsername == "" {
		return
	}
	clients, err := store.ListXrayClientsByOwner(ctx, ownerUsername)
	if err != nil {
		log.Printf("xray owner cleanup: list %s: %v", ownerUsername, err)
		return
	}
	needRestart := false
	for _, m := range clients {
		if m.InboundTag != "" {
			if err := xrayMgr.RemoveXrayClient(m.InboundTag, m.UUID); err != nil {
				log.Printf("xray owner cleanup: remove %s from %s: %v", m.UUID, m.InboundTag, err)
			} else {
				needRestart = true
			}
		}
		if err := store.DeleteXrayClientMeta(ctx, m.UUID); err != nil {
			log.Printf("xray owner cleanup: delete meta %s: %v", m.UUID, err)
		}
	}
	if needRestart {
		xrayMgr.restartIfExternalRunning()
	}
}

// suspendOwnerXrayClients removes an owner's clients from the live Xray config
// while keeping their metadata. That makes reseller suspension reversible.
func suspendOwnerXrayClients(ctx context.Context, store *Store, ownerUsername string) error {
	if store == nil || ownerUsername == "" {
		return nil
	}
	clients, err := store.ListXrayClientsByOwner(ctx, ownerUsername)
	if err != nil {
		return err
	}
	inbounds, err := xrayMgr.ListInbounds()
	if err != nil {
		return err
	}
	present := make(map[string]map[string]bool)
	for _, inbound := range inbounds {
		present[inbound.Tag] = make(map[string]bool)
		for _, client := range inbound.Clients {
			present[inbound.Tag][client.UUID] = true
		}
	}
	changed := false
	var failures []string
	for _, client := range clients {
		if client.InboundTag == "" || !present[client.InboundTag][client.UUID] {
			continue
		}
		if err := xrayMgr.RemoveXrayClient(client.InboundTag, client.UUID); err != nil {
			failures = append(failures, client.UUID+": "+err.Error())
			continue
		}
		changed = true
	}
	if changed {
		xrayMgr.restartIfExternalRunning()
	}
	if len(failures) > 0 {
		return fmt.Errorf("suspend Xray clients: %s", strings.Join(failures, "; "))
	}
	return nil
}

// restoreOwnerXrayClients restores metadata-backed clients after a reseller is
// reactivated. Existing entries are left untouched, so retries are idempotent.
func restoreOwnerXrayClients(ctx context.Context, store *Store, ownerUsername string) error {
	if store == nil || ownerUsername == "" {
		return nil
	}
	clients, err := store.ListXrayClientsByOwner(ctx, ownerUsername)
	if err != nil {
		return err
	}
	inbounds, err := xrayMgr.ListInbounds()
	if err != nil {
		return err
	}
	present := make(map[string]map[string]bool)
	for _, inbound := range inbounds {
		present[inbound.Tag] = make(map[string]bool)
		for _, client := range inbound.Clients {
			present[inbound.Tag][client.UUID] = true
		}
	}
	changed := false
	var failures []string
	for _, client := range clients {
		if client.ExpiresAt != nil && time.Now().After(*client.ExpiresAt) {
			continue
		}
		if client.InboundTag == "" {
			continue
		}
		clientsForInbound, ok := present[client.InboundTag]
		if !ok {
			failures = append(failures, client.UUID+": inbound "+client.InboundTag+" no longer exists")
			continue
		}
		if clientsForInbound[client.UUID] {
			continue
		}
		email := strings.TrimSpace(client.Email)
		if email == "" {
			email = strings.TrimSpace(client.Name)
		}
		if email == "" {
			email = client.UUID
		}
		if err := xrayMgr.AddXrayClient(client.InboundTag, client.UUID, email); err != nil {
			failures = append(failures, client.UUID+": "+err.Error())
			continue
		}
		clientsForInbound[client.UUID] = true
		changed = true
	}
	if changed {
		xrayMgr.restartIfExternalRunning()
	}
	if len(failures) > 0 {
		return fmt.Errorf("restore Xray clients: %s", strings.Join(failures, "; "))
	}
	return nil
}

// startXrayClientExpiryChecker removes expired clients from the live config.
// Reseller-owned metadata is retained so a paid renewal can restore access.
func startXrayClientExpiryChecker(store *Store) {
	if store == nil {
		return
	}
	go func() {
		ticker := time.NewTicker(5 * time.Minute)
		defer ticker.Stop()
		for range ticker.C {
			ctx := context.Background()
			expired, err := store.ListExpiredXrayClients(ctx)
			if err != nil {
				log.Printf("xray expiry checker: list error: %v", err)
				continue
			}
			if len(expired) == 0 {
				continue
			}
			needRestart := false
			present := make(map[string]map[string]bool)
			if inbounds, listErr := xrayMgr.ListInbounds(); listErr == nil {
				for _, inbound := range inbounds {
					present[inbound.Tag] = make(map[string]bool)
					for _, client := range inbound.Clients {
						present[inbound.Tag][client.UUID] = true
					}
				}
			}
			for _, m := range expired {
				tag := m.InboundTag
				if tag == "" {
					if m.OwnerUsername == "" {
						_ = store.DeleteXrayClientMeta(ctx, m.UUID)
					}
					continue
				}
				if present[tag][m.UUID] {
					if err := xrayMgr.RemoveXrayClient(tag, m.UUID); err != nil {
						log.Printf("xray expiry: remove %s from %s: %v", m.UUID, tag, err)
					} else {
						needRestart = true
					}
				}
				if m.OwnerUsername == "" {
					if err := store.DeleteXrayClientMeta(ctx, m.UUID); err != nil {
						log.Printf("xray expiry: delete meta %s: %v", m.UUID, err)
					}
				}
				log.Printf("xray expiry: removed expired client %q (%s) from inbound %s", m.Name, m.UUID, tag)
			}
			if needRestart {
				xrayMgr.restartIfExternalRunning()
			}
		}
	}()
}