github.com/google/martian/v3@v3.3.3/body/body_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 body allows for the replacement of message body on responses. 16 package body 17 18 import ( 19 "bytes" 20 "crypto/rand" 21 "encoding/json" 22 "fmt" 23 "io" 24 "io/ioutil" 25 "mime/multipart" 26 "net/http" 27 "net/textproto" 28 "strconv" 29 "strings" 30 31 "github.com/google/martian/v3/log" 32 "github.com/google/martian/v3/parse" 33 ) 34 35 func init() { 36 parse.Register("body.Modifier", modifierFromJSON) 37 } 38 39 // Modifier substitutes the body on an HTTP response. 40 type Modifier struct { 41 contentType string 42 body []byte 43 boundary string 44 } 45 46 type modifierJSON struct { 47 ContentType string `json:"contentType"` 48 Body []byte `json:"body"` // Body is expected to be a Base64 encoded string. 49 Scope []parse.ModifierType `json:"scope"` 50 } 51 52 // NewModifier constructs and returns a body.Modifier. 53 func NewModifier(b []byte, contentType string) *Modifier { 54 log.Debugf("body.NewModifier: len(b): %d, contentType %s", len(b), contentType) 55 return &Modifier{ 56 contentType: contentType, 57 body: b, 58 boundary: randomBoundary(), 59 } 60 } 61 62 // modifierFromJSON takes a JSON message as a byte slice and returns a 63 // body.Modifier and an error. 64 // 65 // Example JSON Configuration message: 66 // { 67 // "scope": ["request", "response"], 68 // "contentType": "text/plain", 69 // "body": "c29tZSBkYXRhIHdpdGggACBhbmQg77u/" // Base64 encoded body 70 // } 71 func modifierFromJSON(b []byte) (*parse.Result, error) { 72 msg := &modifierJSON{} 73 if err := json.Unmarshal(b, msg); err != nil { 74 return nil, err 75 } 76 77 mod := NewModifier(msg.Body, msg.ContentType) 78 return parse.NewResult(mod, msg.Scope) 79 } 80 81 // ModifyRequest sets the Content-Type header and overrides the request body. 82 func (m *Modifier) ModifyRequest(req *http.Request) error { 83 log.Debugf("body.ModifyRequest: request: %s", req.URL) 84 req.Body.Close() 85 86 req.Header.Set("Content-Type", m.contentType) 87 88 // Reset the Content-Encoding since we know that the new body isn't encoded. 89 req.Header.Del("Content-Encoding") 90 91 req.ContentLength = int64(len(m.body)) 92 req.Body = ioutil.NopCloser(bytes.NewReader(m.body)) 93 94 return nil 95 } 96 97 // SetBoundary set the boundary string used for multipart range responses. 98 func (m *Modifier) SetBoundary(boundary string) { 99 m.boundary = boundary 100 } 101 102 // ModifyResponse sets the Content-Type header and overrides the response body. 103 func (m *Modifier) ModifyResponse(res *http.Response) error { 104 log.Debugf("body.ModifyResponse: request: %s", res.Request.URL) 105 // Replace the existing body, close it first. 106 res.Body.Close() 107 108 res.Header.Set("Content-Type", m.contentType) 109 110 // Reset the Content-Encoding since we know that the new body isn't encoded. 111 res.Header.Del("Content-Encoding") 112 113 // If no range request header is present, return the body as the response body. 114 if res.Request.Header.Get("Range") == "" { 115 res.ContentLength = int64(len(m.body)) 116 res.Body = ioutil.NopCloser(bytes.NewReader(m.body)) 117 118 return nil 119 } 120 121 rh := res.Request.Header.Get("Range") 122 rh = strings.ToLower(rh) 123 sranges := strings.Split(strings.TrimLeft(rh, "bytes="), ",") 124 var ranges [][]int 125 for _, rng := range sranges { 126 if strings.HasSuffix(rng, "-") { 127 rng = fmt.Sprintf("%s%d", rng, len(m.body)-1) 128 } 129 130 rs := strings.Split(rng, "-") 131 if len(rs) != 2 { 132 res.StatusCode = http.StatusRequestedRangeNotSatisfiable 133 return nil 134 } 135 start, err := strconv.Atoi(strings.TrimSpace(rs[0])) 136 if err != nil { 137 return err 138 } 139 140 end, err := strconv.Atoi(strings.TrimSpace(rs[1])) 141 if err != nil { 142 return err 143 } 144 145 if start > end { 146 res.StatusCode = http.StatusRequestedRangeNotSatisfiable 147 return nil 148 } 149 150 ranges = append(ranges, []int{start, end}) 151 } 152 153 // Range request. 154 res.StatusCode = http.StatusPartialContent 155 156 // Single range request. 157 if len(ranges) == 1 { 158 start := ranges[0][0] 159 end := ranges[0][1] 160 seg := m.body[start : end+1] 161 res.ContentLength = int64(len(seg)) 162 res.Body = ioutil.NopCloser(bytes.NewReader(seg)) 163 res.Header.Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end, len(m.body))) 164 165 return nil 166 } 167 168 // Multipart range request. 169 var mpbody bytes.Buffer 170 mpw := multipart.NewWriter(&mpbody) 171 mpw.SetBoundary(m.boundary) 172 173 for _, rng := range ranges { 174 start, end := rng[0], rng[1] 175 mimeh := make(textproto.MIMEHeader) 176 mimeh.Set("Content-Type", m.contentType) 177 mimeh.Set("Content-Range", fmt.Sprintf("bytes %d-%d/%d", start, end, len(m.body))) 178 179 seg := m.body[start : end+1] 180 181 pw, err := mpw.CreatePart(mimeh) 182 if err != nil { 183 return err 184 } 185 186 if _, err := pw.Write(seg); err != nil { 187 return err 188 } 189 } 190 mpw.Close() 191 192 res.ContentLength = int64(len(mpbody.Bytes())) 193 res.Body = ioutil.NopCloser(bytes.NewReader(mpbody.Bytes())) 194 res.Header.Set("Content-Type", fmt.Sprintf("multipart/byteranges; boundary=%s", m.boundary)) 195 196 return nil 197 } 198 199 // randomBoundary generates a 30 character string for boundaries for mulipart range 200 // requests. This func panics if io.Readfull fails. 201 // Borrowed from: https://golang.org/src/mime/multipart/writer.go?#L73 202 func randomBoundary() string { 203 var buf [30]byte 204 _, err := io.ReadFull(rand.Reader, buf[:]) 205 if err != nil { 206 panic(err) 207 } 208 return fmt.Sprintf("%x", buf[:]) 209 }