github.com/google/martian/v3@v3.3.3/stash/stash_modifier.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 provides a modifier that stores the request URL in a 16 // specified header. 17 package stash 18 19 import ( 20 "encoding/json" 21 "fmt" 22 "net/http" 23 24 "github.com/google/martian/v3/parse" 25 ) 26 27 func init() { 28 parse.Register("stash.Modifier", modifierFromJSON) 29 } 30 31 // Modifier adds a header to the request containing the current state of the URL. 32 // The header will be named with the value stored in headerName. 33 // There will be no validation done on this header name. 34 type Modifier struct { 35 headerName string 36 } 37 38 type modifierJSON struct { 39 HeaderName string `json:"headerName"` 40 Scope []parse.ModifierType `json:"scope"` 41 } 42 43 // NewModifier returns a RequestModifier that write the current URL into a header. 44 func NewModifier(headerName string) *Modifier { 45 return &Modifier{headerName: headerName} 46 } 47 48 // ModifyRequest writes the current URL into a header. 49 func (m *Modifier) ModifyRequest(req *http.Request) error { 50 req.Header.Set(m.headerName, req.URL.String()) 51 return nil 52 } 53 54 // ModifyResponse writes the same header written in the request into the response. 55 func (m *Modifier) ModifyResponse(res *http.Response) error { 56 res.Header.Set(m.headerName, res.Request.Header.Get(m.headerName)) 57 return nil 58 } 59 60 func modifierFromJSON(b []byte) (*parse.Result, error) { 61 // If you would like the saved state of the URL to be written in the response you must specify 62 // this modifier's scope as both request and response. 63 msg := &modifierJSON{} 64 if err := json.Unmarshal(b, msg); err != nil { 65 return nil, err 66 } 67 68 mod := NewModifier(msg.HeaderName) 69 r, err := parse.NewResult(mod, msg.Scope) 70 if err != nil { 71 return nil, err 72 } 73 74 if r.ResponseModifier() != nil && r.RequestModifier() == nil { 75 return nil, fmt.Errorf("to write header on a response, specify scope as both request and response") 76 } 77 78 return r, nil 79 }