3x-ui/web/service/custom_geo.go

package service

import (
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

	"github.com/mhsanaei/3x-ui/v2/config"
	"github.com/mhsanaei/3x-ui/v2/database"
	"github.com/mhsanaei/3x-ui/v2/database/model"
	"github.com/mhsanaei/3x-ui/v2/logger"
)

const (
	customGeoTypeGeosite  = "geosite"
	customGeoTypeGeoip    = "geoip"
	minDatBytes           = 64
	customGeoProbeTimeout = 12 * time.Second
)

var (
	customGeoAliasPattern = regexp.MustCompile(`^[a-z0-9_-]+$`)
	reservedCustomAliases = map[string]struct{}{
		"geoip": {}, "geosite": {},
		"geoip_ir": {}, "geosite_ir": {},
		"geoip_ru": {}, "geosite_ru": {},
	}
	ErrCustomGeoInvalidType    = errors.New("custom_geo_invalid_type")
	ErrCustomGeoAliasRequired  = errors.New("custom_geo_alias_required")
	ErrCustomGeoAliasPattern   = errors.New("custom_geo_alias_pattern")
	ErrCustomGeoAliasReserved  = errors.New("custom_geo_alias_reserved")
	ErrCustomGeoURLRequired    = errors.New("custom_geo_url_required")
	ErrCustomGeoInvalidURL     = errors.New("custom_geo_invalid_url")
	ErrCustomGeoURLScheme      = errors.New("custom_geo_url_scheme")
	ErrCustomGeoURLHost        = errors.New("custom_geo_url_host")
	ErrCustomGeoDuplicateAlias = errors.New("custom_geo_duplicate_alias")
	ErrCustomGeoNotFound       = errors.New("custom_geo_not_found")
	ErrCustomGeoDownload       = errors.New("custom_geo_download")
)

type CustomGeoUpdateAllItem struct {
	Id       int    `json:"id"`
	Alias    string `json:"alias"`
	FileName string `json:"fileName"`
}

type CustomGeoUpdateAllFailure struct {
	Id       int    `json:"id"`
	Alias    string `json:"alias"`
	FileName string `json:"fileName"`
	Err      string `json:"error"`
}

type CustomGeoUpdateAllResult struct {
	Succeeded []CustomGeoUpdateAllItem    `json:"succeeded"`
	Failed    []CustomGeoUpdateAllFailure `json:"failed"`
}

type CustomGeoService struct {
	serverService    *ServerService
	updateAllGetAll  func() ([]model.CustomGeoResource, error)
	updateAllApply   func(id int, onStartup bool) (string, error)
	updateAllRestart func() error
}

func NewCustomGeoService() *CustomGeoService {
	s := &CustomGeoService{
		serverService: &ServerService{},
	}
	s.updateAllGetAll = s.GetAll
	s.updateAllApply = s.applyDownloadAndPersist
	s.updateAllRestart = func() error { return s.serverService.RestartXrayService() }
	return s
}

func NormalizeAliasKey(alias string) string {
	return strings.ToLower(strings.ReplaceAll(alias, "-", "_"))
}

func (s *CustomGeoService) fileNameFor(typ, alias string) string {
	if typ == customGeoTypeGeoip {
		return fmt.Sprintf("geoip_%s.dat", alias)
	}
	return fmt.Sprintf("geosite_%s.dat", alias)
}

func (s *CustomGeoService) validateType(typ string) error {
	if typ != customGeoTypeGeosite && typ != customGeoTypeGeoip {
		return ErrCustomGeoInvalidType
	}
	return nil
}

func (s *CustomGeoService) validateAlias(alias string) error {
	if alias == "" {
		return ErrCustomGeoAliasRequired
	}
	if !customGeoAliasPattern.MatchString(alias) {
		return ErrCustomGeoAliasPattern
	}
	if _, ok := reservedCustomAliases[NormalizeAliasKey(alias)]; ok {
		return ErrCustomGeoAliasReserved
	}
	return nil
}

func (s *CustomGeoService) validateURL(raw string) error {
	if raw == "" {
		return ErrCustomGeoURLRequired
	}
	u, err := url.Parse(raw)
	if err != nil {
		return ErrCustomGeoInvalidURL
	}
	if u.Scheme != "http" && u.Scheme != "https" {
		return ErrCustomGeoURLScheme
	}
	if u.Host == "" {
		return ErrCustomGeoURLHost
	}
	return nil
}

func localDatFileNeedsRepair(path string) bool {
	fi, err := os.Stat(path)
	if err != nil {
		return true
	}
	if fi.IsDir() {
		return true
	}
	return fi.Size() < int64(minDatBytes)
}

func CustomGeoLocalFileNeedsRepair(path string) bool {
	return localDatFileNeedsRepair(path)
}

func probeCustomGeoURLWithGET(rawURL string) error {
	client := &http.Client{Timeout: customGeoProbeTimeout}
	req, err := http.NewRequest(http.MethodGet, rawURL, nil)
	if err != nil {
		return err
	}
	req.Header.Set("Range", "bytes=0-0")
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	_, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, 256))
	switch resp.StatusCode {
	case http.StatusOK, http.StatusPartialContent:
		return nil
	default:
		return fmt.Errorf("get range status %d", resp.StatusCode)
	}
}

func probeCustomGeoURL(rawURL string) error {
	client := &http.Client{Timeout: customGeoProbeTimeout}
	req, err := http.NewRequest(http.MethodHead, rawURL, nil)
	if err != nil {
		return err
	}
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	_ = resp.Body.Close()
	sc := resp.StatusCode
	if sc >= 200 && sc < 300 {
		return nil
	}
	if sc == http.StatusMethodNotAllowed || sc == http.StatusNotImplemented {
		return probeCustomGeoURLWithGET(rawURL)
	}
	return fmt.Errorf("head status %d", sc)
}

func (s *CustomGeoService) EnsureOnStartup() {
	list, err := s.GetAll()
	if err != nil {
		logger.Warning("custom geo startup: load list:", err)
		return
	}
	n := len(list)
	if n == 0 {
		logger.Info("custom geo startup: no custom geofiles configured")
		return
	}
	logger.Infof("custom geo startup: checking %d custom geofile(s)", n)
	for i := range list {
		r := &list[i]
		if err := s.validateURL(r.Url); err != nil {
			logger.Warningf("custom geo startup id=%d: invalid url: %v", r.Id, err)
			continue
		}
		s.syncLocalPath(r)
		localPath := r.LocalPath
		if !localDatFileNeedsRepair(localPath) {
			logger.Infof("custom geo startup id=%d alias=%s path=%s: present", r.Id, r.Alias, localPath)
			continue
		}
		logger.Infof("custom geo startup id=%d alias=%s path=%s: missing or needs repair, probing source", r.Id, r.Alias, localPath)
		if err := probeCustomGeoURL(r.Url); err != nil {
			logger.Warningf("custom geo startup id=%d alias=%s url=%s: probe: %v (attempting download anyway)", r.Id, r.Alias, r.Url, err)
		}
		_, _ = s.applyDownloadAndPersist(r.Id, true)
	}
}

func (s *CustomGeoService) downloadToPath(resourceURL, destPath string, lastModifiedHeader string) (skipped bool, newLastModified string, err error) {
	skipped, lm, err := s.downloadToPathOnce(resourceURL, destPath, lastModifiedHeader, false)
	if err != nil {
		return false, "", err
	}
	if skipped {
		if _, statErr := os.Stat(destPath); statErr == nil && !localDatFileNeedsRepair(destPath) {
			return true, lm, nil
		}
		return s.downloadToPathOnce(resourceURL, destPath, lastModifiedHeader, true)
	}
	return false, lm, nil
}

func (s *CustomGeoService) downloadToPathOnce(resourceURL, destPath string, lastModifiedHeader string, forceFull bool) (skipped bool, newLastModified string, err error) {
	var req *http.Request
	req, err = http.NewRequest(http.MethodGet, resourceURL, nil)
	if err != nil {
		return false, "", fmt.Errorf("%w: %v", ErrCustomGeoDownload, err)
	}

	if !forceFull {
		if fi, statErr := os.Stat(destPath); statErr == nil && !localDatFileNeedsRepair(destPath) {
			if !fi.ModTime().IsZero() {
				req.Header.Set("If-Modified-Since", fi.ModTime().UTC().Format(http.TimeFormat))
			} else if lastModifiedHeader != "" {
				if t, perr := time.Parse(http.TimeFormat, lastModifiedHeader); perr == nil {
					req.Header.Set("If-Modified-Since", t.UTC().Format(http.TimeFormat))
				}
			}
		}
	}

	client := &http.Client{Timeout: 10 * time.Minute}
	resp, err := client.Do(req)
	if err != nil {
		return false, "", fmt.Errorf("%w: %v", ErrCustomGeoDownload, err)
	}
	defer resp.Body.Close()

	var serverModTime time.Time
	if lm := resp.Header.Get("Last-Modified"); lm != "" {
		if parsed, perr := time.Parse(http.TimeFormat, lm); perr == nil {
			serverModTime = parsed
			newLastModified = lm
		}
	}

	updateModTime := func() {
		if !serverModTime.IsZero() {
			_ = os.Chtimes(destPath, serverModTime, serverModTime)
		}
	}

	if resp.StatusCode == http.StatusNotModified {
		if forceFull {
			return false, "", fmt.Errorf("%w: unexpected 304 on unconditional get", ErrCustomGeoDownload)
		}
		updateModTime()
		return true, newLastModified, nil
	}
	if resp.StatusCode != http.StatusOK {
		return false, "", fmt.Errorf("%w: unexpected status %d", ErrCustomGeoDownload, resp.StatusCode)
	}

	binDir := filepath.Dir(destPath)
	if err = os.MkdirAll(binDir, 0o755); err != nil {
		return false, "", fmt.Errorf("%w: %v", ErrCustomGeoDownload, err)
	}

	tmpPath := destPath + ".tmp"
	out, err := os.Create(tmpPath)
	if err != nil {
		return false, "", fmt.Errorf("%w: %v", ErrCustomGeoDownload, err)
	}
	n, err := io.Copy(out, resp.Body)
	closeErr := out.Close()
	if err != nil {
		_ = os.Remove(tmpPath)
		return false, "", fmt.Errorf("%w: %v", ErrCustomGeoDownload, err)
	}
	if closeErr != nil {
		_ = os.Remove(tmpPath)
		return false, "", fmt.Errorf("%w: %v", ErrCustomGeoDownload, closeErr)
	}
	if n < minDatBytes {
		_ = os.Remove(tmpPath)
		return false, "", fmt.Errorf("%w: file too small", ErrCustomGeoDownload)
	}

	if err = os.Rename(tmpPath, destPath); err != nil {
		_ = os.Remove(tmpPath)
		return false, "", fmt.Errorf("%w: %v", ErrCustomGeoDownload, err)
	}

	updateModTime()
	if newLastModified == "" && resp.Header.Get("Last-Modified") != "" {
		newLastModified = resp.Header.Get("Last-Modified")
	}
	return false, newLastModified, nil
}

func (s *CustomGeoService) resolveDestPath(r *model.CustomGeoResource) string {
	if r.LocalPath != "" {
		return r.LocalPath
	}
	return filepath.Join(config.GetBinFolderPath(), s.fileNameFor(r.Type, r.Alias))
}

func (s *CustomGeoService) syncLocalPath(r *model.CustomGeoResource) {
	p := filepath.Join(config.GetBinFolderPath(), s.fileNameFor(r.Type, r.Alias))
	r.LocalPath = p
}

func (s *CustomGeoService) Create(r *model.CustomGeoResource) error {
	if err := s.validateType(r.Type); err != nil {
		return err
	}
	if err := s.validateAlias(r.Alias); err != nil {
		return err
	}
	if err := s.validateURL(r.Url); err != nil {
		return err
	}
	var existing int64
	database.GetDB().Model(&model.CustomGeoResource{}).
		Where("geo_type = ? AND alias = ?", r.Type, r.Alias).Count(&existing)
	if existing > 0 {
		return ErrCustomGeoDuplicateAlias
	}
	s.syncLocalPath(r)
	skipped, lm, err := s.downloadToPath(r.Url, r.LocalPath, r.LastModified)
	if err != nil {
		return err
	}
	now := time.Now().Unix()
	r.LastUpdatedAt = now
	r.LastModified = lm
	if err = database.GetDB().Create(r).Error; err != nil {
		_ = os.Remove(r.LocalPath)
		return err
	}
	logger.Infof("custom geo created id=%d type=%s alias=%s skipped=%v", r.Id, r.Type, r.Alias, skipped)
	if err = s.serverService.RestartXrayService(); err != nil {
		logger.Warning("custom geo create: restart xray:", err)
	}
	return nil
}

func (s *CustomGeoService) Update(id int, r *model.CustomGeoResource) error {
	var cur model.CustomGeoResource
	if err := database.GetDB().First(&cur, id).Error; err != nil {
		if database.IsNotFound(err) {
			return ErrCustomGeoNotFound
		}
		return err
	}
	if err := s.validateType(r.Type); err != nil {
		return err
	}
	if err := s.validateAlias(r.Alias); err != nil {
		return err
	}
	if err := s.validateURL(r.Url); err != nil {
		return err
	}
	if cur.Type != r.Type || cur.Alias != r.Alias {
		var cnt int64
		database.GetDB().Model(&model.CustomGeoResource{}).
			Where("geo_type = ? AND alias = ? AND id <> ?", r.Type, r.Alias, id).
			Count(&cnt)
		if cnt > 0 {
			return ErrCustomGeoDuplicateAlias
		}
	}
	oldPath := s.resolveDestPath(&cur)
	s.syncLocalPath(r)
	r.Id = id
	r.LocalPath = filepath.Join(config.GetBinFolderPath(), s.fileNameFor(r.Type, r.Alias))
	if oldPath != r.LocalPath && oldPath != "" {
		if _, err := os.Stat(oldPath); err == nil {
			_ = os.Remove(oldPath)
		}
	}
	_, lm, err := s.downloadToPath(r.Url, r.LocalPath, cur.LastModified)
	if err != nil {
		return err
	}
	r.LastUpdatedAt = time.Now().Unix()
	r.LastModified = lm
	err = database.GetDB().Model(&model.CustomGeoResource{}).Where("id = ?", id).Updates(map[string]any{
		"geo_type":        r.Type,
		"alias":           r.Alias,
		"url":             r.Url,
		"local_path":      r.LocalPath,
		"last_updated_at": r.LastUpdatedAt,
		"last_modified":   r.LastModified,
	}).Error
	if err != nil {
		return err
	}
	logger.Infof("custom geo updated id=%d", id)
	if err = s.serverService.RestartXrayService(); err != nil {
		logger.Warning("custom geo update: restart xray:", err)
	}
	return nil
}

func (s *CustomGeoService) Delete(id int) (displayName string, err error) {
	var r model.CustomGeoResource
	if err := database.GetDB().First(&r, id).Error; err != nil {
		if database.IsNotFound(err) {
			return "", ErrCustomGeoNotFound
		}
		return "", err
	}
	displayName = s.fileNameFor(r.Type, r.Alias)
	p := s.resolveDestPath(&r)
	if err := database.GetDB().Delete(&model.CustomGeoResource{}, id).Error; err != nil {
		return displayName, err
	}
	if p != "" {
		if _, err := os.Stat(p); err == nil {
			if rmErr := os.Remove(p); rmErr != nil {
				logger.Warningf("custom geo delete file %s: %v", p, rmErr)
			}
		}
	}
	logger.Infof("custom geo deleted id=%d", id)
	if err := s.serverService.RestartXrayService(); err != nil {
		logger.Warning("custom geo delete: restart xray:", err)
	}
	return displayName, nil
}

func (s *CustomGeoService) GetAll() ([]model.CustomGeoResource, error) {
	var list []model.CustomGeoResource
	err := database.GetDB().Order("id asc").Find(&list).Error
	return list, err
}

func (s *CustomGeoService) applyDownloadAndPersist(id int, onStartup bool) (displayName string, err error) {
	var r model.CustomGeoResource
	if err := database.GetDB().First(&r, id).Error; err != nil {
		if database.IsNotFound(err) {
			return "", ErrCustomGeoNotFound
		}
		return "", err
	}
	displayName = s.fileNameFor(r.Type, r.Alias)
	s.syncLocalPath(&r)
	skipped, lm, err := s.downloadToPath(r.Url, r.LocalPath, r.LastModified)
	if err != nil {
		if onStartup {
			logger.Warningf("custom geo startup download id=%d: %v", id, err)
		} else {
			logger.Warningf("custom geo manual update id=%d: %v", id, err)
		}
		return displayName, err
	}
	now := time.Now().Unix()
	updates := map[string]any{
		"last_modified":   lm,
		"local_path":      r.LocalPath,
		"last_updated_at": now,
	}
	if err = database.GetDB().Model(&model.CustomGeoResource{}).Where("id = ?", id).Updates(updates).Error; err != nil {
		if onStartup {
			logger.Warningf("custom geo startup id=%d: persist metadata: %v", id, err)
		} else {
			logger.Warningf("custom geo manual update id=%d: persist metadata: %v", id, err)
		}
		return displayName, err
	}
	if skipped {
		if onStartup {
			logger.Infof("custom geo startup download skipped (not modified) id=%d", id)
		} else {
			logger.Infof("custom geo manual update skipped (not modified) id=%d", id)
		}
	} else {
		if onStartup {
			logger.Infof("custom geo startup download ok id=%d", id)
		} else {
			logger.Infof("custom geo manual update ok id=%d", id)
		}
	}
	return displayName, nil
}

func (s *CustomGeoService) TriggerUpdate(id int) (string, error) {
	displayName, err := s.applyDownloadAndPersist(id, false)
	if err != nil {
		return displayName, err
	}
	if err = s.serverService.RestartXrayService(); err != nil {
		logger.Warning("custom geo manual update: restart xray:", err)
	}
	return displayName, nil
}

func (s *CustomGeoService) TriggerUpdateAll() (*CustomGeoUpdateAllResult, error) {
	var list []model.CustomGeoResource
	var err error
	if s.updateAllGetAll != nil {
		list, err = s.updateAllGetAll()
	} else {
		list, err = s.GetAll()
	}
	if err != nil {
		return nil, err
	}
	res := &CustomGeoUpdateAllResult{}
	if len(list) == 0 {
		return res, nil
	}
	for _, r := range list {
		var name string
		var applyErr error
		if s.updateAllApply != nil {
			name, applyErr = s.updateAllApply(r.Id, false)
		} else {
			name, applyErr = s.applyDownloadAndPersist(r.Id, false)
		}
		if applyErr != nil {
			res.Failed = append(res.Failed, CustomGeoUpdateAllFailure{
				Id: r.Id, Alias: r.Alias, FileName: name, Err: applyErr.Error(),
			})
			continue
		}
		res.Succeeded = append(res.Succeeded, CustomGeoUpdateAllItem{
			Id: r.Id, Alias: r.Alias, FileName: name,
		})
	}
	if len(res.Succeeded) > 0 {
		var restartErr error
		if s.updateAllRestart != nil {
			restartErr = s.updateAllRestart()
		} else {
			restartErr = s.serverService.RestartXrayService()
		}
		if restartErr != nil {
			logger.Warning("custom geo update all: restart xray:", restartErr)
		}
	}
	return res, nil
}

type CustomGeoAliasItem struct {
	Alias      string `json:"alias"`
	Type       string `json:"type"`
	FileName   string `json:"fileName"`
	ExtExample string `json:"extExample"`
}

type CustomGeoAliasesResponse struct {
	Geosite []CustomGeoAliasItem `json:"geosite"`
	Geoip   []CustomGeoAliasItem `json:"geoip"`
}

func (s *CustomGeoService) GetAliasesForUI() (CustomGeoAliasesResponse, error) {
	list, err := s.GetAll()
	if err != nil {
		logger.Warning("custom geo GetAliasesForUI:", err)
		return CustomGeoAliasesResponse{}, err
	}
	var out CustomGeoAliasesResponse
	for _, r := range list {
		fn := s.fileNameFor(r.Type, r.Alias)
		ex := fmt.Sprintf("ext:%s:tag", fn)
		item := CustomGeoAliasItem{
			Alias:      r.Alias,
			Type:       r.Type,
			FileName:   fn,
			ExtExample: ex,
		}
		if r.Type == customGeoTypeGeoip {
			out.Geoip = append(out.Geoip, item)
		} else {
			out.Geosite = append(out.Geosite, item)
		}
	}
	return out, nil
}