Wormhole/internal/store/query.go

package store

import (
	"fmt"
	"strconv"

	"gorm.io/gorm"
)

type PageParams struct {
	Page     int
	PageSize int
	Q        string
}

func NormalizePage(p PageParams) (offset, limit int) {
	if p.Page < 1 {
		p.Page = 1
	}
	if p.PageSize < 1 {
		p.PageSize = 20
	}
	if p.PageSize > 200 {
		p.PageSize = 200
	}
	return (p.Page - 1) * p.PageSize, p.PageSize
}

func tunnelOwnerScope(db *gorm.DB, clientID, ownerID uint, isAdmin bool) *gorm.DB {
	q := db.Model(&Tunnel{})
	if clientID > 0 {
		return q.Where("client_id = ?", clientID)
	}
	if !isAdmin {
		var ids []uint
		db.Model(&Client{}).Where("owner_user_id = ?", ownerID).Pluck("id", &ids)
		if len(ids) == 0 {
			return q.Where("1 = 0")
		}
		return q.Where("client_id IN ?", ids)
	}
	return q
}

func hostOwnerScope(db *gorm.DB, clientID, ownerID uint, isAdmin bool) *gorm.DB {
	q := db.Model(&Host{})
	if clientID > 0 {
		return q.Where("client_id = ?", clientID)
	}
	if !isAdmin {
		var ids []uint
		db.Model(&Client{}).Where("owner_user_id = ?", ownerID).Pluck("id", &ids)
		if len(ids) == 0 {
			return q.Where("1 = 0")
		}
		return q.Where("client_id IN ?", ids)
	}
	return q
}

func applyTunnelSearch(q *gorm.DB, keyword string) *gorm.DB {
	if keyword == "" {
		return q
	}
	like := "%" + keyword + "%"
	portLike := "%" + keyword + "%"
	return q.Where(
		"name LIKE ? OR target_addr LIKE ? OR CAST(listen_port AS TEXT) LIKE ?",
		like, like, portLike,
	)
}

func applyHostSearch(q *gorm.DB, keyword string) *gorm.DB {
	if keyword == "" {
		return q
	}
	like := "%" + keyword + "%"
	return q.Where("host LIKE ? OR target_addr LIKE ? OR location LIKE ?", like, like, like)
}

func applyDomainForwardSearch(q *gorm.DB, keyword string) *gorm.DB {
	if keyword == "" {
		return q
	}
	like := "%" + keyword + "%"
	return q.Where(
		"source_host LIKE ? OR target_url LIKE ? OR source_location LIKE ?",
		like, like, like,
	)
}

func (s *Store) ListDomainForwardsPaged(p PageParams) ([]DomainForward, int64, error) {
	var items []DomainForward
	var total int64
	q := applyDomainForwardSearch(s.db.Model(&DomainForward{}), p.Q)
	if err := q.Count(&total).Error; err != nil {
		return nil, 0, err
	}
	offset, limit := NormalizePage(p)
	err := q.Order("id asc").Offset(offset).Limit(limit).Find(&items).Error
	return items, total, err
}

func (s *Store) ListTunnelsPaged(clientID, ownerID uint, isAdmin bool, p PageParams) ([]Tunnel, int64, error) {
	var tunnels []Tunnel
	var total int64
	q := tunnelOwnerScope(s.db, clientID, ownerID, isAdmin)
	q = applyTunnelSearch(q, p.Q)
	if err := q.Count(&total).Error; err != nil {
		return nil, 0, err
	}
	offset, limit := NormalizePage(p)
	err := q.Order("id asc").Offset(offset).Limit(limit).Find(&tunnels).Error
	return tunnels, total, err
}

func (s *Store) ListHostsPaged(clientID, ownerID uint, isAdmin bool, p PageParams) ([]Host, int64, error) {
	var hosts []Host
	var total int64
	q := hostOwnerScope(s.db, clientID, ownerID, isAdmin)
	q = applyHostSearch(q, p.Q)
	if err := q.Count(&total).Error; err != nil {
		return nil, 0, err
	}
	offset, limit := NormalizePage(p)
	err := q.Order("id asc").Offset(offset).Limit(limit).Find(&hosts).Error
	return hosts, total, err
}

func (s *Store) TunnelPortInUse(port int, excludeID uint) (bool, error) {
	if port <= 0 {
		return false, fmt.Errorf("invalid port")
	}
	q := s.db.Model(&Tunnel{}).Where("listen_port = ?", port)
	if excludeID > 0 {
		q = q.Where("id <> ?", excludeID)
	}
	var count int64
	if err := q.Count(&count).Error; err != nil {
		return false, err
	}
	return count > 0, nil
}

func (s *Store) HostRouteInUse(host, location string, excludeID uint) (bool, error) {
	if host == "" {
		return false, fmt.Errorf("host required")
	}
	if location == "" {
		location = "/"
	}
	q := s.db.Model(&Host{}).Where("host = ? AND location = ?", host, location)
	if excludeID > 0 {
		q = q.Where("id <> ?", excludeID)
	}
	var count int64
	if err := q.Count(&count).Error; err != nil {
		return false, err
	}
	return count > 0, nil
}

func (s *Store) ListTunnelsByFlow(limit, offset int) ([]Tunnel, int64, error) {
	var tunnels []Tunnel
	var total int64
	q := s.db.Model(&Tunnel{})
	if err := q.Count(&total).Error; err != nil {
		return nil, 0, err
	}
	if limit <= 0 {
		limit = 10
	}
	err := q.Order("(inlet_flow + export_flow) desc").Offset(offset).Limit(limit).Find(&tunnels).Error
	return tunnels, total, err
}

func (s *Store) ListHostsByFlow(limit, offset int) ([]Host, int64, error) {
	var hosts []Host
	var total int64
	q := s.db.Model(&Host{})
	if err := q.Count(&total).Error; err != nil {
		return nil, 0, err
	}
	if limit <= 0 {
		limit = 10
	}
	err := q.Order("(inlet_flow + export_flow) desc").Offset(offset).Limit(limit).Find(&hosts).Error
	return hosts, total, err
}

func (s *Store) ListAllRequestIPs(limit, offset int, qstr string) ([]RequestIP, int64, error) {
	var items []RequestIP
	var total int64
	q := s.db.Model(&RequestIP{})
	if qstr != "" {
		like := "%" + qstr + "%"
		q = q.Where("ip LIKE ? OR direct_ip LIKE ?", like, like)
	}
	if err := q.Count(&total).Error; err != nil {
		return nil, 0, err
	}
	if limit <= 0 {
		limit = 50
	}
	err := q.Order("hit_count desc").Offset(offset).Limit(limit).Find(&items).Error
	return items, total, err
}

func ParsePageQuery(pageStr, pageSizeStr, q string) PageParams {
	page, _ := strconv.Atoi(pageStr)
	pageSize, _ := strconv.Atoi(pageSizeStr)
	return PageParams{Page: page, PageSize: pageSize, Q: q}
}