explore: Use fewer cookies

Originally the site was largely stateless and I used small cookies to
smuggle things around, but this means I can't really use credentials
without sharing them with the client, so let's do away with that.

Signed-off-by: Jon Johnson <jon.johnson@chainguard.dev>

Author: jonjohnsonjr

internal/explore/cookies.go

@@ -1,9 +1,6 @@
 package explore
 
 import (
-	"encoding/base64"
-	"encoding/json"
-	"log"
 	"net/http"
 	"time"
 
@@ -13,13 +10,6 @@ import (
 	"github.com/google/go-containerregistry/pkg/v1/remote/transport"
 )
 
-type CookieValue struct {
-	Reg           string
-	PingResp      *transport.PingResp
-	Repo          string
-	TokenResponse *transport.TokenResponse
-}
-
 type RedirectCookie struct {
 	Digest string
 	Url    string
@@ -33,102 +23,60 @@ func (h *handler) transportFromCookie(w http.ResponseWriter, r *http.Request, re
 	scopes := []string{parsed.Scope(transport.PullScope)}
 	reg := parsed.Registry
 
-	var (
-		pr  *transport.PingResp
-		tok *transport.TokenResponse
-	)
-	if regCookie, err := r.Cookie("registry_token"); err == nil {
-		b, err := base64.URLEncoding.DecodeString(regCookie.Value)
-		if err != nil {
-			return nil, err
-		}
-		var v CookieValue
-		if err := json.Unmarshal(b, &v); err != nil {
-			return nil, err
-		}
-		if v.Reg == reg.String() {
-			pr = v.PingResp
-			if v.Repo == repo {
-				tok = v.TokenResponse
-			}
-		}
-	}
-
 	t := remote.DefaultTransport
 	t = transport.NewRetry(t)
 	t = transport.NewUserAgent(t, h.userAgent)
 	if r.URL.Query().Get("trace") != "" {
 		t = transport.NewTracer(t)
 	}
 
-	if pr == nil {
-		if cpr, ok := h.pings[reg.String()]; ok {
-			if debug {
-				log.Printf("cached ping: %v", cpr)
-			}
-			pr = cpr
-		} else {
-			if debug {
-				log.Printf("pinging %s", reg.String())
-			}
-			pr, err = transport.Ping(r.Context(), reg, t)
-			if err != nil {
-				return nil, err
-			}
-			h.pings[reg.String()] = pr
-		}
-	}
-
-	if tok == nil {
-		if debug {
-			log.Printf("getting token %s", reg.String())
-		}
-		t, tok, err = transport.NewBearer(r.Context(), pr, reg, auth, t, scopes)
+	h.Lock()
+	pr, ok := h.pings[reg.String()]
+	h.Unlock()
+	if !ok {
+		pr, err = transport.Ping(r.Context(), reg, t)
 		if err != nil {
 			return nil, err
 		}
+		h.Lock()
+		h.pings[reg.String()] = pr
+		h.Unlock()
+	}
 
-		// Probably no auth needed.
-		if tok == nil {
-			return t, nil
-		}
+	h.Lock()
+	tok, ok := h.tokens[parsed.String()]
+	h.Unlock()
+	if ok && !tok.Expires.Before(time.Now().Add(30*time.Second)) {
+		// If this won't expire within 30 seconds, reuse it.
+		return transport.OldBearer(pr, tok.TokenResponse, reg, auth, t, scopes)
+	}
 
-		// Clear this to make cookies smaller.
-		tok.AccessToken = ""
+	// We don't have a cached token or it's expired (or about to), so get a new one.
+	rt, tr, err := transport.NewBearer(r.Context(), pr, reg, auth, t, scopes)
+	if err != nil {
+		return nil, err
+	}
 
-		v := &CookieValue{
-			Reg:           reg.String(),
-			PingResp:      pr,
-			Repo:          repo,
-			TokenResponse: tok,
-		}
-		b, err := json.Marshal(v)
-		if err != nil {
-			return nil, err
-		}
-		cv := base64.URLEncoding.EncodeToString(b)
-		cookie := &http.Cookie{
-			Name:     "registry_token",
-			Value:    cv,
-			Secure:   true,
-			HttpOnly: true,
-			SameSite: http.SameSiteLaxMode,
-		}
-		if tok.ExpiresIn == 0 {
-			tok.ExpiresIn = 60
-		}
-		exp := time.Now().Add(time.Second * time.Duration(tok.ExpiresIn))
-		cookie.Expires = exp
-		http.SetCookie(w, cookie)
-	} else {
-		if debug {
-			log.Printf("restoring bearer %s", reg.String())
-		}
-		t, err = transport.OldBearer(pr, tok, reg, auth, t, scopes)
-		if err != nil {
-			return nil, err
-		}
+	// Probably no auth needed.
+	if tr == nil {
+		return rt, nil
+	}
+
+	// Clear this to make cache smaller (sometimes this duplicates Token).
+	tr.AccessToken = ""
+
+	if tr.ExpiresIn == 0 {
+		tr.ExpiresIn = 60
 	}
+	exp := time.Now().Add(time.Second * time.Duration(tr.ExpiresIn))
+	tok = token{
+		TokenResponse: tr,
+		Expires:       exp,
+	}
+
+	h.Lock()
+	h.tokens[parsed.String()] = tok
+	h.Unlock()
 
-	return t, nil
+	return rt, nil
 }

internal/explore/explore.go

@@ -47,6 +47,11 @@ import (
 const tooBig = 1 << 22
 const respTooBig = 1 << 25
 
+type token struct {
+	Expires       time.Time
+	TokenResponse *transport.TokenResponse
+}
+
 type handler struct {
 	mux       http.Handler
 	keychain  authn.Keychain
@@ -58,6 +63,12 @@ type handler struct {
 	// reg.String() -> ping resp
 	pings map[string]*transport.PingResp
 
+	// repo.String() -> token
+	tokens map[string]token
+
+	// blob.Digest() -> url
+	redirects map[string]string
+
 	tocCache   cache
 	indexCache cache
 
@@ -86,6 +97,8 @@ func New(opts ...Option) http.Handler {
 	h := handler{
 		manifests:  map[string]*remote.Descriptor{},
 		pings:      map[string]*transport.PingResp{},
+		tokens:     map[string]token{},
+		redirects:  map[string]string{},
 		sawTags:    map[string][]string{},
 		inflight:   map[string]*soci.Indexer{},
 		tocCache:   buildTocCache(),