Basic cookie injecting proxy in Go

proxy.go

package proxy

import (
	"log"
	"net/http"
	"net/http/httputil"
	"net/url"
	"time"
)

func Start(listen string, origin *url.URL, stop chan int) {
	director := basicDirector(origin)
	transport := &cookieInjector{Transport: http.DefaultTransport}

	proxy := &httputil.ReverseProxy{Director: director, Transport: transport}

	go http.ListenAndServe(listen, proxy)

	log.Printf("Listening on %s, forwarding to %s\n", listen, origin)

	<-stop
}

// basic single origin proxy director

func basicDirector(origin *url.URL) func(*http.Request) {
	return func(req *http.Request) {
		req.URL.Scheme = origin.Scheme
		req.URL.Host = origin.Host
		req.Host = origin.Host
	}
}

// Cookie injecting transport

type cookieInjector struct {
	Transport http.RoundTripper
}

func (injector *cookieInjector) RoundTrip(req *http.Request) (*http.Response, error) {
	testCookie := findCookie(req.Cookies(), "Test")

	var cookie *http.Cookie
	if testCookie == nil {
		cookie = injectCookie(req)
	}

	t := time.Now()

	resp, error := injector.Transport.RoundTrip(req)

	log.Printf("%s %s -> %s: %s(took: %v)", req.Method, req.URL.Path, req.Host, resp.Status, time.Now().Sub(t))

	if testCookie == nil {
		setCookie := append(resp.Header["Set-Cookie"], cookie.String())
		resp.Header["Set-Cookie"] = setCookie
	}

	return resp, error
}

// Helper functions

func findCookie(cookies []*http.Cookie, name string) *http.Cookie {
	for _, c := range cookies {
		if c.Name == "Test" {
			return c
		}
	}

	return nil
}

func injectCookie(req *http.Request) *http.Cookie {
	time := time.Now().Add(24 * time.Hour)
	cookie := &http.Cookie{Name: "Test", Value: "TestValue", Path: "/", Expires: time}

	req.AddCookie(cookie)

	return cookie
}