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  }