traefik/pkg/middlewares/extractor.go
Jean-Baptiste Doumenjou c0f1e74bed
chore: move to Traefik organization.
Co-authored-by: Romain <rtribotte@users.noreply.github.com>
2020-09-16 15:46:04 +02:00

62 lines
1.9 KiB
Go

package middlewares
import (
"context"
"errors"
"fmt"
"net/http"
"github.com/traefik/traefik/v2/pkg/config/dynamic"
"github.com/traefik/traefik/v2/pkg/log"
"github.com/vulcand/oxy/utils"
)
// GetSourceExtractor returns the SourceExtractor function corresponding to the given sourceMatcher.
// It defaults to a RemoteAddrStrategy IPStrategy if need be.
// It returns an error if more than one source criterion is provided.
func GetSourceExtractor(ctx context.Context, sourceMatcher *dynamic.SourceCriterion) (utils.SourceExtractor, error) {
if sourceMatcher != nil {
if sourceMatcher.IPStrategy != nil && sourceMatcher.RequestHeaderName != "" {
return nil, errors.New("iPStrategy and RequestHeaderName are mutually exclusive")
}
if sourceMatcher.IPStrategy != nil && sourceMatcher.RequestHost {
return nil, errors.New("iPStrategy and RequestHost are mutually exclusive")
}
if sourceMatcher.RequestHeaderName != "" && sourceMatcher.RequestHost {
return nil, errors.New("requestHost and RequestHeaderName are mutually exclusive")
}
}
if sourceMatcher == nil ||
sourceMatcher.IPStrategy == nil &&
sourceMatcher.RequestHeaderName == "" && !sourceMatcher.RequestHost {
sourceMatcher = &dynamic.SourceCriterion{
IPStrategy: &dynamic.IPStrategy{},
}
}
logger := log.FromContext(ctx)
if sourceMatcher.IPStrategy != nil {
strategy, err := sourceMatcher.IPStrategy.Get()
if err != nil {
return nil, err
}
logger.Debug("Using IPStrategy")
return utils.ExtractorFunc(func(req *http.Request) (string, int64, error) {
return strategy.GetIP(req), 1, nil
}), nil
}
if sourceMatcher.RequestHeaderName != "" {
logger.Debug("Using RequestHeaderName")
return utils.NewExtractor(fmt.Sprintf("request.header.%s", sourceMatcher.RequestHeaderName))
}
if sourceMatcher.RequestHost {
logger.Debug("Using RequestHost")
return utils.NewExtractor("request.host")
}
return nil, errors.New("no SourceCriterion criterion defined")
}