Add JIT access fallback for GitHub API requests

internal/handlers/github_api.go

@@ -13,51 +13,86 @@ import (
 	"github.com/dependabot/proxy/internal/ctxdata"
 	"github.com/dependabot/proxy/internal/helpers"
 	"github.com/dependabot/proxy/internal/logging"
+	"github.com/dependabot/proxy/internal/threadsafe"
 )
 
 // GitHubAPIHandler handles requests destined for the GitHub API, adding auth
 // This allows git credentials for "github.com" to apply to "api.github.com" and
 // will allow git credentials for "<tenant>.ghe.com" to apply to "api.<tenant>.ghe.com" in Proxima.
 type GitHubAPIHandler struct {
-	credentials *gitCredentialsMap
+	credentials     *gitCredentialsMap
+	jitAccessByHost map[string]jitAccessConfig
+	client          ScopeRequester
+
+	reposAlreadyTried *threadsafe.Map[string, struct{}]
 }
 
 const ghAPIAddedAuthCtxKey = "gh-api.added-auth"
 const reservedProximaIdentity = "proxima-service-identity"
 
 // NewGitHubAPIHandler returns a new GitHubAPIHandler, extracting the app
 // access token from the array of credentials
-func NewGitHubAPIHandler(creds config.Credentials) *GitHubAPIHandler {
+func NewGitHubAPIHandler(creds config.Credentials, client ScopeRequester) *GitHubAPIHandler {
 	handler := GitHubAPIHandler{
-		credentials: newGitCredentialsMap(),
+		credentials:       newGitCredentialsMap(),
+		jitAccessByHost:   map[string]jitAccessConfig{},
+		client:            client,
+		reposAlreadyTried: threadsafe.NewMap[string, struct{}](),
 	}
+	hasGitSourceCredentials := false
 
 	for _, cred := range creds {
-		host := cred.Host()
-		if host == "" {
-			continue
-		}
-		if cred["type"] != "git_source" || (host != "github.com" && !(strings.HasSuffix(fmt.Sprint(host), ".ghe.com"))) {
-			continue
+		switch cred["type"] {
+		case "git_source":
+			host := cred.Host()
+			if host == "" {
+				continue
+			}
+			if host != "github.com" && !strings.HasSuffix(fmt.Sprint(host), ".ghe.com") {
+				continue
+			}
+			handler.credentials.addGitSourceCredentials("api."+host, cred)
+			hasGitSourceCredentials = true
+		case "jit_access":
+			if client != nil {
+				handler.addJITAccess(cred)
+			}
 		}
-		handler.credentials.addGitSourceCredentials("api."+host, cred)
 	}
 
-	if len(handler.credentials.data) == 0 {
+	if !hasGitSourceCredentials && len(handler.jitAccessByHost) == 0 {
 		logrus.Warn("GitHubAPIHandler has no app access tokens")
 	}
 
 	return &handler
 }
 
+func (h *GitHubAPIHandler) addJITAccess(cred config.Credential) {
+	if cred.GetString("credential-type") != "git_source" {
+		return
+	}
+
+	host := strings.ToLower(cred.GetString("host"))
+	if host == "" {
+		return
+	}
+
+	if !strings.HasPrefix(host, "api.") {
+		host = "api." + host
+	}
+
+	h.jitAccessByHost[host] = jitAccessConfig{
+		endpoint: cred.GetString("endpoint"),
+		username: cred.GetString("username"),
+		password: cred.GetString("password"),
+	}
+}
+
 // HandleRequest adds auth to a GitHub API request
 func (h *GitHubAPIHandler) HandleRequest(req *http.Request, ctx *goproxy.ProxyCtx) (*http.Request, *http.Response) {
 	if !h.isHandledGitHubAPIRequest(req) {
 		return req, nil
 	}
-	if len(h.credentials.data) == 0 {
-		return req, nil
-	}
 
 	host := helpers.GetHost(req)
 	creds := getCredentialsForRequest(req, h.credentials, gitHubAPIExtractOrgAndRepo)
@@ -84,15 +119,23 @@ func (h *GitHubAPIHandler) HandleResponse(rsp *http.Response, ctx *goproxy.Proxy
 	if rsp == nil {
 		return rsp
 	}
-	if addedAuth, ok := ctxdata.GetBool(ctx, ghAPIAddedAuthCtxKey); !ok || !addedAuth {
+	if !h.isHandledGitHubAPIRequest(ctx.Req) {
 		return rsp
 	}
-	if !h.isHandledGitHubAPIRequest(ctx.Req) {
+	addedAuth, _ := ctxdata.GetBool(ctx, ghAPIAddedAuthCtxKey)
+	if !addedAuth {
 		return rsp
 	}
 	if !isPotentialAuthFailure(rsp.StatusCode) {
 		return rsp
 	}
+	repoKey, hasRepo := h.jitRepositoryKey(ctx.Req)
+	if hasRepo {
+		if _, ok := h.reposAlreadyTried.Get(repoKey); ok {
+			logging.RequestLogf(ctx, "* github api repository previously retried, won't retry again")
+			return rsp
+		}
+	}
 
 	username, password, reqWasAuthed := ctx.Req.BasicAuth()
 	for _, creds := range getCredentialsForRequest(ctx.Req, h.credentials, gitHubAPIExtractOrgAndRepo) {
@@ -121,9 +164,76 @@ func (h *GitHubAPIHandler) HandleResponse(rsp *http.Response, ctx *goproxy.Proxy
 		logging.RequestLogf(ctx, "* re-auth'd request returned %d, ignoring response", newRsp.StatusCode)
 		helpers.DrainAndClose(newRsp)
 	}
+
+	// All known credentials have been tried, try to JIT create access credentials.
+	jitCreds := h.getJITCredentialsForRequest(ctx)
+	if jitCreds != nil {
+		newReq := ctx.Req.Clone(ctx.Req.Context())
+		logging.RequestLogf(ctx, "* auth'd github api request failed authentication, retrying with jit access auth")
+		newReq.Header.Del("Authorization")
+		newReq.Header.Del("X-GitHub-PSI-JWT")
+		if jitCreds.username == reservedProximaIdentity {
+			newReq.Header.Set("X-GitHub-PSI-JWT", jitCreds.password)
+		} else {
+			newReq.Header.Set("Authorization", "token "+jitCreds.password)
+		}
+		newRsp, err := ctx.RoundTrip(newReq)
+		if err != nil {
+			return rsp
+		}
+
+		if !isPotentialAuthFailure(newRsp.StatusCode) {
+			helpers.DrainAndClose(rsp)
+			logging.RequestLogf(ctx, "* re-auth'd jit request returned %d, replacing response", newRsp.StatusCode)
+			return newRsp
+		}
+		logging.RequestLogf(ctx, "* re-auth'd jit request returned %d, ignoring response", newRsp.StatusCode)
+		helpers.DrainAndClose(newRsp)
+	}
+	if hasRepo {
+		h.reposAlreadyTried.Set(repoKey, struct{}{})
+	}
+
 	return rsp
 }
 
+func (h *GitHubAPIHandler) getJITCredentialsForRequest(ctx *goproxy.ProxyCtx) *gitCredentials {
+	host := helpers.GetHost(ctx.Req)
+	jitConfig := h.jitAccessByHost[host]
+	if jitConfig.endpoint == "" {
+		return nil
+	}
+
+	org, repo, ok := gitHubAPIExtractOrgAndRepo(ctx.Req.URL.Path)
+	if !ok {
+		return nil
+	}
+
+	if h.client == nil {
+		return nil
+	}
+	logging.RequestLogf(ctx, "* requesting JIT access for github api request")
+	credential, err := h.client.RequestJITAccess(ctx, jitConfig.endpoint, jitConfig.username, jitConfig.password, org, repo)
+	if credential == nil || err != nil {
+		return nil
+	}
+
+	repoNWO := fmt.Sprintf("%s/%s", org, repo)
+
+	// Add the returned credentials to the beginning of the repo-scoped list, so that
+	// they are prioritized over existing tokens for future requests.
+	hostCreds := h.credentials.get(host)
+	return hostCreds.addToken(repoNWO, credential.GetString("username"), credential.GetString("password"), true)
+}
+
+func (h *GitHubAPIHandler) jitRepositoryKey(req *http.Request) (string, bool) {
+	org, repo, ok := gitHubAPIExtractOrgAndRepo(req.URL.Path)
+	if !ok {
+		return "", false
+	}
+	return fmt.Sprintf("%s/%s/%s", helpers.GetHost(req), org, repo), true
+}
+
 func (h *GitHubAPIHandler) isHandledGitHubAPIRequest(req *http.Request) bool {
 	return req.URL.Scheme == "https" && helpers.MethodPermitted(req, "GET", "HEAD") && helpers.CheckGitHubAPIHost(req)
 }