add support for x-forwaded-for and x-forwarded-host

Author: Michal Witkowski

runtime/context.go

@@ -4,12 +4,16 @@ import (
 	"net/http"
 	"strings"
 
+	"net"
+
 	"golang.org/x/net/context"
 	"google.golang.org/grpc/metadata"
 )
 
 const metadataHeaderPrefix = "Grpc-Metadata-"
 const metadataTrailerPrefix = "Grpc-Trailer-"
+const xForwardedFor = "X-Forwarded-For"
+const xForwardedHost = "X-Forwarded-Host"
 
 /*
 AnnotateContext adds context information such as metadata from the request.
@@ -22,14 +26,27 @@ func AnnotateContext(ctx context.Context, req *http.Request) context.Context {
 	for key, vals := range req.Header {
 		for _, val := range vals {
 			if key == "Authorization" {
-				pairs = append(pairs, key, val)
+				pairs = append(pairs, "authorization", val)
 				continue
 			}
 			if strings.HasPrefix(key, metadataHeaderPrefix) {
 				pairs = append(pairs, key[len(metadataHeaderPrefix):], val)
 			}
 		}
 	}
+	if host := req.Header.Get(xForwardedHost); host != "" {
+		pairs = append(pairs, strings.ToLower(xForwardedHost), host)
+	} else if req.Host != "" {
+		pairs = append(pairs, strings.ToLower(xForwardedHost), req.Host)
+	}
+	remoteIp, _, err := net.SplitHostPort(req.RemoteAddr)
+	if err == nil {
+		if req.Header.Get(xForwardedFor) == "" {
+			pairs = append(pairs, strings.ToLower(xForwardedFor), remoteIp)
+		} else {
+			pairs = append(pairs, strings.ToLower(xForwardedFor), req.Header.Get(xForwardedFor)+", "+remoteIp)
+		}
+	}
 
 	if len(pairs) == 0 {
 		return ctx

runtime/context_test.go

@@ -10,27 +10,45 @@ import (
 	"google.golang.org/grpc/metadata"
 )
 
-func TestAnnotateContext(t *testing.T) {
+const (
+	emptyForwardMetaCount = 2
+)
+
+func TestAnnotateContext_WorksWithEmpty(t *testing.T) {
 	ctx := context.Background()
 
-	request, err := http.NewRequest("GET", "http://localhost", nil)
+	request, err := http.NewRequest("GET", "http://www.example.com", nil)
 	if err != nil {
-		t.Fatalf("http.NewRequest(%q, %q, nil) failed with %v; want success", "GET", "http://localhost", err)
+		t.Fatalf("http.NewRequest(%q, %q, nil) failed with %v; want success", "GET", "http://www.example.com", err)
 	}
+	// Make sure we set a remote.
+	request.RemoteAddr = "192.168.0.1:12345"
+
 	request.Header.Add("Some-Irrelevant-Header", "some value")
 	annotated := runtime.AnnotateContext(ctx, request)
-	if annotated != ctx {
-		t.Errorf("AnnotateContext(ctx, request) = %v; want %v", annotated, ctx)
+	md, ok := metadata.FromContext(annotated)
+	if !ok || len(md) != emptyForwardMetaCount {
+		t.Errorf("Expected 2 metadata items in context; got %v", md)
 	}
+}
 
+func TestAnnotateContext_ForwardsGrpcMetadata(t *testing.T) {
+	ctx := context.Background()
+	request, err := http.NewRequest("GET", "http://www.example.com", nil)
+	if err != nil {
+		t.Fatalf("http.NewRequest(%q, %q, nil) failed with %v; want success", "GET", "http://www.example.com", err)
+	}
+	request.RemoteAddr = "192.168.0.1:12345"
+
+	request.Header.Add("Some-Irrelevant-Header", "some value")
 	request.Header.Add("Grpc-Metadata-FooBar", "Value1")
 	request.Header.Add("Grpc-Metadata-Foo-BAZ", "Value2")
 	request.Header.Add("Grpc-Metadata-foo-bAz", "Value3")
 	request.Header.Add("Authorization", "Token 1234567890")
-	annotated = runtime.AnnotateContext(ctx, request)
+	annotated := runtime.AnnotateContext(ctx, request)
 	md, ok := metadata.FromContext(annotated)
-	if !ok || len(md) != 3 {
-		t.Errorf("Expected 3 metadata items in context; got %v", md)
+	if !ok || len(md) != emptyForwardMetaCount+3 {
+		t.Errorf("Expected 5 metadata items in context; got %v", md)
 	}
 	if got, want := md["foobar"], []string{"Value1"}; !reflect.DeepEqual(got, want) {
 		t.Errorf(`md["foobar"] = %q; want %q`, got, want)
@@ -42,3 +60,26 @@ func TestAnnotateContext(t *testing.T) {
 		t.Errorf(`md["authorization"] = %q want %q`, got, want)
 	}
 }
+
+func TestAnnotateContext_XForwardedFor(t *testing.T) {
+	ctx := context.Background()
+	request, err := http.NewRequest("GET", "http://bar.foo.example.com", nil)
+	if err != nil {
+		t.Fatalf("http.NewRequest(%q, %q, nil) failed with %v; want success", "GET", "http://bar.foo.example.com", err)
+	}
+	request.Header.Add("X-Forwarded-For", "192.168.0.100") // client
+	request.RemoteAddr = "8.8.8.8:12345"                   // proxy
+
+	annotated := runtime.AnnotateContext(ctx, request)
+	md, ok := metadata.FromContext(annotated)
+	if !ok || len(md) != emptyForwardMetaCount {
+		t.Errorf("Expected 2 metadata items in context; got %v", md)
+	}
+	if got, want := md["x-forwarded-host"], []string{"bar.foo.example.com"}; !reflect.DeepEqual(got, want) {
+		t.Errorf("md[\"host\"] = %v; want %v", got, want)
+	}
+	// Note: it must be in order client, proxy1, proxy2
+	if got, want := md["x-forwarded-for"], []string{"192.168.0.100, 8.8.8.8"}; !reflect.DeepEqual(got, want) {
+		t.Errorf("md[\"x-forwarded-for\"] = %v want %v", got, want)
+	}
+}