Merge pull request #769 from luraproject/sc-691/customizable-backend-filtering

Customizable backend filtering

Author: kpacha

proxy/merging.go

@@ -28,7 +28,7 @@ func NewMergeDataMiddleware(logger logging.Logger, endpointConfig *config.Endpoi
 	}
 	serviceTimeout := time.Duration(85*endpointConfig.Timeout.Nanoseconds()/100) * time.Nanosecond
 	combiner := getResponseCombiner(endpointConfig.ExtraConfig)
-	isSequential, propagatedParams := sequentialMergerConfig(endpointConfig)
+	isSequential, sequentialReplacements := sequentialMergerConfig(endpointConfig)
 
 	logger.Debug(
 		fmt.Sprintf(
@@ -40,6 +40,8 @@ func NewMergeDataMiddleware(logger logging.Logger, endpointConfig *config.Endpoi
 		),
 	)
 
+	bfFactory := backendFiltererFactory.filtererFactory
+
 	return func(next ...Proxy) Proxy {
 		if len(next) != totalBackends {
 			// we leave the panic here, because we do not want to continue
@@ -50,66 +52,61 @@ func NewMergeDataMiddleware(logger logging.Logger, endpointConfig *config.Endpoi
 		}
 		reqClone := func(r *Request) *Request { res := r.Clone(); return &res }
 
+		filters, err := bfFactory(endpointConfig)
+		if err != nil {
+			logger.Error(fmt.Sprintf("[ENDPOINT: %s]%s %s", endpointConfig.Endpoint, backendFiltererFactory.logPrefix, err))
+			return func(_ context.Context, _ *Request) (*Response, error) { return nil, err }
+		}
+
 		if hasUnsafeBackends(endpointConfig) {
 			reqClone = CloneRequest
 		}
 
 		if !isSequential {
-			return parallelMerge(reqClone, serviceTimeout, combiner, next...)
+			return parallelMerge(reqClone, serviceTimeout, combiner, filters, next...)
 		}
 
-		sequentialReplacements := make([][]sequentialBackendReplacement, totalBackends)
+		return sequentialMerge(reqClone, serviceTimeout, combiner, sequentialReplacements, filters, next...)
+	}
+}
 
-		var rePropagatedParams = regexp.MustCompile(`[Rr]esp(\d+)_?([\w-.]+)?`)
-		var reUrlPatterns = regexp.MustCompile(`\{\{\.Resp(\d+)_([\w-.]+)\}\}`)
-		destKeyGenerator := func(i string, t string) string {
-			key := "Resp" + i
-			if t != "" {
-				key += "_" + t
-			}
-			return key
-		}
+// BackendFiltererFactory is a factory function that returns a list of BackendFilterer
+// based on the provided EndpointConfig.
+// The returned list must be sorted by the backend index.
+// The list can contain nil values, which means that the backend in that index is untouched.
+type BackendFiltererFactory func(*config.EndpointConfig) ([]BackendFilterer, error)
 
-		for i, b := range endpointConfig.Backend {
-			for _, match := range reUrlPatterns.FindAllStringSubmatch(b.URLPattern, -1) {
-				if len(match) > 1 {
-					backendIndex, err := strconv.Atoi(match[1])
-					if err != nil {
-						continue
-					}
+// BackendFilterer evalutes the request and returns true if the backend should be used,
+// otherwise the backend is skipped in both normal and sequential merging.
+// If the backend is skipped, the response will not be merged into the final response.
+type BackendFilterer func(*Request) bool
 
-					sequentialReplacements[i] = append(sequentialReplacements[i], sequentialBackendReplacement{
-						backendIndex: backendIndex,
-						destination:  destKeyGenerator(match[1], match[2]),
-						source:       strings.Split(match[2], "."),
-						fullResponse: match[2] == "",
-					})
-				}
-			}
+func defaultBackendFiltererFactory(_ *config.EndpointConfig) ([]BackendFilterer, error) {
+	return []BackendFilterer{}, nil
+}
 
-			if i > 0 {
-				for _, p := range propagatedParams {
-					for _, match := range rePropagatedParams.FindAllStringSubmatch(p, -1) {
-						if len(match) > 1 {
-							backendIndex, err := strconv.Atoi(match[1])
-							if err != nil || backendIndex >= totalBackends {
-								continue
-							}
+type backendFiltererRegistry struct {
+	logPrefix       string
+	filtererFactory BackendFiltererFactory
+}
 
-							sequentialReplacements[i] = append(sequentialReplacements[i], sequentialBackendReplacement{
-								backendIndex: backendIndex,
-								destination:  destKeyGenerator(match[1], match[2]),
-								source:       strings.Split(match[2], "."),
-								fullResponse: match[2] == "",
-							})
-						}
-					}
-				}
-			}
-		}
+var backendFiltererFactory = backendFiltererRegistry{
+	filtererFactory: defaultBackendFiltererFactory,
+}
 
-		return sequentialMerge(reqClone, sequentialReplacements, serviceTimeout, combiner, next...)
-	}
+// RegisterBackendFiltererFactory registers a new backend filterer factory
+// to be used by the merging middleware.
+// This factory is used to create a list of BackendFilterer
+// functions that will be used to filter backends based on the request.
+// Important: this function should be called everytime the middleware is created.
+func RegisterBackendFiltererFactory(logPrefix string, f BackendFiltererFactory) {
+	backendFiltererFactory.logPrefix = logPrefix
+	backendFiltererFactory.filtererFactory = f
+}
+
+func ResetBackendFiltererFactory() {
+	backendFiltererFactory.logPrefix = ""
+	backendFiltererFactory.filtererFactory = defaultBackendFiltererFactory
 }
 
 type sequentialBackendReplacement struct {
@@ -119,9 +116,12 @@ type sequentialBackendReplacement struct {
 	fullResponse bool
 }
 
-func sequentialMergerConfig(cfg *config.EndpointConfig) (bool, []string) {
+func sequentialMergerConfig(cfg *config.EndpointConfig) (bool, [][]sequentialBackendReplacement) { // skipcq: GO-R1005
 	enabled := false
+	totalBackends := len(cfg.Backend)
+	sequentialReplacements := make([][]sequentialBackendReplacement, totalBackends)
 	var propagatedParams []string
+
 	if v, ok := cfg.ExtraConfig[Namespace]; ok {
 		if e, ok := v.(map[string]interface{}); ok {
 			if v, ok := e[isSequentialKey]; ok {
@@ -137,7 +137,54 @@ func sequentialMergerConfig(cfg *config.EndpointConfig) (bool, []string) {
 			}
 		}
 	}
-	return enabled, propagatedParams
+	var rePropagatedParams = regexp.MustCompile(`[Rr]esp(\d+)_?([\w-.]+)?`)
+	var reUrlPatterns = regexp.MustCompile(`\{\{\.Resp(\d+)_([\w-.]+)\}\}`)
+	destKeyGenerator := func(i string, t string) string {
+		key := "Resp" + i
+		if t != "" {
+			key += "_" + t
+		}
+		return key
+	}
+
+	for i, b := range cfg.Backend {
+		for _, match := range reUrlPatterns.FindAllStringSubmatch(b.URLPattern, -1) {
+			if len(match) > 1 {
+				backendIndex, err := strconv.Atoi(match[1])
+				if err != nil {
+					continue
+				}
+
+				sequentialReplacements[i] = append(sequentialReplacements[i], sequentialBackendReplacement{
+					backendIndex: backendIndex,
+					destination:  destKeyGenerator(match[1], match[2]),
+					source:       strings.Split(match[2], "."),
+					fullResponse: match[2] == "",
+				})
+			}
+		}
+
+		if i > 0 {
+			for _, p := range propagatedParams {
+				for _, match := range rePropagatedParams.FindAllStringSubmatch(p, -1) {
+					if len(match) > 1 {
+						backendIndex, err := strconv.Atoi(match[1])
+						if err != nil || backendIndex >= totalBackends {
+							continue
+						}
+
+						sequentialReplacements[i] = append(sequentialReplacements[i], sequentialBackendReplacement{
+							backendIndex: backendIndex,
+							destination:  destKeyGenerator(match[1], match[2]),
+							source:       strings.Split(match[2], "."),
+							fullResponse: match[2] == "",
+						})
+					}
+				}
+			}
+		}
+	}
+	return enabled, sequentialReplacements
 }
 
 func hasUnsafeBackends(cfg *config.EndpointConfig) bool {
@@ -154,19 +201,32 @@ func hasUnsafeBackends(cfg *config.EndpointConfig) bool {
 	return false
 }
 
-func parallelMerge(reqCloner func(*Request) *Request, timeout time.Duration, rc ResponseCombiner, next ...Proxy) Proxy {
+func parallelMerge(
+	reqCloner func(*Request) *Request,
+	timeout time.Duration,
+	rc ResponseCombiner,
+	filters []BackendFilterer,
+	next ...Proxy,
+) Proxy {
 	return func(ctx context.Context, request *Request) (*Response, error) {
 		localCtx, cancel := context.WithTimeout(ctx, timeout)
 
-		parts := make(chan *Response, len(next))
-		failed := make(chan error, len(next))
+		proxyCount := len(next)
+		filterCount := len(filters)
+
+		parts := make(chan *Response, proxyCount)
+		failed := make(chan error, proxyCount)
 
-		for _, n := range next {
+		for i, n := range next {
+			if (i < filterCount) && (filters[i] != nil) && !filters[i](request) {
+				proxyCount--
+				continue
+			}
 			go requestPart(localCtx, n, reqCloner(request), parts, failed)
 		}
 
-		acc := newIncrementalMergeAccumulator(len(next), rc)
-		for i := 0; i < len(next); i++ {
+		acc := newIncrementalMergeAccumulator(proxyCount, rc)
+		for i := 0; i < proxyCount; i++ {
 			select {
 			case err := <-failed:
 				acc.Merge(nil, err)
@@ -181,10 +241,18 @@ func parallelMerge(reqCloner func(*Request) *Request, timeout time.Duration, rc
 	}
 }
 
-func sequentialMerge(reqCloner func(*Request) *Request, sequentialReplacements [][]sequentialBackendReplacement, timeout time.Duration, rc ResponseCombiner, next ...Proxy) Proxy { // skipcq: GO-R1005
+func sequentialMerge( // skipcq: GO-R1005
+	reqCloner func(*Request) *Request,
+	timeout time.Duration,
+	rc ResponseCombiner,
+	sequentialReplacements [][]sequentialBackendReplacement,
+	filters []BackendFilterer,
+	next ...Proxy,
+) Proxy {
 	return func(ctx context.Context, request *Request) (*Response, error) {
 		localCtx, cancel := context.WithTimeout(ctx, timeout)
 
+		filterCount := len(filters)
 		parts := make([]*Response, len(next))
 		out := make(chan *Response, 1)
 		errCh := make(chan error, 1)
@@ -270,6 +338,12 @@ func sequentialMerge(reqCloner func(*Request) *Request, sequentialReplacements [
 				}
 			}
 
+			if (i < filterCount) && (filters[i] != nil) && !filters[i](request) {
+				parts[i] = &Response{IsComplete: true, Data: make(map[string]interface{})}
+				acc.pending--
+				continue
+			}
+
 			sequentialRequestPart(localCtx, n, reqCloner(request), out, errCh)
 
 			select {
@@ -335,7 +409,7 @@ func (i *incrementalMergeAccumulator) Result() (*Response, error) {
 		return nil, newMergeError(i.errs)
 	}
 
-	if i.pending != 0 || len(i.errs) != 0 {
+	if i.pending > 0 || len(i.errs) > 0 {
 		i.data.IsComplete = false
 	}
 	return i.data, newMergeError(i.errs)