github.com/google/martian/v3@v3.3.3/header/forwarded_modifier_test.go (about) 1 // Copyright 2015 Google Inc. All rights reserved. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 package header 16 17 import ( 18 "net/http" 19 "testing" 20 ) 21 22 func TestSetForwardHeaders(t *testing.T) { 23 xfp := "X-Forwarded-Proto" 24 xff := "X-Forwarded-For" 25 xfh := "X-Forwarded-Host" 26 xfu := "X-Forwarded-Url" 27 28 m := NewForwardedModifier() 29 req, err := http.NewRequest("GET", "http://martian.local?key=value", nil) 30 if err != nil { 31 t.Fatalf("http.NewRequest(): got %v, want no error", err) 32 } 33 req.RemoteAddr = "10.0.0.1:8112" 34 35 if m.ModifyRequest(req); err != nil { 36 t.Fatalf("ModifyRequest(): got %v, want no error", err) 37 } 38 39 if got, want := req.Header.Get(xfp), "http"; got != want { 40 t.Errorf("req.Header.Get(%q): got %q, want %q", xfp, got, want) 41 } 42 if got, want := req.Header.Get(xff), "10.0.0.1"; got != want { 43 t.Errorf("req.Header.Get(%q): got %q, want %q", xff, got, want) 44 } 45 if got, want := req.Header.Get(xfh), "martian.local"; got != want { 46 t.Errorf("req.Header.Get(%q): got %q, want %q", xfh, got, want) 47 } 48 if got, want := req.Header.Get(xfu), "http://martian.local?key=value"; got != want { 49 t.Errorf("req.Header.Get(%q): got %q, want %q", xfh, got, want) 50 } 51 52 // Test with existing X-Forwarded-For. 53 req.RemoteAddr = "12.12.12.12" 54 55 if m.ModifyRequest(req); err != nil { 56 t.Fatalf("ModifyRequest(): got %v, want no error", err) 57 } 58 59 if got, want := req.Header.Get(xff), "10.0.0.1, 12.12.12.12"; got != want { 60 t.Errorf("req.Header.Get(%q): got %q, want %q", xff, got, want) 61 } 62 63 // Test that proto, host, and URL headers are preserved if already present. 64 req, err = http.NewRequest("GET", "http://example.com/path?k=v", nil) 65 if err != nil { 66 t.Fatalf("http.NewRequest(): got %v, want no error", err) 67 } 68 req.Header.Set(xfp, "https") 69 req.Header.Set(xfh, "preserved.host.com") 70 req.Header.Set(xfu, "https://preserved.host.com/foo?x=y") 71 72 if m.ModifyRequest(req); err != nil { 73 t.Fatalf("ModifyRequest(): got %v, want no error", err) 74 } 75 76 if got, want := req.Header.Get(xfp), "https"; got != want { 77 t.Errorf("req.Header.Get(%q): got %q, want %q", xfh, got, want) 78 } 79 if got, want := req.Header.Get(xfh), "preserved.host.com"; got != want { 80 t.Errorf("req.Header.Get(%q): got %q, want %q", xfh, got, want) 81 } 82 if got, want := req.Header.Get(xfu), "https://preserved.host.com/foo?x=y"; got != want { 83 t.Errorf("req.Header.Get(%q): got %q, want %q", xfh, got, want) 84 } 85 }