geosite_compact.go 2.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100
  1. package router
  2. import (
  3. "encoding/gob"
  4. "errors"
  5. "io"
  6. "runtime"
  7. "github.com/xtls/xray-core/common/strmatcher"
  8. )
  9. type geoSiteListGob struct {
  10. Sites map[string][]byte
  11. Deps map[string][]string
  12. Hosts map[string][]string
  13. }
  14. func SerializeGeoSiteList(sites []*GeoSite, deps map[string][]string, hosts map[string][]string, w io.Writer) error {
  15. data := geoSiteListGob{
  16. Sites: make(map[string][]byte),
  17. Deps: deps,
  18. Hosts: hosts,
  19. }
  20. for _, site := range sites {
  21. if site == nil {
  22. continue
  23. }
  24. var buf bytesWriter
  25. if err := SerializeDomainMatcher(site.Domain, &buf); err != nil {
  26. return err
  27. }
  28. data.Sites[site.CountryCode] = buf.Bytes()
  29. }
  30. return gob.NewEncoder(w).Encode(data)
  31. }
  32. type bytesWriter struct {
  33. data []byte
  34. }
  35. func (w *bytesWriter) Write(p []byte) (n int, err error) {
  36. w.data = append(w.data, p...)
  37. return len(p), nil
  38. }
  39. func (w *bytesWriter) Bytes() []byte {
  40. return w.data
  41. }
  42. func LoadGeoSiteMatcher(r io.Reader, countryCode string) (strmatcher.IndexMatcher, error) {
  43. var data geoSiteListGob
  44. if err := gob.NewDecoder(r).Decode(&data); err != nil {
  45. return nil, err
  46. }
  47. return loadWithDeps(&data, countryCode, make(map[string]bool))
  48. }
  49. func loadWithDeps(data *geoSiteListGob, code string, visited map[string]bool) (strmatcher.IndexMatcher, error) {
  50. if visited[code] {
  51. return nil, errors.New("cyclic dependency")
  52. }
  53. visited[code] = true
  54. var matchers []strmatcher.IndexMatcher
  55. if siteData, ok := data.Sites[code]; ok {
  56. m, err := NewDomainMatcherFromBuffer(siteData)
  57. if err == nil {
  58. matchers = append(matchers, m)
  59. }
  60. }
  61. if deps, ok := data.Deps[code]; ok {
  62. for _, dep := range deps {
  63. m, err := loadWithDeps(data, dep, visited)
  64. if err == nil {
  65. matchers = append(matchers, m)
  66. }
  67. }
  68. }
  69. if len(matchers) == 0 {
  70. return nil, errors.New("matcher not found for: " + code)
  71. }
  72. if len(matchers) == 1 {
  73. return matchers[0], nil
  74. }
  75. runtime.GC()
  76. return &strmatcher.IndexMatcherGroup{Matchers: matchers}, nil
  77. }
  78. func LoadGeoSiteHosts(r io.Reader) (map[string][]string, error) {
  79. var data geoSiteListGob
  80. if err := gob.NewDecoder(r).Decode(&data); err != nil {
  81. return nil, err
  82. }
  83. return data.Hosts, nil
  84. }