pattern_matcher.go 4.49 KB
Newer Older
1 2 3 4 5 6 7
package main

import (
	"fmt"
	"path/filepath"
	"strings"

Frank Denis's avatar
Frank Denis committed
8 9
	"github.com/k-sone/critbitgo"

10 11 12 13 14 15 16 17 18 19 20
	"github.com/jedisct1/dlog"
)

type PatternType int

const (
	PatternTypeNone PatternType = iota
	PatternTypePrefix
	PatternTypeSuffix
	PatternTypeSubstring
	PatternTypePattern
21
	PatternTypeExact
22 23 24
)

type PatternMatcher struct {
Frank Denis's avatar
Frank Denis committed
25 26
	blockedPrefixes   *critbitgo.Trie
	blockedSuffixes   *critbitgo.Trie
27 28
	blockedSubstrings []string
	blockedPatterns   []string
29
	blockedExact      map[string]interface{}
30 31 32 33 34
	indirectVals      map[string]interface{}
}

func NewPatternPatcher() *PatternMatcher {
	patternMatcher := PatternMatcher{
Frank Denis's avatar
Frank Denis committed
35 36
		blockedPrefixes: critbitgo.NewTrie(),
		blockedSuffixes: critbitgo.NewTrie(),
37
		blockedExact:    make(map[string]interface{}),
38
		indirectVals:    make(map[string]interface{}),
39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56
	}
	return &patternMatcher
}

func isGlobCandidate(str string) bool {
	for i, c := range str {
		if c == '?' || c == '[' {
			return true
		} else if c == '*' && i != 0 && i != len(str)-1 {
			return true
		}
	}
	return false
}

func (patternMatcher *PatternMatcher) Add(pattern string, val interface{}, position int) (PatternType, error) {
	leadingStar := strings.HasPrefix(pattern, "*")
	trailingStar := strings.HasSuffix(pattern, "*")
57
	exact := strings.HasPrefix(pattern, "=")
58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76
	patternType := PatternTypeNone
	if isGlobCandidate(pattern) {
		patternType = PatternTypePattern
		_, err := filepath.Match(pattern, "example.com")
		if len(pattern) < 2 || err != nil {
			return patternType, fmt.Errorf("Syntax error in block rules at pattern %d", position)
		}
	} else if leadingStar && trailingStar {
		patternType = PatternTypeSubstring
		if len(pattern) < 3 {
			return patternType, fmt.Errorf("Syntax error in block rules at pattern %d", position)
		}
		pattern = pattern[1 : len(pattern)-1]
	} else if trailingStar {
		patternType = PatternTypePrefix
		if len(pattern) < 2 {
			return patternType, fmt.Errorf("Syntax error in block rules at pattern %d", position)
		}
		pattern = pattern[:len(pattern)-1]
77 78 79 80 81 82
	} else if exact {
		patternType = PatternTypeExact
		if len(pattern) < 2 {
			return patternType, fmt.Errorf("Syntax error in block rules at pattern %d", position)
		}
		pattern = pattern[1:]
83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106
	} else {
		patternType = PatternTypeSuffix
		if leadingStar {
			pattern = pattern[1:]
		}
		pattern = strings.TrimPrefix(pattern, ".")
	}
	if len(pattern) == 0 {
		dlog.Errorf("Syntax error in block rule at line %d", position)
	}

	pattern = strings.ToLower(pattern)
	switch patternType {
	case PatternTypeSubstring:
		patternMatcher.blockedSubstrings = append(patternMatcher.blockedSubstrings, pattern)
		if val != nil {
			patternMatcher.indirectVals[pattern] = val
		}
	case PatternTypePattern:
		patternMatcher.blockedPatterns = append(patternMatcher.blockedPatterns, pattern)
		if val != nil {
			patternMatcher.indirectVals[pattern] = val
		}
	case PatternTypePrefix:
Frank Denis's avatar
Frank Denis committed
107
		patternMatcher.blockedPrefixes.Insert([]byte(pattern), val)
108
	case PatternTypeSuffix:
Frank Denis's avatar
Frank Denis committed
109
		patternMatcher.blockedSuffixes.Insert([]byte(StringReverse(pattern)), val)
110 111
	case PatternTypeExact:
		patternMatcher.blockedExact[pattern] = val
112 113 114 115 116 117 118 119 120 121 122 123
	default:
		dlog.Fatal("Unexpected block type")
	}
	return patternType, nil
}

func (patternMatcher *PatternMatcher) Eval(qName string) (reject bool, reason string, val interface{}) {
	if len(qName) < 2 {
		return false, "", nil
	}

	revQname := StringReverse(qName)
Frank Denis's avatar
Frank Denis committed
124
	if match, xval, found := patternMatcher.blockedSuffixes.LongestPrefix([]byte(revQname)); found {
125 126 127 128 129 130
		if len(match) == len(qName) || revQname[len(match)] == '.' {
			return true, "*." + StringReverse(string(match)), xval
		}
		if len(match) < len(revQname) && len(revQname) > 0 {
			if i := strings.LastIndex(revQname, "."); i > 0 {
				pName := revQname[:i]
Frank Denis's avatar
Frank Denis committed
131
				if match, _, found := patternMatcher.blockedSuffixes.LongestPrefix([]byte(pName)); found {
132 133 134 135 136 137 138 139
					if len(match) == len(pName) || pName[len(match)] == '.' {
						return true, "*." + StringReverse(string(match)), xval
					}
				}
			}
		}
	}

Frank Denis's avatar
Frank Denis committed
140
	if match, xval, found := patternMatcher.blockedPrefixes.LongestPrefix([]byte(qName)); found {
141 142 143 144 145 146 147 148 149 150 151 152 153 154
		return true, string(match) + "*", xval
	}

	for _, substring := range patternMatcher.blockedSubstrings {
		if strings.Contains(qName, substring) {
			return true, "*" + substring + "*", patternMatcher.indirectVals[substring]
		}
	}

	for _, pattern := range patternMatcher.blockedPatterns {
		if found, _ := filepath.Match(pattern, qName); found {
			return true, pattern, patternMatcher.indirectVals[pattern]
		}
	}
155 156 157 158 159

	if xval := patternMatcher.blockedExact[qName]; xval != nil {
		return true, qName, xval
	}

160 161
	return false, "", nil
}