github.com/google/martian/v3@v3.3.3/port/port_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 port 16 17 import ( 18 "encoding/json" 19 "fmt" 20 "net" 21 "net/http" 22 "strconv" 23 "strings" 24 25 "github.com/google/martian/v3/parse" 26 ) 27 28 func init() { 29 parse.Register("port.Modifier", modifierFromJSON) 30 } 31 32 // Modifier alters the request URL and Host header to use the provided port. 33 // Only one of port, defaultForScheme, or remove may be specified. Whichever is set last is the one that will take effect. 34 // If remove is true, remove the port from the host string ('example.com'). 35 // If defaultForScheme is true, explicitly specify 80 for HTTP or 443 for HTTPS ('http://example.com:80'). Do nothing for a scheme that is not 'http' or 'https'. 36 // If port is specified, explicitly add it to the host string ('example.com:1234'). 37 // If port is zero and the other fields are false, the request will not be modified. 38 type Modifier struct { 39 port int 40 defaultForScheme bool 41 remove bool 42 } 43 44 type modifierJSON struct { 45 Port int `json:"port"` 46 DefaultForScheme bool `json:"defaultForScheme"` 47 Remove bool `json:"remove"` 48 Scope []parse.ModifierType `json:"scope"` 49 } 50 51 // NewModifier returns a RequestModifier that can be configured to alter the request URL and Host header's port. 52 // One of DefaultPortForScheme, UsePort, or RemovePort should be called to configure this modifier. 53 func NewModifier() *Modifier { 54 return &Modifier{} 55 } 56 57 // DefaultPortForScheme configures the modifier to explicitly specify 80 for HTTP or 443 for HTTPS ('http://example.com:80'). 58 // The modifier will not modify requests with a scheme that is not 'http' or 'https'. 59 // This overrides any previous configuration for this modifier. 60 func (m *Modifier) DefaultPortForScheme() { 61 m.defaultForScheme = true 62 m.remove = false 63 } 64 65 // UsePort configures the modifier to add the specified port to the host string ('example.com:1234'). 66 // This overrides any previous configuration for this modifier. 67 func (m *Modifier) UsePort(port int) { 68 m.port = port 69 m.remove = false 70 m.defaultForScheme = false 71 } 72 73 // RemovePort configures the modifier to remove the port from the host string ('example.com'). 74 // This overrides any previous configuration for this modifier. 75 func (m *Modifier) RemovePort() { 76 m.remove = true 77 m.defaultForScheme = false 78 } 79 80 // ModifyRequest alters the request URL and Host header to modify the port as specified. 81 // See docs for Modifier for details. 82 func (m *Modifier) ModifyRequest(req *http.Request) error { 83 if m.port == 0 && !m.defaultForScheme && !m.remove { 84 return nil 85 } 86 87 host := req.URL.Host 88 if strings.Contains(host, ":") { 89 h, _, err := net.SplitHostPort(host) 90 if err != nil { 91 return err 92 } 93 host = h 94 } 95 96 if m.remove { 97 req.URL.Host = host 98 req.Header.Set("Host", host) 99 return nil 100 } 101 102 if m.defaultForScheme { 103 switch req.URL.Scheme { 104 case "http": 105 hp := net.JoinHostPort(host, "80") 106 req.URL.Host = hp 107 req.Header.Set("Host", hp) 108 return nil 109 case "https": 110 hp := net.JoinHostPort(host, "443") 111 req.URL.Host = hp 112 req.Header.Set("Host", hp) 113 return nil 114 default: 115 // Unknown scheme, do nothing. 116 return nil 117 } 118 } 119 120 // Not removing or using default for the scheme, so use the provided port number. 121 hp := net.JoinHostPort(host, strconv.Itoa(m.port)) 122 req.URL.Host = hp 123 req.Header.Set("Host", hp) 124 125 return nil 126 } 127 128 func modifierFromJSON(b []byte) (*parse.Result, error) { 129 msg := &modifierJSON{} 130 if err := json.Unmarshal(b, msg); err != nil { 131 return nil, err 132 } 133 134 errMsg := fmt.Errorf("Must specify only one of port, defaultForScheme or remove") 135 136 mod := NewModifier() 137 // Check that exactly one field of port, defaultForScheme, and remove is set. 138 switch { 139 case msg.Port != 0: 140 if msg.DefaultForScheme || msg.Remove { 141 return nil, errMsg 142 } 143 mod.UsePort(msg.Port) 144 case msg.DefaultForScheme: 145 if msg.Port != 0 || msg.Remove { 146 return nil, errMsg 147 } 148 mod.DefaultPortForScheme() 149 case msg.Remove: 150 if msg.Port != 0 || msg.DefaultForScheme { 151 return nil, errMsg 152 } 153 mod.RemovePort() 154 default: 155 return nil, errMsg 156 } 157 158 return parse.NewResult(mod, msg.Scope) 159 }