Merge pull request #296 from erbesharat/feature/predefined-regex

[path]: Add predefined regex mode

fallback.go

@@ -64,7 +64,7 @@ func fallback(w http.ResponseWriter, r *http.Request, fallbackType string, code
 func (f *Fallback) countFallback(recType string) {
 	if f.config.Prometheus.Enable {
 		FallbacksCount.WithLabelValues(f.request.Host, recType, f.fallbackType).Add(1)
-		RequestsByStatus.WithLabelValues(f.request.URL.Host, string(f.code)).Add(1)
+		RequestsByStatus.WithLabelValues(f.request.URL.Host, strconv.Itoa(f.code)).Add(1)
 	}
 }
 

path.go

@@ -34,6 +34,17 @@ type Path struct {
 	rec  record
 }
 
+// RegexRecords connects each zone address to a RegexRecord
+// which contains raw TXT record and the re= field
+type RegexRecords map[string]RegexRecord
+
+// RegexRecord holds the TXT record and re= field of a predefined regex record
+type RegexRecord struct {
+	TXT        string
+	Regex      string
+	Submatches []string
+}
+
 var PathRegex = regexp.MustCompile("\\/([A-Za-z0-9-._~!$'()*+,;=:@]+)")
 var FromRegex = regexp.MustCompile("\\/\\$(\\d+)")
 var GroupRegex = regexp.MustCompile("P<[a-zA-Z]+[a-zA-Z0-9]*>")
@@ -83,6 +94,83 @@ func (p *Path) RedirectRoot() error {
 	return nil
 }
 
+// SpecificRecord finds the most specific match using the custom regexes from subzones
+// It goes through all the custom regexes specified in each subzone and uses the
+// most specific match to return the final record.
+func (p *Path) SpecificRecord() (*record, error) {
+	// Iterate subzones and parse the records
+	regexes, err := p.fetchRegexes()
+	if err != nil {
+		return nil, err
+	}
+
+	record, err := p.specificMatch(regexes)
+	if err != nil {
+		return nil, err
+	}
+
+	return record, nil
+}
+
+func (p *Path) specificMatch(regexes RegexRecords) (*record, error) {
+	recordMatch := make(map[int]RegexRecord)
+
+	// Run each regex on the path and list them in a map
+	for _, zone := range regexes {
+		regex, err := regexp.Compile(zone.Regex)
+		if err != nil {
+			return nil, fmt.Errorf("Couldn't compile the regex: %s", err.Error())
+		}
+		matches := regex.FindAllStringSubmatch(p.path, -1)
+		if len(matches) > 0 {
+			zone.Submatches = matches[0]
+		}
+		recordMatch[len(zone.Submatches)] = zone
+	}
+
+	// Sort the map keys to find the most specific match
+	var keys []int
+	for k := range recordMatch {
+		keys = append(keys, k)
+	}
+	sort.Ints(keys)
+
+	// Add the most specific match's path slice to the request context to use in placeholders
+	*p.req = *p.req.WithContext(context.WithValue(p.req.Context(), "regexMatches", recordMatch[keys[len(keys)-1]].Submatches))
+
+	rec := record{}
+	if err := rec.Parse(recordMatch[keys[len(keys)-1]].TXT, p.rw, p.req, p.c); err != nil {
+		return nil, fmt.Errorf("Could not parse record: %s", err)
+	}
+
+	return &rec, nil
+}
+
+func (p *Path) fetchRegexes() (RegexRecords, error) {
+	regexes := make(RegexRecords)
+	for i, loop := 1, true; loop != false; i++ {
+		txts, err := query(fmt.Sprintf("%d.%s", i, p.req.Host), p.req.Context(), p.c)
+		if err != nil && len(regexes) >= 1 {
+			break
+		}
+		if err != nil {
+			return nil, fmt.Errorf("Couldn't fetch the subzones for predefined regex: %s", err.Error())
+		}
+
+		if !strings.Contains(txts[0], "re=") {
+			return nil, fmt.Errorf("Couldn't find the re= field in records: %s", err.Error())
+		}
+
+		// Extract the re= field from record and add it to the map
+		regexes[fmt.Sprintf("%d.%s", i, p.req.Host)] = RegexRecord{
+			TXT:   txts[0],
+			Regex: strings.TrimPrefix(strings.Split(txts[0][strings.Index(txts[0], "re="):], ";")[0], "re="),
+		}
+	}
+
+	return regexes, nil
+}
+
 // zoneFromPath generates a DNS zone with the given request's path and host
 // It will use custom regex to parse the path if it's provided in
 // the given record.
@@ -96,26 +184,35 @@ func zoneFromPath(r *http.Request, rec record) (string, int, []string, error) {
 
 	path = fmt.Sprintf("%s?%s", path, r.URL.RawQuery)
 
+	// Normalize the path to follow RFC1034 rules
 	if strings.ContainsAny(path, ".") {
 		path = strings.Replace(path, ".", "-", -1)
 	}
+
 	pathSubmatchs := PathRegex.FindAllStringSubmatch(path, -1)
 	if rec.Re != "" {
+		// Compile the record regex and find path submatches
 		CustomRegex, err := regexp.Compile(rec.Re)
 		if err != nil {
 			log.Printf("<%s> [txtdirect]: the given regex doesn't work as expected: %s", time.Now().String(), rec.Re)
 		}
 		pathSubmatchs = CustomRegex.FindAllStringSubmatch(path, -1)
+
+		// Only generate the zone if the custom regex contains a group
 		if GroupRegex.MatchString(rec.Re) {
+			//
 			pathSlice := []string{}
 			unordered := make(map[string]string)
 			for _, item := range pathSubmatchs[0] {
 				pathSlice = append(pathSlice, item)
 			}
+
+			// Order the path slice using groups order in custom regex
 			order := GroupOrderRegex.FindAllStringSubmatch(rec.Re, -1)
 			for i, group := range order {
 				unordered[group[1]] = pathSlice[i+1]
 			}
+
 			url := sortMap(unordered)
 			*r = *r.WithContext(context.WithValue(r.Context(), "regexMatches", unordered))
 			reverse(url)

txtdirect.go

@@ -84,9 +84,7 @@ func customResolver(c Config) net.Resolver {
 	}
 }
 
-// query checks the given zone using net.LookupTXT to
-// find TXT records in that zone
-func query(zone string, ctx context.Context, c Config) ([]string, error) {
+func absoluteZone(zone string) string {
 	// Removes port from zone
 	if strings.Contains(zone, ":") {
 		zoneSlice := strings.Split(zone, ":")
@@ -97,25 +95,30 @@ func query(zone string, ctx context.Context, c Config) ([]string, error) {
 		zone = strings.Join([]string{basezone, zone}, ".")
 	}
 
-	// Use absolute zone
-	var absoluteZone string
 	if strings.HasSuffix(zone, ".") {
-		absoluteZone = zone
-	} else {
-		absoluteZone = strings.Join([]string{zone, "."}, "")
+		return zone
 	}
 
+	return strings.Join([]string{zone, "."}, "")
+}
+
+// query checks the given zone using net.LookupTXT to
+// find TXT records in that zone
+func query(zone string, ctx context.Context, c Config) ([]string, error) {
 	var txts []string
 	var err error
 	if c.Resolver != "" {
 		net := customResolver(c)
-		txts, err = net.LookupTXT(ctx, absoluteZone)
+		txts, err = net.LookupTXT(ctx, absoluteZone(zone))
 	} else {
-		txts, err = net.LookupTXT(absoluteZone)
+		txts, err = net.LookupTXT(absoluteZone(zone))
 	}
 	if err != nil {
 		return nil, fmt.Errorf("could not get TXT record: %s", err)
 	}
+	if txts[0] == "" {
+		return nil, fmt.Errorf("TXT record doesn't exist or is empty")
+	}
 	return txts, nil
 }
 
@@ -206,14 +209,25 @@ func Redirect(w http.ResponseWriter, r *http.Request, c Config) error {
 			return path.RedirectRoot()
 		}
 
-		if path.path != "" {
+		if path.path != "" && rec.Re != "record" {
 			record := path.Redirect()
 			// It means fallback got triggered, If record is nil
 			if record == nil {
 				return nil
 			}
 			rec = *record
 		}
+
+		// Use predefined regexes if custom regex is set to "record"
+		if path.rec.Re == "record" {
+			record, err := path.SpecificRecord()
+			if err != nil {
+				log.Printf("[txtdirect]: Fallback is triggered because redirect to the most specific match failed: %s", err.Error())
+				fallback(path.rw, path.req, "to", path.rec.Code, path.c)
+				return nil
+			}
+			rec = *record
+		}
 	}
 
 	if rec.Type == "proxy" {

txtdirect_test.go

@@ -45,6 +45,10 @@ var txts = map[string]string{
 	"_redirect.noto.path.e2e.test.":      "v=txtv0;type=path",
 	"_redirect.noroot.path.e2e.test.":    "v=txtv0;to=https://noroot.fallback.path.test;type=path;code=302",
 	"_redirect.metapath.e2e.test.":       "v=txtv0;type=path",
+	"_redirect.regex.path.e2e.test.":     "v=txtv0;re=record;to=https://example.com;type=path",
+	"_redirect.1.regex.path.e2e.test.":   "v=txtv0;re=\\/test1;to=https://example.com/first/predefined{1};type=host",
+	"_redirect.2.regex.path.e2e.test.":   "v=txtv0;re=\\/test1\\/test2;to=https://example.com/second/predefined{1};type=host",
+
 	// type=gometa
 	"_redirect.pkg.txtdirect.test.":           "v=txtv0;to=https://example.com/example/example;type=gometa;vcs=git",
 	"_redirect.pkgweb.metapath.e2e.test.":     "v=txtv0;to=https://example.com/example/example;type=gometa;website=https://godoc.org/go.txtdirect.org/txtdirect",
@@ -263,6 +267,16 @@ func TestRedirectE2e(t *testing.T) {
 			enable:   []string{"host"},
 			referer:  true,
 		},
+		{
+			url:      "https://regex.path.e2e.test/test1",
+			expected: "https://example.com/first/predefined/test1",
+			enable:   []string{"host", "path"},
+		},
+		{
+			url:      "https://regex.path.e2e.test/test1/test2",
+			expected: "https://example.com/second/predefined/test1/test2",
+			enable:   []string{"host", "path"},
+		},
 	}
 	for _, test := range tests {
 		req := httptest.NewRequest("GET", test.url, nil)