github.com/google/martian/v3@v3.3.3/stash/stash_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 stash 16 17 import ( 18 "net" 19 "net/http" 20 "testing" 21 22 "github.com/google/martian/v3/fifo" 23 "github.com/google/martian/v3/parse" 24 "github.com/google/martian/v3/port" 25 "github.com/google/martian/v3/proxyutil" 26 ) 27 28 func TestStashRequest(t *testing.T) { 29 fg := fifo.NewGroup() 30 fg.AddRequestModifier(NewModifier("stashed-url")) 31 pmod := port.NewModifier() 32 pmod.UsePort(8080) 33 fg.AddRequestModifier(pmod) 34 35 req, err := http.NewRequest("GET", "http://example.com", nil) 36 if err != nil { 37 t.Fatalf("NewRequest(): got %v, want no error", err) 38 } 39 40 if err := fg.ModifyRequest(req); err != nil { 41 t.Fatalf("smod.ModifyRequest(): got %v, want no error", err) 42 } 43 44 _, port, err := net.SplitHostPort(req.URL.Host) 45 if err != nil { 46 t.Fatalf("net.SplitHostPort(%q): got %v, want no error", req.URL.Host, err) 47 } 48 49 if got, want := port, "8080"; got != want { 50 t.Errorf("port: got %v, want %v", got, want) 51 } 52 53 if got, want := req.Header.Get("stashed-url"), "http://example.com"; got != want { 54 t.Errorf("stashed-url header: got %v, want %v", got, want) 55 } 56 57 } 58 59 func TestStashRequestResponse(t *testing.T) { 60 headerName := "stashed-url" 61 originalURL := "http://example.com" 62 fg := fifo.NewGroup() 63 fg.AddRequestModifier(NewModifier(headerName)) 64 fg.AddResponseModifier(NewModifier(headerName)) 65 pmod := port.NewModifier() 66 pmod.UsePort(8080) 67 fg.AddRequestModifier(pmod) 68 69 req, err := http.NewRequest("GET", originalURL, nil) 70 if err != nil { 71 t.Fatalf("NewRequest(): got %v, want no error", err) 72 } 73 74 if err := fg.ModifyRequest(req); err != nil { 75 t.Fatalf("smod.ModifyRequest(): got %v, want no error", err) 76 } 77 78 _, port, err := net.SplitHostPort(req.URL.Host) 79 if err != nil { 80 t.Fatalf("net.SplitHostPort(%q): got %v, want no error", req.URL.Host, err) 81 } 82 83 if got, want := port, "8080"; got != want { 84 t.Errorf("port: got %v, want %v", got, want) 85 } 86 87 if got, want := req.Header.Get(headerName), originalURL; got != want { 88 t.Errorf("res.Header.Get(%q): got %v, want %v", headerName, got, want) 89 } 90 91 res := proxyutil.NewResponse(200, nil, req) 92 if err := fg.ModifyResponse(res); err != nil { 93 t.Fatalf("resmod.ModifyResponse(): got %v, want no error", err) 94 } 95 96 if got, want := res.Header.Get(headerName), originalURL; got != want { 97 t.Errorf("res.Header.Get(%q): got %q, want %q", headerName, got, want) 98 } 99 } 100 101 func TestStashInvalidHeaderName(t *testing.T) { 102 mod := NewModifier("invalid-chars-actually-work-;><@") 103 104 req, err := http.NewRequest("GET", "http://example.com", nil) 105 if err != nil { 106 t.Fatalf("NewRequest(): got %v, want no error", err) 107 } 108 109 if err := mod.ModifyRequest(req); err != nil { 110 t.Fatalf("smod.ModifyRequest(): got %v, want no error", err) 111 } 112 113 if got, want := req.Header.Get("invalid-chars-actually-work-;><@"), "http://example.com"; got != want { 114 t.Errorf("stashed-url header: got %v, want %v", got, want) 115 } 116 } 117 118 func TestModiferFromJSON(t *testing.T) { 119 headerName := "stashed-url" 120 originalURL := "http://example.com" 121 msg := []byte(`{ 122 "fifo.Group": { 123 "scope": ["request", "response"], 124 "modifiers": [ 125 { 126 "stash.Modifier": { 127 "scope": ["request", "response"], 128 "headerName": "stashed-url" 129 } 130 }, 131 { 132 "port.Modifier": { 133 "scope": ["request"], 134 "port": 8080 135 } 136 } 137 ] 138 } 139 }`) 140 141 r, err := parse.FromJSON(msg) 142 if err != nil { 143 t.Fatalf("parse.FromJSON(): got %v, want no error", err) 144 } 145 146 reqmod := r.RequestModifier() 147 if reqmod == nil { 148 t.Fatal("reqmod: got nil, want not nil") 149 } 150 151 req, err := http.NewRequest("GET", originalURL, nil) 152 if err != nil { 153 t.Fatalf("NewRequest(): got %v, want no error", err) 154 } 155 156 if err := reqmod.ModifyRequest(req); err != nil { 157 t.Fatalf("ModifyRequest(): got %v, want no error", err) 158 } 159 160 _, port, err := net.SplitHostPort(req.URL.Host) 161 if err != nil { 162 t.Fatalf("net.SplitHostPort(%q): got %v, want no error", req.URL.Host, err) 163 } 164 165 if got, want := port, "8080"; got != want { 166 t.Errorf("port: got %v, want %v", got, want) 167 } 168 169 if got, want := req.Header.Get(headerName), originalURL; got != want { 170 t.Errorf("req.Header.Get(%q) header: got %v, want %v", headerName, got, want) 171 } 172 173 resmod := r.ResponseModifier() 174 res := proxyutil.NewResponse(200, nil, req) 175 if err := resmod.ModifyResponse(res); err != nil { 176 t.Fatalf("resmod.ModifyResponse(): got %v, want no error", err) 177 } 178 179 if got, want := res.Header.Get(headerName), originalURL; got != want { 180 t.Errorf("res.Header.Get(%q): got %q, want %q", headerName, got, want) 181 } 182 } 183 184 func TestModiferFromJSONInvalidConfigurations(t *testing.T) { 185 msg := []byte(`{ 186 "stash.Modifier": { 187 "scope": ["response"], 188 "headerName": "stash-header" 189 } 190 }`) 191 192 _, err := parse.FromJSON(msg) 193 if err == nil { 194 t.Fatalf("parseFromJSON(msg): Got no error, but should have gotten one.") 195 } 196 }