mirror of
https://github.com/mpolden/echoip.git
synced 2025-06-26 20:57:51 +02:00
Add support for multiple trusted headers
This commit is contained in:
18
http/http.go
18
http/http.go
@ -22,7 +22,7 @@ const (
|
||||
|
||||
type Server struct {
|
||||
Template string
|
||||
IPHeader string
|
||||
IPHeaders []string
|
||||
LookupAddr func(net.IP) (string, error)
|
||||
LookupPort func(net.IP, uint64) error
|
||||
db database.Client
|
||||
@ -47,8 +47,14 @@ func New(db database.Client) *Server {
|
||||
return &Server{db: db}
|
||||
}
|
||||
|
||||
func ipFromRequest(header string, r *http.Request) (net.IP, error) {
|
||||
remoteIP := r.Header.Get(header)
|
||||
func ipFromRequest(headers []string, r *http.Request) (net.IP, error) {
|
||||
remoteIP := ""
|
||||
for _, header := range headers {
|
||||
remoteIP = r.Header.Get(header)
|
||||
if remoteIP != "" {
|
||||
break
|
||||
}
|
||||
}
|
||||
if remoteIP == "" {
|
||||
host, _, err := net.SplitHostPort(r.RemoteAddr)
|
||||
if err != nil {
|
||||
@ -64,7 +70,7 @@ func ipFromRequest(header string, r *http.Request) (net.IP, error) {
|
||||
}
|
||||
|
||||
func (s *Server) newResponse(r *http.Request) (Response, error) {
|
||||
ip, err := ipFromRequest(s.IPHeader, r)
|
||||
ip, err := ipFromRequest(s.IPHeaders, r)
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
@ -91,7 +97,7 @@ func (s *Server) newPortResponse(r *http.Request) (PortResponse, error) {
|
||||
if err != nil || port < 1 || port > 65355 {
|
||||
return PortResponse{Port: port}, fmt.Errorf("invalid port: %d", port)
|
||||
}
|
||||
ip, err := ipFromRequest(s.IPHeader, r)
|
||||
ip, err := ipFromRequest(s.IPHeaders, r)
|
||||
if err != nil {
|
||||
return PortResponse{Port: port}, err
|
||||
}
|
||||
@ -104,7 +110,7 @@ func (s *Server) newPortResponse(r *http.Request) (PortResponse, error) {
|
||||
}
|
||||
|
||||
func (s *Server) CLIHandler(w http.ResponseWriter, r *http.Request) *appError {
|
||||
ip, err := ipFromRequest(s.IPHeader, r)
|
||||
ip, err := ipFromRequest(s.IPHeaders, r)
|
||||
if err != nil {
|
||||
return internalServerError(err)
|
||||
}
|
||||
|
@ -149,16 +149,17 @@ func TestJSONHandlers(t *testing.T) {
|
||||
|
||||
func TestIPFromRequest(t *testing.T) {
|
||||
var tests = []struct {
|
||||
remoteAddr string
|
||||
headerKey string
|
||||
headerValue string
|
||||
trustedHeader string
|
||||
out string
|
||||
remoteAddr string
|
||||
headerKey string
|
||||
headerValue string
|
||||
trustedHeaders []string
|
||||
out string
|
||||
}{
|
||||
{"127.0.0.1:9999", "", "", "", "127.0.0.1"}, // No header given
|
||||
{"127.0.0.1:9999", "X-Real-IP", "1.3.3.7", "", "127.0.0.1"}, // Trusted header is empty
|
||||
{"127.0.0.1:9999", "X-Real-IP", "1.3.3.7", "X-Foo-Bar", "127.0.0.1"}, // Trusted header does not match
|
||||
{"127.0.0.1:9999", "X-Real-IP", "1.3.3.7", "X-Real-IP", "1.3.3.7"}, // Trusted header matches
|
||||
{"127.0.0.1:9999", "", "", nil, "127.0.0.1"}, // No header given
|
||||
{"127.0.0.1:9999", "X-Real-IP", "1.3.3.7", nil, "127.0.0.1"}, // Trusted header is empty
|
||||
{"127.0.0.1:9999", "X-Real-IP", "1.3.3.7", []string{"X-Foo-Bar"}, "127.0.0.1"}, // Trusted header does not match
|
||||
{"127.0.0.1:9999", "X-Real-IP", "1.3.3.7", []string{"X-Real-IP", "X-Forwarded-For"}, "1.3.3.7"}, // Trusted header matches
|
||||
{"127.0.0.1:9999", "X-Forwarded-For", "1.3.3.7", []string{"X-Real-IP", "X-Forwarded-For"}, "1.3.3.7"}, // Second trusted header matches
|
||||
}
|
||||
for _, tt := range tests {
|
||||
r := &http.Request{
|
||||
@ -166,7 +167,7 @@ func TestIPFromRequest(t *testing.T) {
|
||||
Header: http.Header{},
|
||||
}
|
||||
r.Header.Add(tt.headerKey, tt.headerValue)
|
||||
ip, err := ipFromRequest(tt.trustedHeader, r)
|
||||
ip, err := ipFromRequest(tt.trustedHeaders, r)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
Reference in New Issue
Block a user