1 // Copyright 2020 Palantir Technologies, Inc. 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 // Code sourced from the following and modified in a few ways: 16 // https://github.com/zenazn/goji/blob/a16712d37ba72246f71f9c8012974d46f8e61d16/web/mutil/writer_proxy.go 17 18 // Copyright (c) 2014, 2015, 2016 Carl Jackson (carl@avtok.com) 19 // 20 // MIT License 21 // 22 // Permission is hereby granted, free of charge, to any person obtaining a copy of 23 // this software and associated documentation files (the "Software"), to deal in 24 // the Software without restriction, including without limitation the rights to 25 // use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of 26 // the Software, and to permit persons to whom the Software is furnished to do so, 27 // subject to the following conditions: 28 // 29 // The above copyright notice and this permission notice shall be included in all 30 // copies or substantial portions of the Software. 31 // 32 // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 33 // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS 34 // FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR 35 // COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER 36 // IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN 37 // CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. 38 39 package baseapp 40 41 import ( 42 "bufio" 43 "io" 44 "net" 45 "net/http" 46 ) 47 48 // RecordingResponseWriter is a proxy for an http.ResponseWriter that 49 // counts bytes written and http status send to the underlying ResponseWriter. 50 type RecordingResponseWriter interface { 51 http.ResponseWriter 52 53 // Status returns the HTTP status of the request, or 0 if one has not 54 // yet been sent. 55 Status() int 56 57 // BytesWritten returns the total number of bytes sent to the client. 58 BytesWritten() int64 59 } 60 61 func WrapWriter(w http.ResponseWriter) RecordingResponseWriter { 62 _, cn := w.(http.CloseNotifier) 63 _, fl := w.(http.Flusher) 64 _, hj := w.(http.Hijacker) 65 _, rf := w.(io.ReaderFrom) 66 67 bp := basicRecorder{ResponseWriter: w} 68 if cn && fl && hj && rf { 69 return &fancyRecorder{bp} 70 } 71 if fl { 72 return &flushRecorder{bp} 73 } 74 return &bp 75 } 76 77 type basicRecorder struct { 78 http.ResponseWriter 79 code int 80 bytesWritten int64 81 } 82 83 func (b *basicRecorder) WriteHeader(code int) { 84 if b.code == 0 { 85 b.code = code 86 } 87 b.ResponseWriter.WriteHeader(code) 88 } 89 90 func (b *basicRecorder) Write(buf []byte) (int, error) { 91 if b.code == 0 { 92 b.code = http.StatusOK 93 } 94 n, err := b.ResponseWriter.Write(buf) 95 b.bytesWritten += int64(n) 96 return n, err 97 } 98 99 func (b *basicRecorder) Status() int { 100 return b.code 101 } 102 103 func (b *basicRecorder) BytesWritten() int64 { 104 return b.bytesWritten 105 } 106 107 // fancyRecorder is a writer that additionally satisfies http.CloseNotifier, 108 // http.Flusher, http.Hijacker, and io.ReaderFrom. It exists for the common case 109 // of wrapping the http.ResponseWriter that package http gives you, in order to 110 // make the proxied object support the full method set of the proxied object. 111 type fancyRecorder struct { 112 basicRecorder 113 } 114 115 func (f *fancyRecorder) CloseNotify() <-chan bool { 116 cn := f.basicRecorder.ResponseWriter.(http.CloseNotifier) 117 return cn.CloseNotify() 118 } 119 func (f *fancyRecorder) Flush() { 120 fl := f.basicRecorder.ResponseWriter.(http.Flusher) 121 fl.Flush() 122 } 123 func (f *fancyRecorder) Hijack() (net.Conn, *bufio.ReadWriter, error) { 124 hj := f.basicRecorder.ResponseWriter.(http.Hijacker) 125 return hj.Hijack() 126 } 127 func (f *fancyRecorder) ReadFrom(r io.Reader) (int64, error) { 128 if f.code == 0 { 129 f.code = http.StatusOK 130 } 131 rf := f.basicRecorder.ResponseWriter.(io.ReaderFrom) 132 n, err := rf.ReadFrom(r) 133 f.bytesWritten += n 134 return n, err 135 } 136 137 var _ http.CloseNotifier = &fancyRecorder{} 138 var _ http.Flusher = &fancyRecorder{} 139 var _ http.Hijacker = &fancyRecorder{} 140 var _ io.ReaderFrom = &fancyRecorder{} 141 142 type flushRecorder struct { 143 basicRecorder 144 } 145 146 func (f *flushRecorder) Flush() { 147 fl := f.basicRecorder.ResponseWriter.(http.Flusher) 148 fl.Flush() 149 } 150 151 var _ http.Flusher = &flushRecorder{} 152