github.com/franciscocpg/up@v0.1.10/http/inject/inject.go (about) 1 // Package inject provides script and style injection. 2 package inject 3 4 import ( 5 "bytes" 6 "io" 7 "net/http" 8 "strconv" 9 "strings" 10 11 "github.com/apex/up" 12 "github.com/apex/up/internal/inject" 13 ) 14 15 // response wrapper. 16 type response struct { 17 http.ResponseWriter 18 rules inject.Rules 19 body bytes.Buffer 20 header bool 21 ignore bool 22 code int 23 } 24 25 // Write implementation. 26 func (r *response) Write(b []byte) (int, error) { 27 if !r.header { 28 r.WriteHeader(200) 29 return r.Write(b) 30 } 31 32 return r.body.Write(b) 33 } 34 35 // WriteHeader implementation. 36 func (r *response) WriteHeader(code int) { 37 r.header = true 38 w := r.ResponseWriter 39 kind := w.Header().Get("Content-Type") 40 r.ignore = !strings.HasPrefix(kind, "text/html") || code >= 300 41 r.code = code 42 } 43 44 // end injects if necessary. 45 func (r *response) end() { 46 w := r.ResponseWriter 47 48 if r.ignore { 49 w.WriteHeader(r.code) 50 r.body.WriteTo(w) 51 return 52 } 53 54 body := r.rules.Apply(r.body.String()) 55 w.Header().Set("Content-Length", strconv.Itoa(len(body))) 56 io.WriteString(w, body) 57 } 58 59 // New inject handler. 60 func New(c *up.Config, next http.Handler) (http.Handler, error) { 61 if len(c.Inject) == 0 { 62 return next, nil 63 } 64 65 h := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { 66 res := &response{ResponseWriter: w, rules: c.Inject} 67 next.ServeHTTP(res, r) 68 res.end() 69 }) 70 71 return h, nil 72 }