support non-gcr registries

- GCR_HOST --> REGISTRY_HOST
- PROJECT_ID --> REPO_PREFIX
- add support for TLS_KEY, TLS_CERT files to listen with tls
- add /_token endpoint to perform docker registry token handshake protocol

Signed-off-by: Ahmet Alp Balkan <ahmetb@google.com>

.gitignore

@@ -1 +1,4 @@
 key.json
+*.pem
+*.crt
+*.cert

README.md

@@ -24,7 +24,7 @@ Then, push to a registry like:
 
     docker push gcr.io/[YOUR_PROJECT]/gcr-proxy
 
-## Deploying (to Google Cloud Run)
+## Deploying (to Google Cloud Run) for GCR.io
 
 You can easily deploy this as a serverless container to [Google Cloud Run][run].
 This handles many of the heavy-lifting for you.
@@ -44,11 +44,11 @@ of the GCR registry you want to expose publicly:
 gcloud beta run deploy \
     --allow-unauthenticated \
     --image "[IMAGE]" \
-    --set-env-vars "GCR_PROJECT_ID=[PROJECT_ID]"
+    --set-env-vars "REGISTRY_HOST=gcr.io,REPO_PREFIX=[GCP_PROJECT_ID]"
 ```
 
-> This will deploy a proxy for your `gcr.io/[PROJECT_ID]` public registry. If
-> your GCR registry is private, see the section below on "Exposing private
+> This will deploy a proxy for your `gcr.io/[GCP_PROJECT_ID]` public registry.
+> If your GCR registry is private, see the section below on "Exposing private
 > registries".
 
 Then create a domain mapping by running (replace the `--domain` value):
@@ -69,7 +69,22 @@ some 15-20 minutes to actually provision TLS certificates for your domain name.
 Kubernetes, obtain a valid TLS certificate for your domain name, and make it
 publicly accessible.
 
-### Exposing private registries publicly
+### Using with other Docker Registries
+
+If you set `REGISTRY_HOST` and `REGISTRY_PREFIX` environment variables, you can
+also use this proxy for other docker registries.
+
+For example, to proxy `docker pull ahmet/example` to Docker Hub, specify
+environment variables:
+
+- `REGISTRY_HOST=index.docker.io`
+- `REPO_PREFIX=ahmet`
+
+> **Note:** This is not tested with registries other than Docker Hub and GCR.io.
+> If you can make it work with Azure Container Registry or AWS Elastic Container
+> Registry, contribute examples here.
+
+### Exposing private registries publicly (GCR.io)
 
 > ⚠️ This will make images in your private GCR registries publicly accessible on
 > the internet.
@@ -95,17 +110,16 @@ publicly accessible.
 
    You need to rebuild and deploy the updated image.
 
-### Advanced Customization
+### Configuration
 
 While deploying, you can set additional environment variables for customization:
 
-- **`GCR_HOST`**: defaults to `gcr.io`.
-- **`DISABLE_BROWSER_REDIRECTS`**: if you set this variable to any value,
-  visiting `example.com/image` on this browser will not redirect to
-  `gcr.io/[PROJECT_ID/image` to allow your users to browse the image on GCR. If
-  you're exposing private registries, you might want to set this variable.
-- **`GOOGLE_APPLICATION_CREDENTIALS`**: path to the IAM service account JSON key
-  file to expose the private GCR registries publicly.
+| Key | Value |
+|-----|-------|
+| `REGISTRY_HOST` | specify  hostname for target registry, e.g. `gcr.io`. |
+| `DISABLE_BROWSER_REDIRECTS` |  if you set this variable to any value,   visiting `example.com/image` on this browser will not redirect to  `[REGISTRY_HOST]/[REPO_PREFIX]/image` to allow your users to browse the image on GCR. If you're exposing private registries, you might want to set this variable. |
+| `AUTH_HEADER` | The `Authentication: [...]` header’s value to authenticate to the target registry |
+| `GOOGLE_APPLICATION_CREDENTIALS` | (For `gcr.io`) Path to the IAM service account JSON key  file to expose the private GCR registries publicly. |
 
 -----
 

main.go

@@ -23,8 +23,10 @@ import (
 	"log"
 	"net/http"
 	"net/http/httputil"
+	"net/url"
 	"os"
 	"regexp"
+	"strings"
 )
 
 const (
@@ -33,12 +35,13 @@ const (
 
 var (
 	re                 = regexp.MustCompile(`^/v2/`)
+	realm              = regexp.MustCompile(`realm="(.*?)"`)
 	ctxKeyOriginalHost = struct{}{}
 )
 
-type gcrConfig struct {
-	host      string
-	projectID string
+type registryConfig struct {
+	host       string
+	repoPrefix string
 }
 
 func main() {
@@ -48,42 +51,79 @@ func main() {
 	}
 	browserRedirects := os.Getenv("DISABLE_BROWSER_REDIRECTS") == ""
 
-	gcrHost := defaultGCRHost
-	if v := os.Getenv("GCR_HOST"); v != "" {
-		gcrHost = v
+	registryHost := os.Getenv("REGISTRY_HOST")
+	if registryHost == "" {
+		log.Fatal("REGISTRY_HOST environment variable not specified (example: gcr.io)")
 	}
-	gcrProjectID := os.Getenv("GCR_PROJECT_ID")
-	if gcrProjectID == "" {
-		log.Fatal("GCR_PROJECT_ID environment variable not specified")
+	repoPrefix := os.Getenv("REPO_PREFIX")
+	if repoPrefix == "" {
+		log.Fatal("REPO_PREFIX environment variable not specified")
 	}
 
-	gcr := gcrConfig{
-		host:      gcrHost,
-		projectID: gcrProjectID,
+	reg := registryConfig{
+		host:       registryHost,
+		repoPrefix: repoPrefix,
 	}
 
-	var authHeader string
-	if keyPath := os.Getenv("GOOGLE_APPLICATION_CREDENTIALS"); keyPath != "" {
-		b, err := ioutil.ReadFile(keyPath)
+	tokenEndpoint, err := discoverTokenService(reg.host)
+	if err != nil {
+		log.Fatalf("target registry's token endpoint could not be discovered: %+v", err)
+	}
+	log.Printf("discovered token endpoint for backend registry: %s", tokenEndpoint)
+
+	var auth authenticator
+	if basic := os.Getenv("AUTH_HEADER"); basic != "" {
+		auth = authHeader(basic)
+	} else if gcpKey := os.Getenv("GOOGLE_APPLICATION_CREDENTIALS"); gcpKey != "" {
+		b, err := ioutil.ReadFile(gcpKey)
 		if err != nil {
-			log.Fatalf("could not read key file from %s: %+v", keyPath, err)
+			log.Fatalf("could not read key file from %s: %+v", gcpKey, err)
 		}
 		log.Printf("using specified service account json key to authenticate proxied requests")
-		authHeader = "Basic " + base64.StdEncoding.EncodeToString([]byte(fmt.Sprintf("_json_key:%s", string(b))))
+		auth = authHeader("Basic " + base64.StdEncoding.EncodeToString([]byte(fmt.Sprintf("_json_key:%s", string(b)))))
 	}
 
-	addr := ":" + port
+	mux := http.NewServeMux()
 	if browserRedirects {
-		http.Handle("/", browserRedirectHandler(gcr))
+		mux.Handle("/", browserRedirectHandler(reg))
+	}
+	if tokenEndpoint != "" {
+		mux.Handle("/_token", tokenProxyHandler(tokenEndpoint, repoPrefix))
 	}
-	http.Handle("/v2/", captureHostHeader(registryAPIMux(gcr, authHeader)))
+	mux.Handle("/v2/", registryAPIProxy(reg, auth))
+
+	addr := ":" + port
+	handler := captureHostHeader(mux)
 	log.Printf("starting to listen on %s", addr)
-	if err := http.ListenAndServe(addr, nil); err != nil && err != http.ErrServerClosed {
+	if cert, key := os.Getenv("TLS_CERT"), os.Getenv("TLS_KEY"); cert != "" && key != "" {
+		err = http.ListenAndServeTLS(addr, cert, key, handler)
+	} else {
+		err = http.ListenAndServe(addr, handler)
+	}
+	if err != http.ErrServerClosed {
 		log.Fatalf("listen error: %+v", err)
 	}
+
 	log.Printf("server shutdown successfully")
 }
 
+func discoverTokenService(registryHost string) (string, error) {
+	url := fmt.Sprintf("https://%s/v2/", registryHost)
+	resp, err := http.Get(url)
+	if err != nil {
+		return "", fmt.Errorf("failed to query the registry host %s: %+v", registryHost, err)
+	}
+	hdr := resp.Header.Get("www-authenticate")
+	if hdr == "" {
+		return "", fmt.Errorf("www-authenticate header not returned from %s, cannot locate token endpoint", url)
+	}
+	matches := realm.FindStringSubmatch(hdr)
+	if len(matches) == 0 {
+		return "", fmt.Errorf("cannot locate 'realm' in %s response header www-authenticate: %s", url, hdr)
+	}
+	return matches[1], nil
+}
+
 // captureHostHeader is a middleware to capture Host header in a context key.
 func captureHostHeader(next http.Handler) http.Handler {
 	return http.HandlerFunc(func(rw http.ResponseWriter, req *http.Request) {
@@ -93,33 +133,50 @@ func captureHostHeader(next http.Handler) http.Handler {
 	})
 }
 
+// tokenProxyHandler proxies the token requests to the specified token service.
+// It adjusts the ?scope= parameter in the query from "repository:foo:..." to
+// "repository:repoPrefix/foo:.." and reverse proxies the query to the specified
+// tokenEndpoint.
+func tokenProxyHandler(tokenEndpoint, repoPrefix string) http.HandlerFunc {
+	return (&httputil.ReverseProxy{
+		Director: func(r *http.Request) {
+			orig := r.URL.String()
+
+			q := r.URL.Query()
+			scope := q.Get("scope")
+			if scope == "" {
+				return
+			}
+			newScope := strings.Replace(scope, "repository:", fmt.Sprintf("repository:%s/", repoPrefix), 1)
+			q.Set("scope", newScope)
+			u, _ := url.Parse(tokenEndpoint)
+			u.RawQuery = q.Encode()
+			r.URL = u
+			log.Printf("tokenProxyHandler: rewrote url:%s into:%s", orig, r.URL)
+			r.Host = u.Host
+		},
+	}).ServeHTTP
+}
+
 // browserRedirectHandler redirects a request like example.com/my-image to
-// gcr.io/my-image, which shows a public UI for browsing the registry.
-func browserRedirectHandler(c gcrConfig) http.HandlerFunc {
+// REGISTRY_HOST/my-image, which shows a public UI for browsing the registry.
+// This works only on registries that support a web UI when the image name is
+// entered into the browser, like GCR (gcr.io/google-containers/busybox).
+func browserRedirectHandler(cfg registryConfig) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
-		url := fmt.Sprintf("https://%s/%s%s", c.host, c.projectID, r.RequestURI)
+		url := fmt.Sprintf("https://%s/%s%s", cfg.host, cfg.repoPrefix, r.RequestURI)
 		http.Redirect(w, r, url, http.StatusTemporaryRedirect)
 	}
 }
 
-// registryAPIMux returns a handler for Docker Registry v2 API requests
-// (/v2/). Request to path=/v2/ is handled-locally, other /v2/* requests are
-// proxied back to GCR endpoint.
-func registryAPIMux(c gcrConfig, authHeader string) http.HandlerFunc {
-	reverseProxy := &httputil.ReverseProxy{
-		Director: rewriteRegistryV2URL(c),
-		Transport: &gcrRoundtripper{
-			authHeader: authHeader,
+// registryAPIProxy returns a reverse proxy to the specified registry.
+func registryAPIProxy(cfg registryConfig, auth authenticator) http.HandlerFunc {
+	return (&httputil.ReverseProxy{
+		Director: rewriteRegistryV2URL(cfg),
+		Transport: &registryRoundtripper{
+			auth: auth,
 		},
-	}
-
-	return func(w http.ResponseWriter, req *http.Request) {
-		if req.URL.Path == "/v2/" {
-			handleRegistryAPIVersion(w, req)
-			return
-		}
-		reverseProxy.ServeHTTP(w, req)
-	}
+	}).ServeHTTP
 }
 
 // handleRegistryAPIVersion signals docker-registry v2 API on /v2/ endpoint.
@@ -130,43 +187,66 @@ func handleRegistryAPIVersion(w http.ResponseWriter, r *http.Request) {
 
 // rewriteRegistryV2URL rewrites request.URL like /v2/* that come into the server
 // into https://[GCR_HOST]/v2/[PROJECT_ID]/*. It leaves /v2/ as is.
-func rewriteRegistryV2URL(c gcrConfig) func(*http.Request) {
+func rewriteRegistryV2URL(c registryConfig) func(*http.Request) {
 	return func(req *http.Request) {
 		u := req.URL.String()
 		req.Host = c.host
 		req.URL.Scheme = "https"
 		req.URL.Host = c.host
 		if req.URL.Path != "/v2/" {
-			req.URL.Path = re.ReplaceAllString(req.URL.Path, fmt.Sprintf("/v2/%s/", c.projectID))
+			req.URL.Path = re.ReplaceAllString(req.URL.Path, fmt.Sprintf("/v2/%s/", c.repoPrefix))
 		}
 		log.Printf("rewrote url: %s into %s", u, req.URL)
 	}
 }
 
-type gcrRoundtripper struct {
-	authHeader string
+type registryRoundtripper struct {
+	auth authenticator
 }
 
-func (g *gcrRoundtripper) RoundTrip(req *http.Request) (*http.Response, error) {
+func (rrt *registryRoundtripper) RoundTrip(req *http.Request) (*http.Response, error) {
 	log.Printf("request received. url=%s", req.URL)
 
-	if g.authHeader != "" {
-		req.Header.Set("Authorization", g.authHeader)
+	if rrt.auth != nil {
+		req.Header.Set("Authorization", rrt.auth.AuthHeader())
 	}
 
+	origHost := req.Context().Value(ctxKeyOriginalHost).(string)
 	if ua := req.Header.Get("user-agent"); ua != "" {
-		origHost := req.Context().Value(ctxKeyOriginalHost).(string)
 		req.Header.Set("user-agent", "gcr-proxy/0.1 customDomain/"+origHost+" "+ua)
 	}
 
-	// TODO(ahmetb) remove after internal bug 129780113 is fixed.
+	// TODO(ahmetb) remove after Google internal bug 129780113 is fixed.
 	req.Header.Set("accept", "*/*")
 
 	resp, err := http.DefaultTransport.RoundTrip(req)
 	if err == nil {
 		log.Printf("request completed (status=%d) url=%s", resp.StatusCode, req.URL)
 	} else {
 		log.Printf("request failed with error: %+v", err)
+		return nil, err
 	}
-	return resp, err
+	updateTokenEndpoint(resp, origHost)
+	return resp, nil
 }
+
+// updateTokenEndpoint modifies the response header like:
+//    Www-Authenticate: Bearer realm="https://auth.docker.io/token",service="registry.docker.io"
+// to point to the https://host/token endpoint to force using local token
+// endpoint proxy.
+func updateTokenEndpoint(resp *http.Response, host string) {
+	v := resp.Header.Get("www-authenticate")
+	if v == "" {
+		return
+	}
+	cur := fmt.Sprintf("https://%s/_token", host)
+	resp.Header.Set("www-authenticate", realm.ReplaceAllString(v, fmt.Sprintf(`realm="%s"`, cur)))
+}
+
+type authenticator interface {
+	AuthHeader() string
+}
+
+type authHeader string
+
+func (b authHeader) AuthHeader() string { return string(b) }