source: code/trunk/vendor/github.com/pires/go-proxyproto/policy.go@ 822

package proxyproto

import (
        "fmt"
        "net"
        "strings"
)

// PolicyFunc can be used to decide whether to trust the PROXY info from
// upstream. If set, the connecting address is passed in as an argument.
//
// See below for the different policies.
//
// In case an error is returned the connection is denied.
type PolicyFunc func(upstream net.Addr) (Policy, error)

// Policy defines how a connection with a PROXY header address is treated.
type Policy int

const (
        // USE address from PROXY header
        USE Policy = iota
        // IGNORE address from PROXY header, but accept connection
        IGNORE
        // REJECT connection when PROXY header is sent
        // Note: even though the first read on the connection returns an error if
        // a PROXY header is present, subsequent reads do not. It is the task of
        // the code using the connection to handle that case properly.
        REJECT
        // REQUIRE connection to send PROXY header, reject if not present
        // Note: even though the first read on the connection returns an error if
        // a PROXY header is not present, subsequent reads do not. It is the task
        // of the code using the connection to handle that case properly.
        REQUIRE
        // SKIP accepts a connection without requiring the PROXY header
        // Note: an example usage can be found in the SkipProxyHeaderForCIDR
        // function.
        SKIP
)

// SkipProxyHeaderForCIDR returns a PolicyFunc which can be used to accept a
// connection from a skipHeaderCIDR without requiring a PROXY header, e.g.
// Kubernetes pods local traffic. The def is a policy to use when an upstream
// address doesn't match the skipHeaderCIDR.
func SkipProxyHeaderForCIDR(skipHeaderCIDR *net.IPNet, def Policy) PolicyFunc {
        return func(upstream net.Addr) (Policy, error) {
                ip, err := ipFromAddr(upstream)
                if err != nil {
                        return def, err
                }

                if skipHeaderCIDR != nil && skipHeaderCIDR.Contains(ip) {
                        return SKIP, nil
                }

                return def, nil
        }
}

// WithPolicy adds given policy to a connection when passed as option to NewConn()
func WithPolicy(p Policy) func(*Conn) {
        return func(c *Conn) {
                c.ProxyHeaderPolicy = p
        }
}

// LaxWhiteListPolicy returns a PolicyFunc which decides whether the
// upstream ip is allowed to send a proxy header based on a list of allowed
// IP addresses and IP ranges. In case upstream IP is not in list the proxy
// header will be ignored. If one of the provided IP addresses or IP ranges
// is invalid it will return an error instead of a PolicyFunc.
func LaxWhiteListPolicy(allowed []string) (PolicyFunc, error) {
        allowFrom, err := parse(allowed)
        if err != nil {
                return nil, err
        }

        return whitelistPolicy(allowFrom, IGNORE), nil
}

// MustLaxWhiteListPolicy returns a LaxWhiteListPolicy but will panic if one
// of the provided IP addresses or IP ranges is invalid.
func MustLaxWhiteListPolicy(allowed []string) PolicyFunc {
        pfunc, err := LaxWhiteListPolicy(allowed)
        if err != nil {
                panic(err)
        }

        return pfunc
}

// StrictWhiteListPolicy returns a PolicyFunc which decides whether the
// upstream ip is allowed to send a proxy header based on a list of allowed
// IP addresses and IP ranges. In case upstream IP is not in list reading on
// the connection will be refused on the first read. Please note: subsequent
// reads do not error. It is the task of the code using the connection to
// handle that case properly. If one of the provided IP addresses or IP
// ranges is invalid it will return an error instead of a PolicyFunc.
func StrictWhiteListPolicy(allowed []string) (PolicyFunc, error) {
        allowFrom, err := parse(allowed)
        if err != nil {
                return nil, err
        }

        return whitelistPolicy(allowFrom, REJECT), nil
}

// MustStrictWhiteListPolicy returns a StrictWhiteListPolicy but will panic
// if one of the provided IP addresses or IP ranges is invalid.
func MustStrictWhiteListPolicy(allowed []string) PolicyFunc {
        pfunc, err := StrictWhiteListPolicy(allowed)
        if err != nil {
                panic(err)
        }

        return pfunc
}

func whitelistPolicy(allowed []func(net.IP) bool, def Policy) PolicyFunc {
        return func(upstream net.Addr) (Policy, error) {
                upstreamIP, err := ipFromAddr(upstream)
                if err != nil {
                        // something is wrong with the source IP, better reject the connection
                        return REJECT, err
                }

                for _, allowFrom := range allowed {
                        if allowFrom(upstreamIP) {
                                return USE, nil
                        }
                }

                return def, nil
        }
}

func parse(allowed []string) ([]func(net.IP) bool, error) {
        a := make([]func(net.IP) bool, len(allowed))
        for i, allowFrom := range allowed {
                if strings.LastIndex(allowFrom, "/") > 0 {
                        _, ipRange, err := net.ParseCIDR(allowFrom)
                        if err != nil {
                                return nil, fmt.Errorf("proxyproto: given string %q is not a valid IP range: %v", allowFrom, err)
                        }

                        a[i] = ipRange.Contains
                } else {
                        allowed := net.ParseIP(allowFrom)
                        if allowed == nil {
                                return nil, fmt.Errorf("proxyproto: given string %q is not a valid IP address", allowFrom)
                        }

                        a[i] = allowed.Equal
                }
        }

        return a, nil
}

func ipFromAddr(upstream net.Addr) (net.IP, error) {
        upstreamString, _, err := net.SplitHostPort(upstream.String())
        if err != nil {
                return nil, err
        }

        upstreamIP := net.ParseIP(upstreamString)
        if nil == upstreamIP {
                return nil, fmt.Errorf("proxyproto: invalid IP address")
        }

        return upstreamIP, nil
}