traefik/pkg/middlewares/extractor.go

63 lines
1.9 KiB
Go
Raw Normal View History

package middlewares
import (
"context"
"errors"
"fmt"
"net/http"
2022-11-21 17:36:05 +00:00
"github.com/rs/zerolog/log"
"github.com/traefik/traefik/v2/pkg/config/dynamic"
2022-11-21 17:36:05 +00:00
"github.com/vulcand/oxy/v2/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{},
}
}
2022-11-21 17:36:05 +00:00
logger := log.Ctx(ctx)
if sourceMatcher.IPStrategy != nil {
strategy, err := sourceMatcher.IPStrategy.Get()
if err != nil {
return nil, err
}
2022-11-21 17:36:05 +00:00
logger.Debug().Msg("Using IPStrategy")
return utils.ExtractorFunc(func(req *http.Request) (string, int64, error) {
return strategy.GetIP(req), 1, nil
}), nil
}
if sourceMatcher.RequestHeaderName != "" {
2022-11-21 17:36:05 +00:00
logger.Debug().Msg("Using RequestHeaderName")
return utils.NewExtractor(fmt.Sprintf("request.header.%s", sourceMatcher.RequestHeaderName))
}
if sourceMatcher.RequestHost {
2022-11-21 17:36:05 +00:00
logger.Debug().Msg("Using RequestHost")
return utils.NewExtractor("request.host")
}
return nil, errors.New("no SourceCriterion criterion defined")
}