feat: support convert mrs format back to text format

This commit is contained in:
wwqgtxx 2024-07-28 10:07:37 +08:00
parent 1db3e4583b
commit c830b8aaf7
10 changed files with 137 additions and 12 deletions

View File

@ -57,6 +57,16 @@ func (set *IpCidrSet) Merge() error {
return nil
}
func (set *IpCidrSet) Foreach(f func(prefix netip.Prefix) bool) {
for _, r := range set.rr {
for _, prefix := range r.Prefixes() {
if !f(prefix) {
return
}
}
}
}
// ToIPSet not safe convert to *netipx.IPSet
// be careful, must be used after Merge
func (set *IpCidrSet) ToIPSet() *netipx.IPSet {

View File

@ -123,16 +123,18 @@ func (t *DomainTrie[T]) Optimize() {
t.root.optimize()
}
func (t *DomainTrie[T]) Foreach(print func(domain string, data T)) {
func (t *DomainTrie[T]) Foreach(fn func(domain string, data T) bool) {
for key, data := range t.root.getChildren() {
recursion([]string{key}, data, print)
recursion([]string{key}, data, fn)
if data != nil && data.inited {
print(joinDomain([]string{key}), data.data)
if !fn(joinDomain([]string{key}), data.data) {
return
}
}
}
}
func recursion[T any](items []string, node *Node[T], fn func(domain string, data T)) {
func recursion[T any](items []string, node *Node[T], fn func(domain string, data T) bool) bool {
for key, data := range node.getChildren() {
newItems := append([]string{key}, items...)
if data != nil && data.inited {
@ -140,10 +142,15 @@ func recursion[T any](items []string, node *Node[T], fn func(domain string, data
if domain[0] == domainStepByte {
domain = complexWildcard + domain
}
fn(domain, data.Data())
if !fn(domain, data.Data()) {
return false
}
}
if !recursion(newItems, data, fn) {
return false
}
recursion(newItems, data, fn)
}
return true
}
func joinDomain(items []string) string {

View File

@ -28,8 +28,9 @@ type qElt struct{ s, e, col int }
// NewDomainSet creates a new *DomainSet struct, from a DomainTrie.
func (t *DomainTrie[T]) NewDomainSet() *DomainSet {
reserveDomains := make([]string, 0)
t.Foreach(func(domain string, data T) {
t.Foreach(func(domain string, data T) bool {
reserveDomains = append(reserveDomains, utils.Reverse(domain))
return true
})
// ensure that the same prefix is continuous
// and according to the ascending sequence of length
@ -136,6 +137,41 @@ func (ss *DomainSet) Has(key string) bool {
}
func (ss *DomainSet) keys(f func(key string) bool) {
var currentKey []byte
var traverse func(int, int) bool
traverse = func(nodeId, bmIdx int) bool {
if getBit(ss.leaves, nodeId) != 0 {
if !f(string(currentKey)) {
return false
}
}
for ; ; bmIdx++ {
if getBit(ss.labelBitmap, bmIdx) != 0 {
return true
}
nextLabel := ss.labels[bmIdx-nodeId]
currentKey = append(currentKey, nextLabel)
nextNodeId := countZeros(ss.labelBitmap, ss.ranks, bmIdx+1)
nextBmIdx := selectIthOne(ss.labelBitmap, ss.ranks, ss.selects, nextNodeId-1) + 1
if !traverse(nextNodeId, nextBmIdx) {
return false
}
currentKey = currentKey[:len(currentKey)-1]
}
}
traverse(0, 0)
return
}
func (ss *DomainSet) Foreach(f func(key string) bool) {
ss.keys(func(key string) bool {
return f(utils.Reverse(key))
})
}
func setBit(bm *[]uint64, i int, v int) {
for i>>6 >= len(*bm) {
*bm = append(*bm, 0)

View File

@ -1,12 +1,29 @@
package trie_test
import (
"golang.org/x/exp/slices"
"testing"
"github.com/metacubex/mihomo/component/trie"
"github.com/stretchr/testify/assert"
)
func testDump(t *testing.T, tree *trie.DomainTrie[struct{}], set *trie.DomainSet) {
var dataSrc []string
tree.Foreach(func(domain string, data struct{}) bool {
dataSrc = append(dataSrc, domain)
return true
})
slices.Sort(dataSrc)
var dataSet []string
set.Foreach(func(key string) bool {
dataSet = append(dataSet, key)
return true
})
slices.Sort(dataSet)
assert.Equal(t, dataSrc, dataSet)
}
func TestDomainSet(t *testing.T) {
tree := trie.New[struct{}]()
domainSet := []string{
@ -33,6 +50,7 @@ func TestDomainSet(t *testing.T) {
assert.True(t, set.Has("google.com"))
assert.False(t, set.Has("qq.com"))
assert.False(t, set.Has("www.baidu.com"))
testDump(t, tree, set)
}
func TestDomainSetComplexWildcard(t *testing.T) {
@ -55,6 +73,7 @@ func TestDomainSetComplexWildcard(t *testing.T) {
assert.False(t, set.Has("google.com"))
assert.True(t, set.Has("www.baidu.com"))
assert.True(t, set.Has("test.test.baidu.com"))
testDump(t, tree, set)
}
func TestDomainSetWildcard(t *testing.T) {
@ -82,4 +101,5 @@ func TestDomainSetWildcard(t *testing.T) {
assert.False(t, set.Has("a.www.google.com"))
assert.False(t, set.Has("test.qq.com"))
assert.False(t, set.Has("test.test.test.qq.com"))
testDump(t, tree, set)
}

View File

@ -121,8 +121,9 @@ func TestTrie_Foreach(t *testing.T) {
assert.NoError(t, tree.Insert(domain, localIP))
}
count := 0
tree.Foreach(func(domain string, data netip.Addr) {
tree.Foreach(func(domain string, data netip.Addr) bool {
count++
return true
})
assert.Equal(t, 7, count)
}

View File

@ -944,10 +944,17 @@ rule-providers:
type: file
rule3:
# mrs类型ruleset目前仅支持domain和ipcidr(即不支持classical
# behavior=domain,format=yaml 可以通过“mihomo convert-ruleset domain yaml XXX.yaml XXX.mrs”转换得到
# behavior=domain,format=text 可以通过“mihomo convert-ruleset domain text XXX.text XXX.mrs”转换得到
# behavior=ipcidr,format=yaml 可以通过“mihomo convert-ruleset ipcidr yaml XXX.yaml XXX.mrs”转换得到
# behavior=ipcidr,format=text 可以通过“mihomo convert-ruleset ipcidr text XXX.text XXX.mrs”转换得到
#
# 对于behavior=domain:
# - format=yaml 可以通过“mihomo convert-ruleset domain yaml XXX.yaml XXX.mrs”转换到mrs格式
# - format=text 可以通过“mihomo convert-ruleset domain text XXX.text XXX.mrs”转换到mrs格式
# - XXX.mrs 可以通过"mihomo convert-ruleset domain mrs XXX.mrs XXX.text"转换回text格式暂不支持转换回ymal格式
#
# 对于behavior=ipcidr:
# - format=yaml 可以通过“mihomo convert-ruleset ipcidr yaml XXX.yaml XXX.mrs”转换到mrs格式
# - format=text 可以通过“mihomo convert-ruleset ipcidr text XXX.text XXX.mrs”转换到mrs格式
# - XXX.mrs 可以通过"mihomo convert-ruleset ipcidr mrs XXX.mrs XXX.text"转换回text格式暂不支持转换回ymal格式
#
type: http
url: "url"
format: mrs

View File

@ -9,6 +9,8 @@ import (
C "github.com/metacubex/mihomo/constant"
P "github.com/metacubex/mihomo/constant/provider"
"github.com/metacubex/mihomo/log"
"golang.org/x/exp/slices"
)
type domainStrategy struct {
@ -78,6 +80,26 @@ func (d *domainStrategy) WriteMrs(w io.Writer) error {
return d.domainSet.WriteBin(w)
}
func (d *domainStrategy) DumpMrs(f func(key string) bool) {
if d.domainSet != nil {
var keys []string
d.domainSet.Foreach(func(key string) bool {
keys = append(keys, key)
return true
})
slices.Sort(keys)
for _, key := range keys {
if _, ok := slices.BinarySearch(keys, "+."+key); ok {
continue // ignore the rules added by trie internal processing
}
if !f(key) {
return
}
}
}
}
var _ mrsRuleStrategy = (*domainStrategy)(nil)
func NewDomainStrategy() *domainStrategy {

View File

@ -3,6 +3,7 @@ package provider
import (
"errors"
"io"
"net/netip"
"github.com/metacubex/mihomo/component/cidr"
C "github.com/metacubex/mihomo/constant"
@ -82,6 +83,14 @@ func (i *ipcidrStrategy) WriteMrs(w io.Writer) error {
return i.cidrSet.WriteBin(w)
}
func (i *ipcidrStrategy) DumpMrs(f func(key string) bool) {
if i.cidrSet != nil {
i.cidrSet.Foreach(func(prefix netip.Prefix) bool {
return f(prefix.String())
})
}
}
func (i *ipcidrStrategy) ToIpCidr() *netipx.IPSet {
return i.cidrSet.ToIPSet()
}

View File

@ -3,6 +3,7 @@ package provider
import (
"encoding/binary"
"errors"
"fmt"
"io"
"os"
@ -21,6 +22,17 @@ func ConvertToMrs(buf []byte, behavior P.RuleBehavior, format P.RuleFormat, w io
return errors.New("empty rule")
}
if _strategy, ok := strategy.(mrsRuleStrategy); ok {
if format == P.MrsRule { // export to TextRule
_strategy.DumpMrs(func(key string) bool {
_, err = fmt.Fprintln(w, key)
if err != nil {
return false
}
return true
})
return nil
}
var encoder *zstd.Encoder
encoder, err = zstd.NewWriter(w)
if err != nil {

View File

@ -58,6 +58,7 @@ type mrsRuleStrategy interface {
ruleStrategy
FromMrs(r io.Reader, count int) error
WriteMrs(w io.Writer) error
DumpMrs(f func(key string) bool)
}
func (rp *ruleSetProvider) Type() P.ProviderType {