Merge pull request #91 from wader/inject-only-html

Only inject if content type is text/html

Author: cortesi

fileserver/fileserver.go

@@ -126,7 +126,7 @@ func serveContent(ci inject.CopyInject, w http.ResponseWriter, r *http.Request,
 		ctype = ctypes[0]
 	}
 
-	injector, err := ci.Sniff(content)
+	injector, err := ci.Sniff(content, ctype)
 	if err != nil {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return err
@@ -138,7 +138,7 @@ func serveContent(ci inject.CopyInject, w http.ResponseWriter, r *http.Request,
 		return err
 	}
 
-	if injector.Found {
+	if injector.Found() {
 		size = size + int64(injector.Extra())
 	}
 

inject/inject.go

@@ -15,37 +15,65 @@ import (
 	"io"
 	"net/http"
 	"regexp"
+	"strings"
 )
 
 // CopyInject copies data, and injects a payload before a specified marker
 type CopyInject struct {
 	// Number of initial bytes within which to search for marker
 	Within int
+	// Only inject in responses with this content type
+	ContentType string
 	// A marker, BEFORE which the payload is inserted
 	Marker *regexp.Regexp
 	// The payload to be inserted
 	Payload []byte
 }
 
-// Injector keeps injection state
-type Injector struct {
-	// Has the marker been found?
-	Found bool
+type Injector interface {
+	Copy(dst io.Writer) (int64, error)
+	Extra() int
+	Found() bool
+}
 
+// realInjector keeps injection state
+type realInjector struct {
+	// Has the marker been found?
+	found       bool
 	conf        *CopyInject
 	src         io.Reader
 	offset      int
 	sniffedData []byte
 }
 
+type nopInjector struct {
+	src io.Reader
+}
+
+func (injector *nopInjector) Copy(dst io.Writer) (int64, error) {
+	return io.Copy(dst, injector.src)
+}
+
+func (injector *nopInjector) Extra() int {
+	return 0
+}
+
+func (injector *nopInjector) Found() bool {
+	return false
+}
+
 // Extra reports the number of extra bytes that will be injected
-func (injector *Injector) Extra() int {
-	if injector.Found {
+func (injector *realInjector) Extra() int {
+	if injector.found {
 		return len(injector.conf.Payload)
 	}
 	return 0
 }
 
+func (injector *realInjector) Found() bool {
+	return injector.found
+}
+
 func min(a int, b int) int {
 	if a > b {
 		return b
@@ -55,8 +83,12 @@ func min(a int, b int) int {
 
 // Sniff reads the first SniffLen bytes of the source, and checks for the
 // marker. Returns an Injector instance.
-func (ci *CopyInject) Sniff(src io.Reader) (*Injector, error) {
-	injector := &Injector{
+func (ci *CopyInject) Sniff(src io.Reader, contentType string) (Injector, error) {
+	if !strings.Contains(contentType, ci.ContentType) {
+		return &nopInjector{src: src}, nil
+	}
+
+	injector := &realInjector{
 		conf: ci,
 		src:  src,
 	}
@@ -74,7 +106,7 @@ func (ci *CopyInject) Sniff(src io.Reader) (*Injector, error) {
 	}
 	loc := ci.Marker.FindIndex(injector.sniffedData[:min(n, ci.Within)])
 	if loc != nil {
-		injector.Found = true
+		injector.found = true
 		injector.offset = loc[0]
 	}
 	return injector, nil
@@ -89,7 +121,7 @@ func (ci *CopyInject) ServeTemplate(statuscode int, w http.ResponseWriter, t *te
 	}
 
 	length := buff.Len()
-	inj, err := ci.Sniff(buff)
+	inj, err := ci.Sniff(buff, "text/html")
 	if err != nil {
 		return err
 	}
@@ -106,9 +138,9 @@ func (ci *CopyInject) ServeTemplate(statuscode int, w http.ResponseWriter, t *te
 
 // Copy copies the data from src to dst, injecting the Payload if Sniff found
 // the marker.
-func (injector *Injector) Copy(dst io.Writer) (int64, error) {
+func (injector *realInjector) Copy(dst io.Writer) (int64, error) {
 	var preludeLen int64
-	if injector.Found {
+	if injector.found {
 		startn, err := io.Copy(
 			dst,
 			bytes.NewBuffer(

inject/inject_test.go

@@ -7,18 +7,31 @@ import (
 	"testing"
 )
 
-func inject(ci CopyInject, data string) (found bool, dstdata string, err error) {
+func inject(ci CopyInject, data string, contentType string) (found bool, dstdata string, err error) {
 	src := bytes.NewBuffer([]byte(data))
 	dst := bytes.NewBuffer(make([]byte, 0))
-	injector, err := ci.Sniff(src)
+	injector, err := ci.Sniff(src, contentType)
 	if err != nil {
 		return false, "", err
 	}
 	_, err = injector.Copy(dst)
 	if err != nil {
 		return false, "", err
 	}
-	return injector.Found, string(dst.Bytes()), nil
+	return injector.Found(), string(dst.Bytes()), nil
+}
+
+func TestReverseProxyNoInject(t *testing.T) {
+	ci := CopyInject{
+		Within:      100,
+		ContentType: "text/html",
+		Marker:      regexp.MustCompile("mark"),
+		Payload:     []byte("inject"),
+	}
+	found, dst, err := inject(ci, "imark", "text/plain")
+	if err != nil || found || dst != "imark" {
+		t.Errorf("Unexpected, found:%v dst:%v error:%v", dst, found, err)
+	}
 }
 
 func TestReverseProxy(t *testing.T) {
@@ -43,11 +56,12 @@ func TestReverseProxy(t *testing.T) {
 	}
 	for i, tt := range sniffTests {
 		ci := CopyInject{
-			Within:  tt.snifflen,
-			Marker:  regexp.MustCompile(tt.marker),
-			Payload: []byte(tt.payload),
+			Within:      tt.snifflen,
+			ContentType: "text/html",
+			Marker:      regexp.MustCompile(tt.marker),
+			Payload:     []byte(tt.payload),
 		}
-		found, dst, err := inject(ci, tt.src)
+		found, dst, err := inject(ci, tt.src, "text/html")
 
 		// Sanity checkss
 		if err != nil {
@@ -74,7 +88,7 @@ func TestReverseProxy(t *testing.T) {
 		}
 
 		// Idempotence
-		found, dst2, err := inject(ci, dst)
+		found, dst2, err := inject(ci, dst, "text/html")
 		if err != nil {
 			t.Errorf("Test %d, unexpected error:\n%s\n", i, err)
 		}
@@ -91,8 +105,8 @@ func TestNil(t *testing.T) {
 	ci := CopyInject{}
 	val := "onetwothree"
 	src := bytes.NewBuffer([]byte(val))
-	injector, err := ci.Sniff(src)
-	if injector.Found || err != nil {
+	injector, err := ci.Sniff(src, "")
+	if injector.Found() || err != nil {
 		t.Error("Unexpected")
 	}
 	dst := bytes.NewBuffer(make([]byte, 0))

livereload/livereload.go

@@ -31,9 +31,10 @@ const (
 
 // Injector for the livereload script
 var Injector = inject.CopyInject{
-	Within:  1024 * 30,
-	Marker:  regexp.MustCompile(`<\/head>`),
-	Payload: []byte(`<script src="/.devd.livereload.js"></script>`),
+	Within:      1024 * 30,
+	ContentType: "text/html",
+	Marker:      regexp.MustCompile(`<\/head>`),
+	Payload:     []byte(`<script src="/.devd.livereload.js"></script>`),
 }
 
 // Server implements a Livereload server

responselogger.go

@@ -15,6 +15,7 @@ import (
 type ResponseLogWriter struct {
 	Log         termlog.Logger
 	Resp        http.ResponseWriter
+	Flusher     http.Flusher
 	Timer       *timer.Timer
 	wroteHeader bool
 }
@@ -80,3 +81,9 @@ func (rl *ResponseLogWriter) WriteHeader(code int) {
 	rl.Resp.WriteHeader(code)
 	rl.Timer.ResponseDone()
 }
+
+func (rl *ResponseLogWriter) Flush() {
+	if rl.Flusher != nil {
+		rl.Flusher.Flush()
+	}
+}

reverseproxy/reverseproxy.go

@@ -173,14 +173,14 @@ func (p *ReverseProxy) ServeHTTPContext(
 		log.Say(fmt.Sprintf("%s uploaded", humanize.Bytes(uint64(req.ContentLength))))
 	}
 
-	inject, err := p.Inject.Sniff(res.Body)
+	inject, err := p.Inject.Sniff(res.Body, res.Header.Get("Content-Type"))
 	if err != nil {
 		log.Shout("reverse proxy error: %v", err)
 		rw.WriteHeader(http.StatusInternalServerError)
 		return
 	}
 
-	if inject.Found {
+	if inject.Found() {
 		cl, err := strconv.ParseInt(res.Header.Get("Content-Length"), 10, 32)
 		if err == nil {
 			cl = cl + int64(inject.Extra())
@@ -196,7 +196,7 @@ func (p *ReverseProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	p.ServeHTTPContext(context.Background(), w, r)
 }
 
-func (p *ReverseProxy) copyResponse(ctx context.Context, dst io.Writer, inject *inject.Injector) {
+func (p *ReverseProxy) copyResponse(ctx context.Context, dst io.Writer, inject inject.Injector) {
 	log := termlog.FromContext(ctx)
 	if p.FlushInterval != 0 {
 		if wf, ok := dst.(writeFlusher); ok {

route.go

@@ -7,6 +7,7 @@ import (
 	"html/template"
 	"net/http"
 	"net/url"
+	"time"
 
 	"github.com/cortesi/devd/fileserver"
 	"github.com/cortesi/devd/httpctx"
@@ -31,6 +32,7 @@ func (ep forwardEndpoint) Handler(prefix string, templates *template.Template, c
 	rp.Transport = &http.Transport{
 		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
 	}
+	rp.FlushInterval = 200 * time.Millisecond
 	return httpctx.StripPrefix(prefix, rp)
 }
 

server.go

@@ -192,9 +192,10 @@ func (dd *Devd) WrapHandler(log termlog.TermLog, next httpctx.Handler) http.Hand
 				}
 			}
 		}
+		flusher, _ := w.(http.Flusher)
 		next.ServeHTTPContext(
 			ctx,
-			&ResponseLogWriter{Log: sublog, Resp: w, Timer: &timr},
+			&ResponseLogWriter{Log: sublog, Resp: w, Flusher: flusher, Timer: &timr},
 			r,
 		)
 	})