Proxy GCE metadata server

ratelimitproxy.go

// go build ratelimitproxy.go
// ./ratelimitproxy -n 2
// export http_proxy=http://localhost:8080
package main

import (
        "sync/atomic"
        "flag"
        "log"
        "net/http"
        "net/http/httputil"
        "net/url"
        "strings"
)

var (
        listen = flag.String("l", "localhost:8080", "address:port to listen on")
        upstream = flag.String("u", "http://169.254.169.254:80/", "Upstream URL to route requests to")
        modulo = flag.Int("n", 2, "fail every nth request, starting with the first")
)

type server struct {
        count int32
        upstream http.Handler
}

func (s *server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
        if strings.HasPrefix(r.URL.Path, "/computeMetadata/v1/instance/service-accounts/") && strings.HasSuffix(r.URL.Path, "/token") && atomic.AddInt32(&s.count, 1) % int32(*modulo) == 0 {
                http.Error(w, "Too many requests", 429)
        } else {
                s.upstream.ServeHTTP(w, r)
        }
}

func main() {
        flag.Parse()
        url, err := url.Parse(*upstream)
        if err != nil {
                log.Fatal(err)
        }
        proxy := httputil.NewSingleHostReverseProxy(url)
        proxy.Director = func(r *http.Request) {
                // The GCE metadata server rejects requests with an X-Forwarded-for header
                r.Header[http.CanonicalHeaderKey("X-Forwarded-For")] = nil
        }
        s := server {
                upstream: proxy,
        }

        http.Handle("/", &s)
        log.Fatal(http.ListenAndServe(*listen, nil))
}