csrf injecting proxy

csrfproxy.go

import (
	"bytes"
	"errors"
	"fmt"
	"github.com/gorilla/csrf"
	"golang.org/x/net/html"
	"io"
	"io/ioutil"
	"net/http"
	"net/http/httputil"
	"strconv"
)

func injectHead(proxy *httputil.ReverseProxy) *httputil.ReverseProxy {
	proxy2 := new(httputil.ReverseProxy)
	proxy2.Director = func(r *http.Request) {
		r.Header.Set("Accept-Encoding", "identity")
		proxy.Director(r)
	}
	proxy2.ModifyResponse = func(r *http.Response) error {
		buffer := new(bytes.Buffer)
		buffer.Grow(int(r.ContentLength) + 117)

		tokenizer := html.NewTokenizer(r.Body)

		if csrf.Token(r.Request) == "" {
			return errors.New("no csrf middleware configured")
		}

		for {
			tt := tokenizer.Next()

			if err := tokenizer.Err(); err != nil {
				if err != io.EOF {
					return err
				}
				break
			}

			if _, err := buffer.Write(tokenizer.Raw()); err != nil {
				return err
			}

			if tt == html.StartTagToken {
				if tag, _ := tokenizer.TagName(); bytes.EqualFold(tag, []byte("head")) {
					if _, err := fmt.Fprintf(buffer, `<meta name="csrf" content="%s" />`, csrf.Token(r.Request)); err != nil {
						return err
					}

					if _, err := buffer.Write(tokenizer.Buffered()); err != nil {
						return err
					}

					if _, err := io.Copy(buffer, r.Body); err != nil {
						return err
					}
					break
				}
			}
		}

		if err := r.Body.Close(); err != nil {
			return err
		}

		r.ContentLength = int64(buffer.Len())
		r.Header.Set("Content-Length", strconv.Itoa(buffer.Len()))

		r.Body = ioutil.NopCloser(buffer)

		return nil
	}

	return proxy2
}