Caching proxy for the npm registry

main.go

package main

import (
	"github.com/golang/groupcache"
	"github.com/gorilla/mux"
	"github.com/pranavraja/front/cache"

	"bytes"
	"flag"
	"io/ioutil"
	"net/http"
	"os"
	"time"
)

// The HTTP client used by all upstream requests
var client *http.Client

const upstream = "http://registry.npmjs.org"

var host string

func getBlob(w http.ResponseWriter, r *http.Request) {
	client.Get(upstream + r.URL.Path)
}

func init() {
	timeoutTransport := &http.Transport{
		Proxy: http.ProxyFromEnvironment,
		ResponseHeaderTimeout: 5 * time.Second,
	}
	client = &http.Client{Transport: timeoutTransport}
	flag.StringVar(&host, "service", "", "The address of this service externally")
}

type cacheGetter interface {
	Get(groupcache.Context, string, groupcache.Sink) error
}

func getAndWriteResponse(w http.ResponseWriter, g cacheGetter, key string) {
	var buf []byte
	err := g.Get(nil, key, groupcache.AllocatingByteSliceSink(&buf))
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	_, err = w.Write(buf)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
}

const TTLSeconds = 3600

func getUpstream(path string) (data []byte, ttl time.Duration) {
	resp, err := client.Get(upstream + path)
	if err != nil {
		return
	}
	defer resp.Body.Close()
	buf, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return
	}
	buf = bytes.Replace(buf, []byte(upstream), []byte(host), -1)
	return buf, time.Duration(TTLSeconds) * time.Second
}

func getUpstreamIntoSink(ctx groupcache.Context, path string, dest groupcache.Sink) error {
	data, _ := getUpstream(path)
	err := dest.SetBytes(data)
	return err
}

func main() {
	flag.Parse()
	if host == "" {
		println("No service address provided. I DON'T KNOW WHO I AM, DUDE")
		os.Exit(1)
	}
	r := mux.NewRouter()
	// Use a TTL cache to get the list of versions, as it may change
	ttlCache := cache.New(getUpstream)
	r.HandleFunc("/{name}", func(w http.ResponseWriter, r *http.Request) {
		buf, _ := ttlCache.Get(r.URL.Path)
		if buf == nil {
			http.Error(w, "Couldn't fetch "+upstream+r.URL.Path, http.StatusBadGateway)
		}
		_, err := w.Write(buf)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
		}
	})
	// Use an immutable cache for the rest, with version numbers in the URL
	upstreamGetterFunc := groupcache.GetterFunc(getUpstreamIntoSink)
	// Allow 1GB for blob storage
	blobGroup := groupcache.NewGroup("blob", 1<<9, upstreamGetterFunc)
	r.HandleFunc("/{name}/-/{filename}", func(w http.ResponseWriter, r *http.Request) {
		getAndWriteResponse(w, blobGroup, r.URL.Path)
	})
	// Allow 10MB for version metadata storage
	metadataGroup := groupcache.NewGroup("versionMetadata", 1<<7, upstreamGetterFunc)
	r.HandleFunc("/{name}/{version}", func(w http.ResponseWriter, r *http.Request) {
		getAndWriteResponse(w, metadataGroup, r.URL.Path)
	})
	http.Handle("/", r)
	port := os.Getenv("PORT")
	if port == "" {
		port = "5000"
	}
	http.ListenAndServe(":"+port, nil)
}