Whitelist Handler

main.go

package main

import (
	_ "expvar"
	"io"
	"log"
	"net/http"
	"strings"
)

func helloServer(w http.ResponseWriter, req *http.Request) {
	io.WriteString(w, "Hello")
}

type WhitelistHandler struct {
	AllowedIPsForURI map[string]map[string]bool
}

func (h *WhitelistHandler) WhitelistIP(requestURI, remoteIP string) {
	if h.AllowedIPsForURI == nil {
		h.AllowedIPsForURI = make(map[string]map[string]bool)
	}

	ips := h.AllowedIPsForURI[requestURI]
	if ips == nil {
		h.AllowedIPsForURI[requestURI] = map[string]bool{remoteIP: true}
	} else {
		ips[remoteIP] = true
	}
}

func parseIP(s string) string {
	if !strings.HasPrefix(s, "[") {
		return s
	}

	index := strings.Index(s, "]")
	if index > 0 {
		return s[1:index]
	}

	return s
}

func (h *WhitelistHandler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
	// log.Println(req.RequestURI, "from", req.RemoteAddr)

	allowedIPs := h.AllowedIPsForURI[req.RequestURI]
	if allowedIPs == nil {
		http.DefaultServeMux.ServeHTTP(w, req)
		return
	}

	if allowedIPs[parseIP(req.RemoteAddr)] {
		http.DefaultServeMux.ServeHTTP(w, req)
		return
	}

	w.WriteHeader(http.StatusForbidden)
	return
}

func main() {
	log.Println("Starting helloServer")
	http.HandleFunc("/", helloServer)

	h := &WhitelistHandler{}
	h.WhitelistIP("/debug/vars", "127.0.0.1")
	h.WhitelistIP("/debug/vars", "127.0.0.2")
	h.WhitelistIP("/debug/vars", "::1")

	http.ListenAndServe(":8080", h)
}