Refactor Terraform and Hex handlers to support multiple registries per organization

internal/handlers/hex_organization.go

@@ -13,12 +13,17 @@ import (
 
 // HexOrganizationHandler handles requests to repo.hex.pm, adding auth.
 type HexOrganizationHandler struct {
-	orgTokens map[string]string
+	credentials []hexOrganizationCredentials
+}
+
+type hexOrganizationCredentials struct {
+	organization string
+	token        string
 }
 
 // NewHexOrganizationHandler returns a new HexOrganizationHandler.
 func NewHexOrganizationHandler(creds config.Credentials) *HexOrganizationHandler {
-	handler := HexOrganizationHandler{orgTokens: map[string]string{}}
+	handler := HexOrganizationHandler{credentials: []hexOrganizationCredentials{}}
 
 	for _, cred := range creds {
 		if cred["type"] != "hex_organization" {
@@ -31,7 +36,11 @@ func NewHexOrganizationHandler(creds config.Credentials) *HexOrganizationHandler
 			continue
 		}
 
-		handler.orgTokens[org] = token
+		hexCred := hexOrganizationCredentials{
+			organization: org,
+			token:        token,
+		}
+		handler.credentials = append(handler.credentials, hexCred)
 	}
 
 	return &handler
@@ -52,13 +61,14 @@ func (h *HexOrganizationHandler) HandleRequest(req *http.Request, ctx *goproxy.P
 		return req, nil
 	}
 
-	token, ok := h.orgTokens[pathParts[1]]
-	if !ok {
-		return req, nil
+	reqOrg := pathParts[1]
+	for _, cred := range h.credentials {
+		if cred.organization == reqOrg {
+			logging.RequestLogf(ctx, "* authenticating hex request (org: %s)", reqOrg)
+			req.Header.Set("authorization", cred.token)
+			return req, nil
+		}
 	}
 
-	logging.RequestLogf(ctx, "* authenticating hex request (org: %s)", pathParts[1])
-	req.Header.Set("authorization", token)
-
 	return req, nil
 }

internal/handlers/terraform_registry.go

@@ -13,14 +13,20 @@ import (
 )
 
 type TerraformRegistryHandler struct {
-	credentials     map[string]string
+	credentials     []terraformRegistryCredentials
 	oidcCredentials map[string]*oidc.OIDCCredential
 	mutex           sync.RWMutex
 }
 
+type terraformRegistryCredentials struct {
+	host  string
+	url   string
+	token string
+}
+
 func NewTerraformRegistryHandler(credentials config.Credentials) *TerraformRegistryHandler {
 	handler := TerraformRegistryHandler{
-		credentials:     make(map[string]string),
+		credentials:     []terraformRegistryCredentials{},
 		oidcCredentials: make(map[string]*oidc.OIDCCredential),
 	}
 
@@ -40,7 +46,12 @@ func NewTerraformRegistryHandler(credentials config.Credentials) *TerraformRegis
 			continue
 		}
 
-		handler.credentials[host] = credential.GetString("token")
+		terraformCred := terraformRegistryCredentials{
+			host:  host,
+			url:   credential.GetString("url"),
+			token: credential.GetString("token"),
+		}
+		handler.credentials = append(handler.credentials, terraformCred)
 	}
 	return &handler
 }
@@ -56,15 +67,22 @@ func (h *TerraformRegistryHandler) HandleRequest(request *http.Request, context
 	}
 
 	// Fall back to static credentials
-	host := request.URL.Hostname()
-	token, ok := h.credentials[host]
-
-	if !ok {
-		return request, nil
+	for _, cred := range h.credentials {
+		// Match by URL first (more specific), then by host
+		if cred.url != "" {
+			if helpers.UrlMatchesRequest(request, cred.url, true) {
+				logging.RequestLogf(context, "* authenticating terraform registry request (url: %s)", cred.url)
+				request.Header.Set("Authorization", "Bearer "+cred.token)
+				return request, nil
+			}
+		} else if cred.host != "" {
+			if helpers.CheckHost(request, cred.host) {
+				logging.RequestLogf(context, "* authenticating terraform registry request (host: %s)", cred.host)
+				request.Header.Set("Authorization", "Bearer "+cred.token)
+				return request, nil
+			}
+		}
 	}
 
-	logging.RequestLogf(context, "* authenticating terraform registry request (host: %s)", host)
-	request.Header.Set("Authorization", "Bearer "+token)
-
 	return request, nil
 }