go-yeah.go

package main

import (
	"flag"
	"io"
	"net/http"
	"os"
	"strings"

	"github.com/gorilla/handlers"
	"golang.org/x/net/html"
	"golang.org/x/net/html/atom"
)

func main() {
	addr := flag.String("addr", ":8080", "ip:port to listen on")
	origin := flag.String("origin", "http://localhost:8080", "this server's url")
	audiof := flag.String("audio_file", "oh_yeah.mp3", "audio file to play on proxied pages")
	flag.Parse()
	p := &Proxy{
		Origin:    *origin,
		AudioFile: *audiof,
	}
	logger := func(h http.Handler) http.Handler {
		return handlers.CombinedLoggingHandler(os.Stderr, h)
	}
	http.Handle("/proxy", logger(p.Handler()))
	http.Handle("/proxy/audio.mp3", logger(p.AudioHandler()))
	http.ListenAndServe(*addr, nil)
}

// Proxy provides an http handler for proxying pages that are modified to
// auto play an audio.
type Proxy struct {
	Origin    string
	AudioFile string
}

// Handler takes a 'url' parameter and proxy it.
//
// e.g. curl -i $server/proxy?url=http://google.com
func (p *Proxy) Handler() http.Handler {
	f := func(w http.ResponseWriter, r *http.Request) {
		if r.Method != "GET" {
			http.Error(w, http.StatusText(http.StatusMethodNotAllowed), http.StatusMethodNotAllowed)
			return
		}
		url := r.FormValue("url")
		if url == "" {
			http.Error(w, "Missing 'url' param", http.StatusBadRequest)
			return
		}
		p.do(w, url)
	}
	return http.HandlerFunc(f)
}

// AudioHandler returns the audio file configured in the Proxy.
func (p *Proxy) AudioHandler() http.Handler {
	f := func(w http.ResponseWriter, r *http.Request) {
		http.ServeFile(w, r, p.AudioFile)
	}
	return http.HandlerFunc(f)
}

// Hop-by-hop headers. These are removed when sent to the backend.
// http://www.w3.org/Protocols/rfc2616/rfc2616-sec13.html
var hopHeaders = []string{
	"Connection",
	"Keep-Alive",
	"Proxy-Authenticate",
	"Proxy-Authorization",
	"Te", // canonicalized version of "TE"
	"Trailers",
	"Transfer-Encoding",
	"Upgrade",
}

// do makes the upstream request and proxy the response back to client w.
func (p *Proxy) do(w http.ResponseWriter, url string) {
	resp, err := http.Get(url)
	if err != nil {
		http.Error(w, err.Error(), http.StatusServiceUnavailable)
		return
	}
	defer resp.Body.Close()
	// remove certain headers from the response
	for _, h := range hopHeaders {
		resp.Header.Del(h)
	}
	// copy headers
	for k, vv := range resp.Header {
		for _, v := range vv {
			w.Header().Set(k, v)
		}
	}
	// direct proxy non http 200 text/html
	ok := strings.HasPrefix(resp.Header.Get("Content-Type"), "text/html")
	if !ok || resp.StatusCode != http.StatusOK {
		w.WriteHeader(resp.StatusCode)
		io.Copy(w, resp.Body)
		return
	}
	p.handleBody(w, resp.Body, url)
}

// handleBody handles the response body and writes to client w.
func (p *Proxy) handleBody(w http.ResponseWriter, body io.Reader, url string) {
	//io.Copy(w, body)
	doc, err := html.Parse(body)
	if err != nil {
		http.Error(w, err.Error(), http.StatusServiceUnavailable)
		return
	}
	var f func(n *html.Node)
	stop := false
	f = func(n *html.Node) {
		if n.Type == html.ElementNode {
			switch n.Data {
			case "head":
				n.InsertBefore(&html.Node{
					Type:     html.ElementNode,
					Data:     "base",
					DataAtom: atom.Base,
					Attr: []html.Attribute{
						{Key: "href", Val: url},
						{Key: "target", Val: "_blank"},
					},
				}, n.FirstChild)
			case "body":
				n.AppendChild(&html.Node{
					Type:     html.ElementNode,
					Data:     "audio",
					DataAtom: atom.Audio,
					Attr: []html.Attribute{
						{Key: "preload", Val: "auto"},
						{Key: "autoplay", Val: "autoplay"},
						{Key: "loop", Val: "loop"},
					},
					FirstChild: &html.Node{
						Type:     html.ElementNode,
						Data:     "source",
						DataAtom: atom.Source,
						Attr: []html.Attribute{
							{Key: "type", Val: "audio/mpeg"},
							{Key: "src", Val: p.Origin + "/proxy/audio.mp3"},
						},
					},
				})
				stop = true
				return
			}
		}
		for c := n.FirstChild; c != nil && !stop; c = c.NextSibling {
			f(c)
		}
	}
	f(doc)
	html.Render(w, doc)
}