provider.go 5.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258
  1. package provider
  2. import (
  3. "bytes"
  4. "encoding/json"
  5. "errors"
  6. "io"
  7. "runtime"
  8. "strings"
  9. "time"
  10. "github.com/metacubex/mihomo/common/pool"
  11. "github.com/metacubex/mihomo/component/resource"
  12. C "github.com/metacubex/mihomo/constant"
  13. P "github.com/metacubex/mihomo/constant/provider"
  14. "gopkg.in/yaml.v3"
  15. )
  16. var tunnel P.Tunnel
  17. func SetTunnel(t P.Tunnel) {
  18. tunnel = t
  19. }
  20. type ruleSetProvider struct {
  21. *resource.Fetcher[ruleStrategy]
  22. behavior P.RuleBehavior
  23. format P.RuleFormat
  24. strategy ruleStrategy
  25. }
  26. type RuleSetProvider struct {
  27. *ruleSetProvider
  28. }
  29. type RulePayload struct {
  30. /**
  31. key: Domain or IP Cidr
  32. value: Rule type or is empty
  33. */
  34. Payload []string `yaml:"payload"`
  35. Rules []string `yaml:"rules"`
  36. }
  37. type ruleStrategy interface {
  38. Behavior() P.RuleBehavior
  39. Match(metadata *C.Metadata) bool
  40. Count() int
  41. ShouldResolveIP() bool
  42. ShouldFindProcess() bool
  43. Reset()
  44. Insert(rule string)
  45. FinishInsert()
  46. }
  47. type mrsRuleStrategy interface {
  48. ruleStrategy
  49. FromMrs(r io.Reader, count int) error
  50. WriteMrs(w io.Writer) error
  51. DumpMrs(f func(key string) bool)
  52. }
  53. func (rp *ruleSetProvider) Type() P.ProviderType {
  54. return P.Rule
  55. }
  56. func (rp *ruleSetProvider) Initial() error {
  57. elm, err := rp.Fetcher.Initial()
  58. if err != nil {
  59. return err
  60. }
  61. rp.OnUpdate(elm)
  62. return nil
  63. }
  64. func (rp *ruleSetProvider) Update() error {
  65. elm, same, err := rp.Fetcher.Update()
  66. if err == nil && !same {
  67. rp.OnUpdate(elm)
  68. return nil
  69. }
  70. return err
  71. }
  72. func (rp *ruleSetProvider) Behavior() P.RuleBehavior {
  73. return rp.behavior
  74. }
  75. func (rp *ruleSetProvider) Match(metadata *C.Metadata) bool {
  76. return rp.strategy != nil && rp.strategy.Match(metadata)
  77. }
  78. func (rp *ruleSetProvider) ShouldResolveIP() bool {
  79. return rp.strategy.ShouldResolveIP()
  80. }
  81. func (rp *ruleSetProvider) ShouldFindProcess() bool {
  82. return rp.strategy.ShouldFindProcess()
  83. }
  84. func (rp *ruleSetProvider) Strategy() any {
  85. return rp.strategy
  86. }
  87. func (rp *ruleSetProvider) MarshalJSON() ([]byte, error) {
  88. return json.Marshal(
  89. map[string]interface{}{
  90. "behavior": rp.behavior.String(),
  91. "format": rp.format.String(),
  92. "name": rp.Name(),
  93. "ruleCount": rp.strategy.Count(),
  94. "type": rp.Type().String(),
  95. "updatedAt": rp.UpdatedAt,
  96. "vehicleType": rp.VehicleType().String(),
  97. })
  98. }
  99. func NewRuleSetProvider(name string, behavior P.RuleBehavior, format P.RuleFormat, interval time.Duration, vehicle P.Vehicle,
  100. parse func(tp, payload, target string, params []string, subRules map[string][]C.Rule) (parsed C.Rule, parseErr error)) P.RuleProvider {
  101. rp := &ruleSetProvider{
  102. behavior: behavior,
  103. format: format,
  104. }
  105. onUpdate := func(strategy ruleStrategy) {
  106. rp.strategy = strategy
  107. tunnel.RuleUpdateCallback().Emit(rp)
  108. }
  109. rp.strategy = newStrategy(behavior, parse)
  110. rp.Fetcher = resource.NewFetcher(name, interval, vehicle, func(bytes []byte) (ruleStrategy, error) {
  111. return rulesParse(bytes, newStrategy(behavior, parse), format)
  112. }, onUpdate)
  113. wrapper := &RuleSetProvider{
  114. rp,
  115. }
  116. final := func(provider *RuleSetProvider) { _ = rp.Fetcher.Destroy() }
  117. runtime.SetFinalizer(wrapper, final)
  118. return wrapper
  119. }
  120. func newStrategy(behavior P.RuleBehavior, parse func(tp, payload, target string, params []string, subRules map[string][]C.Rule) (parsed C.Rule, parseErr error)) ruleStrategy {
  121. switch behavior {
  122. case P.Domain:
  123. strategy := NewDomainStrategy()
  124. return strategy
  125. case P.IPCIDR:
  126. strategy := NewIPCidrStrategy()
  127. return strategy
  128. case P.Classical:
  129. strategy := NewClassicalStrategy(parse)
  130. return strategy
  131. default:
  132. return nil
  133. }
  134. }
  135. var ErrNoPayload = errors.New("file must have a `payload` field")
  136. var ErrInvalidFormat = errors.New("invalid format")
  137. func rulesParse(buf []byte, strategy ruleStrategy, format P.RuleFormat) (ruleStrategy, error) {
  138. strategy.Reset()
  139. if format == P.MrsRule {
  140. return rulesMrsParse(buf, strategy)
  141. }
  142. schema := &RulePayload{}
  143. firstLineBuffer := pool.GetBuffer()
  144. defer pool.PutBuffer(firstLineBuffer)
  145. firstLineLength := 0
  146. s := 0 // search start index
  147. for s < len(buf) {
  148. // search buffer for a new line.
  149. line := buf[s:]
  150. if i := bytes.IndexByte(line, '\n'); i >= 0 {
  151. i += s
  152. line = buf[s : i+1]
  153. s = i + 1
  154. } else {
  155. s = len(buf) // stop loop in next step
  156. if firstLineLength == 0 && format == P.YamlRule { // no head or only one line body
  157. return nil, ErrNoPayload
  158. }
  159. }
  160. var str string
  161. switch format {
  162. case P.TextRule:
  163. str = string(line)
  164. str = strings.TrimSpace(str)
  165. if len(str) == 0 {
  166. continue
  167. }
  168. if str[0] == '#' { // comment
  169. continue
  170. }
  171. if strings.HasPrefix(str, "//") { // comment in Premium core
  172. continue
  173. }
  174. case P.YamlRule:
  175. trimLine := bytes.TrimSpace(line)
  176. if len(trimLine) == 0 {
  177. continue
  178. }
  179. if trimLine[0] == '#' { // comment
  180. continue
  181. }
  182. firstLineBuffer.Write(line)
  183. if firstLineLength == 0 { // find payload head
  184. firstLineLength = firstLineBuffer.Len()
  185. firstLineBuffer.WriteString(" - ''") // a test line
  186. err := yaml.Unmarshal(firstLineBuffer.Bytes(), schema)
  187. firstLineBuffer.Truncate(firstLineLength)
  188. if err == nil && (len(schema.Rules) > 0 || len(schema.Payload) > 0) { // found
  189. continue
  190. }
  191. // not found or err!=nil
  192. firstLineBuffer.Truncate(0)
  193. firstLineLength = 0
  194. continue
  195. }
  196. // parse payload body
  197. err := yaml.Unmarshal(firstLineBuffer.Bytes(), schema)
  198. firstLineBuffer.Truncate(firstLineLength)
  199. if err != nil {
  200. continue
  201. }
  202. if len(schema.Rules) > 0 {
  203. str = schema.Rules[0]
  204. }
  205. if len(schema.Payload) > 0 {
  206. str = schema.Payload[0]
  207. }
  208. default:
  209. return nil, ErrInvalidFormat
  210. }
  211. if str == "" {
  212. continue
  213. }
  214. strategy.Insert(str)
  215. }
  216. strategy.FinishInsert()
  217. return strategy, nil
  218. }