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  }