Merge pull request #529 from shogo82148/support-forwarded-header

Support forwarded header

go.mod

@@ -2,4 +2,7 @@ module github.com/shogo82148/aws-xray-yasdk-go
 
 go 1.18
 
-require github.com/google/go-cmp v0.6.0
+require (
+	github.com/google/go-cmp v0.6.0
+	github.com/shogo82148/forwarded-header v0.1.0
+)

go.sum

@@ -1,2 +1,4 @@
 github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
 github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
+github.com/shogo82148/forwarded-header v0.1.0 h1:OZX131i0B8GGZR+s5QEZdSEYDTUPkUX8ISIIQE+swKM=
+github.com/shogo82148/forwarded-header v0.1.0/go.mod h1:Ge73zEgn9jrdyf65vWTgDmRZQFQPHoFU0TaILJcvs54=

xrayaws-v2/go.mod

@@ -26,4 +26,5 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.21.5 // indirect
 	github.com/aws/aws-sdk-go-v2/service/sts v1.26.6 // indirect
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
+	github.com/shogo82148/forwarded-header v0.1.0 // indirect
 )

xrayaws-v2/go.sum

@@ -41,6 +41,8 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/shogo82148/aws-xray-yasdk-go v1.7.4 h1:lImR7K/WbPoFqb6BGM/xnpGD2s4JnQ3WV7thWjDdjG0=
 github.com/shogo82148/aws-xray-yasdk-go v1.7.4/go.mod h1:Ou0BdxWM2H5u5FAcTqYUn08XSv+opM56jDULW8MLCRA=
+github.com/shogo82148/forwarded-header v0.1.0 h1:OZX131i0B8GGZR+s5QEZdSEYDTUPkUX8ISIIQE+swKM=
+github.com/shogo82148/forwarded-header v0.1.0/go.mod h1:Ge73zEgn9jrdyf65vWTgDmRZQFQPHoFU0TaILJcvs54=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=

xrayaws/go.mod

@@ -8,4 +8,7 @@ require (
 	github.com/shogo82148/aws-xray-yasdk-go v1.7.4
 )
 
-require github.com/jmespath/go-jmespath v0.4.0 // indirect
+require (
+	github.com/jmespath/go-jmespath v0.4.0 // indirect
+	github.com/shogo82148/forwarded-header v0.1.0 // indirect
+)

xrayaws/go.sum

@@ -12,6 +12,8 @@ github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZb
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/shogo82148/aws-xray-yasdk-go v1.7.4 h1:lImR7K/WbPoFqb6BGM/xnpGD2s4JnQ3WV7thWjDdjG0=
 github.com/shogo82148/aws-xray-yasdk-go v1.7.4/go.mod h1:Ou0BdxWM2H5u5FAcTqYUn08XSv+opM56jDULW8MLCRA=
+github.com/shogo82148/forwarded-header v0.1.0 h1:OZX131i0B8GGZR+s5QEZdSEYDTUPkUX8ISIIQE+swKM=
+github.com/shogo82148/forwarded-header v0.1.0/go.mod h1:Ge73zEgn9jrdyf65vWTgDmRZQFQPHoFU0TaILJcvs54=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 golang.org/x/net v0.17.0 h1:pVaXccu2ozPjCXewfr1S7xza/zcXTity9cCdXQYSjIM=
 golang.org/x/text v0.13.0 h1:ablQoSUd0tRdKxZewP80B+BaqeKJuVhuRxj/dkrun3k=

xrayhttp/handler.go

@@ -12,6 +12,7 @@ import (
 	"github.com/shogo82148/aws-xray-yasdk-go/xray"
 	"github.com/shogo82148/aws-xray-yasdk-go/xray/sampling"
 	"github.com/shogo82148/aws-xray-yasdk-go/xray/schema"
+	forwardedheader "github.com/shogo82148/forwarded-header"
 )
 
 //go:generate go run codegen.go
@@ -133,8 +134,20 @@ func (tracer *httpTracer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 }
 
 func getURL(r *http.Request) string {
+	forwarded, err := forwardedheader.Parse(r.Header.Values("Forwarded"))
+	if err == nil && len(forwarded) > 0 {
+		proto := forwarded[0].Proto
+		host := forwarded[0].Host
+		if proto != "" && host != "" {
+			return proto + "://" + host + r.URL.Path
+		}
+	}
 	proto := r.Header.Get("X-Forwarded-Proto")
-	if proto == "" {
+	if strings.EqualFold(proto, "https") {
+		proto = "https"
+	} else if strings.EqualFold(proto, "http") {
+		proto = "http"
+	} else if proto == "" {
 		if r.TLS != nil {
 			proto = "https"
 		} else {
@@ -145,6 +158,13 @@ func getURL(r *http.Request) string {
 }
 
 func clientIP(r *http.Request) (string, bool) {
+	forwarded, err := forwardedheader.Parse(r.Header.Values("Forwarded"))
+	if err == nil && len(forwarded) > 0 {
+		ip := forwarded[0].For.IP
+		if ip.IsValid() {
+			return ip.String(), true
+		}
+	}
 	forwardedFor := r.Header.Get("X-Forwarded-For")
 	if forwardedFor != "" {
 		if idx := strings.IndexByte(forwardedFor, ','); idx > 0 {

xrayhttp/handler_test.go

@@ -2,10 +2,12 @@ package xrayhttp
 
 import (
 	"context"
+	"crypto/tls"
 	"fmt"
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"net/url"
 	"os"
 	"strings"
 	"testing"
@@ -600,3 +602,107 @@ func TestHandler_Panic(t *testing.T) {
 		t.Errorf("mismatch (-want +got):\n%s", diff)
 	}
 }
+
+func TestGetURL(t *testing.T) {
+	tests := []struct {
+		name string
+		req  *http.Request
+		want string
+	}{
+		{
+			name: "simple",
+			req: &http.Request{
+				Host: "example.com",
+				URL:  &url.URL{Path: "/"},
+			},
+			want: "http://example.com/",
+		},
+		{
+			name: "https",
+			req: &http.Request{
+				Host: "example.com",
+				URL:  &url.URL{Path: "/"},
+				TLS:  &tls.ConnectionState{},
+			},
+			want: "https://example.com/",
+		},
+		{
+			name: "x-forwarded-proto",
+			req: &http.Request{
+				Header: http.Header{"X-Forwarded-Proto": []string{"https"}},
+				Host:   "example.com",
+				URL:    &url.URL{Path: "/"},
+			},
+			want: "https://example.com/",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got := getURL(tt.req)
+			if got != tt.want {
+				t.Errorf("want %s, got %s", tt.want, got)
+			}
+		})
+	}
+}
+
+func TestClientIP(t *testing.T) {
+	tests := []struct {
+		name      string
+		req       *http.Request
+		wantIP    string
+		forwarded bool
+	}{
+		{
+			name: "simple",
+			req: &http.Request{
+				RemoteAddr: "192.0.2.1:48011",
+			},
+			wantIP:    "192.0.2.1",
+			forwarded: false,
+		},
+		{
+			name: "ipv6",
+			req: &http.Request{
+				RemoteAddr: "[2001:db8::1]:48011",
+			},
+			wantIP:    "2001:db8::1",
+			forwarded: false,
+		},
+		{
+			name: "xff",
+			req: &http.Request{
+				Header: http.Header{
+					"X-Forwarded-For": []string{"198.51.100.1"},
+				},
+				RemoteAddr: "192.0.2.1:48011",
+			},
+			wantIP:    "198.51.100.1",
+			forwarded: true,
+		},
+		{
+			name: "forwarded-header",
+			req: &http.Request{
+				Header: http.Header{
+					"Forwarded": []string{"for=198.51.100.1"},
+				},
+				RemoteAddr: "192.0.2.1:48011",
+			},
+			wantIP:    "198.51.100.1",
+			forwarded: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			gotIP, forwarded := clientIP(tt.req)
+			if gotIP != tt.wantIP {
+				t.Errorf("want %s, got %s", tt.wantIP, gotIP)
+			}
+			if forwarded != tt.forwarded {
+				t.Errorf("want %t, got %t", tt.forwarded, forwarded)
+			}
+		})
+	}
+}