add rulesetmap type for efficient ruleset lookup for proxychain impl

This commit is contained in:
Kevin Pham
2023-12-01 12:44:12 -06:00
parent 7e455c9d70
commit b996df31b2
3 changed files with 126 additions and 8 deletions

View File

@@ -308,3 +308,40 @@ func debugPrintRule(rule string, err error) {
fmt.Println(rule)
fmt.Println("------------------------------ END DEBUG RULESET -------------------------------")
}
// ======================= RuleSetMap implementation =================================================
// RuleSetMap: A map with domain names as keys and pointers to the corresponding Rules as values.
// This type is used to efficiently access rules based on domain names.
type RuleSetMap map[string]*Rule
// ToMap converts a RuleSet into a RuleSetMap. It transforms each Rule in the RuleSet
// into a map entry where the key is the Rule's domain (lowercase)
// and the value is a pointer to the Rule. This method is used to
// efficiently access rules based on domain names.
// The RuleSetMap may be accessed with or without a "www." prefix in the domain.
func (rs *RuleSet) ToMap() RuleSetMap {
rsm := make(RuleSetMap)
addMapEntry := func(d string, rule *Rule) {
d = strings.ToLower(d)
rsm[d] = rule
if strings.HasPrefix(d, "www.") {
d = strings.TrimPrefix(d, "www.")
rsm[d] = rule
} else {
d = fmt.Sprintf("www.%s", d)
rsm[d] = rule
}
}
for i, rule := range *rs {
rulePtr := &(*rs)[i]
addMapEntry(rule.Domain, rulePtr)
for _, domain := range rule.Domains {
addMapEntry(domain, rulePtr)
}
}
return rsm
}

View File

@@ -171,3 +171,55 @@ func TestLoadRulesFromLocalDir(t *testing.T) {
assert.Equal(t, rule.RegexRules[0].Replace, "https:")
}
}
func TestToMap(t *testing.T) {
// Prepare a ruleset with multiple rules, including "www." prefixed domains
rules := RuleSet{
{
Domain: "Example.com",
RegexRules: []Regex{{Match: "match1", Replace: "replace1"}},
},
{
Domain: "www.AnotherExample.com",
RegexRules: []Regex{{Match: "match2", Replace: "replace2"}},
},
{
Domain: "www.foo.bAr.baz.bOol.quX.com",
RegexRules: []Regex{{Match: "match3", Replace: "replace3"}},
},
}
// Convert to RuleSetMap
rsm := rules.ToMap()
// Test for correct number of entries
if len(rsm) != 6 {
t.Errorf("Expected 6 entries in RuleSetMap, got %d", len(rsm))
}
// Test for correct mapping
testDomains := []struct {
domain string
expectedMatch string
}{
{"example.com", "match1"},
{"www.example.com", "match1"},
{"anotherexample.com", "match2"},
{"www.anotherexample.com", "match2"},
{"foo.bar.baz.bool.qux.com", "match3"},
{"no.ruleset.domain.com", ""},
}
for _, test := range testDomains {
if test.domain == "no.ruleset.domain.com" {
assert.Empty(t, test.expectedMatch)
continue
}
rule, exists := rsm[test.domain]
if !exists {
t.Errorf("Expected domain %s to exist in RuleSetMap", test.domain)
} else if rule.RegexRules[0].Match != test.expectedMatch {
t.Errorf("Expected match for %s to be %s, got %s", test.domain, test.expectedMatch, rule.RegexRules[0].Match)
}
}
}