Refactor Terraform and Hex handlers to support multiple registries per organization

internal/handlers/hex_organization.go

@@ -13,25 +13,38 @@ import (
 
 // HexOrganizationHandler handles requests to repo.hex.pm, adding auth.
 type HexOrganizationHandler struct {
-	orgTokens map[string]string
+	credentials []hexOrganizationCredentials
+}
+
+type hexOrganizationCredentials struct {
+	organization string
+	key          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" {
 			continue
 		}
 
 		org := cred.GetString("organization")
-		token := cred.GetString("token")
-		if org == "" || token == "" {
+		// Support both "key" and "token" (backwards compatibility)
+		key := cred.GetString("key")
+		if key == "" {
+			key = cred.GetString("token")
+		}
+		if org == "" || key == "" {
 			continue
 		}
 
-		handler.orgTokens[org] = token
+		hexCred := hexOrganizationCredentials{
+			organization: org,
+			key:          key,
+		}
+		handler.credentials = append(handler.credentials, hexCred)
 	}
 
 	return &handler
@@ -52,13 +65,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.key)
+			return req, nil
+		}
 	}
 
-	logging.RequestLogf(ctx, "* authenticating hex request (org: %s)", pathParts[1])
-	req.Header.Set("authorization", token)
-
 	return req, nil
 }

internal/handlers/hex_organization_test.go

@@ -8,29 +8,29 @@ import (
 )
 
 func TestHexOrganizationHandler(t *testing.T) {
-	dependabotToken := "123"
-	deltaForceToken := "456"
+	dependabotKey := "123"
+	deltaForceKey := "456"
 	credentials := config.Credentials{
 		config.Credential{
 			"type":         "hex_organization",
 			"organization": "dependabot",
-			"token":        dependabotToken,
+			"key":          dependabotKey,
 		},
 		config.Credential{
 			"type":         "hex_organization",
 			"organization": "deltaforce",
-			"token":        deltaForceToken,
+			"key":          deltaForceKey,
 		},
 	}
 	handler := NewHexOrganizationHandler(credentials)
 
 	req := httptest.NewRequest("GET", "https://repo.hex.pm/repos/dependabot/packages/foo", nil)
 	req = handleRequestAndClose(handler, req, nil)
-	assertHasTokenAuth(t, req, "", dependabotToken, "dependabot registry request")
+	assertHasTokenAuth(t, req, "", dependabotKey, "dependabot registry request")
 
 	req = httptest.NewRequest("GET", "https://repo.hex.pm/repos/deltaforce/packages/foo", nil)
 	req = handleRequestAndClose(handler, req, nil)
-	assertHasTokenAuth(t, req, "", deltaForceToken, "deltaforce registry request")
+	assertHasTokenAuth(t, req, "", deltaForceKey, "deltaforce registry request")
 
 	// Not an org
 	req = httptest.NewRequest("GET", "https://repo.hex.pm/packages/foo", nil)
@@ -52,3 +52,36 @@ func TestHexOrganizationHandler(t *testing.T) {
 	req = handleRequestAndClose(handler, req, nil)
 	assertUnauthenticated(t, req, "post request")
 }
+
+func TestHexOrganizationHandler_BackwardsCompatibility(t *testing.T) {
+	t.Run("supports legacy token field", func(t *testing.T) {
+		credentials := config.Credentials{
+			config.Credential{
+				"type":         "hex_organization",
+				"organization": "legacy-org",
+				"token":        "legacy-token",
+			},
+		}
+		handler := NewHexOrganizationHandler(credentials)
+
+		req := httptest.NewRequest("GET", "https://repo.hex.pm/repos/legacy-org/packages/foo", nil)
+		req = handleRequestAndClose(handler, req, nil)
+		assertHasTokenAuth(t, req, "", "legacy-token", "should support legacy token field")
+	})
+
+	t.Run("key takes precedence over token", func(t *testing.T) {
+		credentials := config.Credentials{
+			config.Credential{
+				"type":         "hex_organization",
+				"organization": "test-org",
+				"key":          "new-key",
+				"token":        "old-token",
+			},
+		}
+		handler := NewHexOrganizationHandler(credentials)
+
+		req := httptest.NewRequest("GET", "https://repo.hex.pm/repos/test-org/packages/foo", nil)
+		req = handleRequestAndClose(handler, req, nil)
+		assertHasTokenAuth(t, req, "", "new-key", "key should take precedence over token")
+	})
+}

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,24 @@ func NewTerraformRegistryHandler(credentials config.Credentials) *TerraformRegis
 			continue
 		}
 
-		handler.credentials[host] = credential.GetString("token")
+		token := credential.GetString("token")
+		url := credential.GetString("url")
+
+		// Skip credentials with empty token or both empty host and url
+		if token == "" || (host == "" && url == "") {
+			continue
+		}
+
+		terraformCred := terraformRegistryCredentials{
+			url:   url,
+			token: token,
+		}
+		// Only set host when url is not provided to ensure URL-prefix matching
+		// takes precedence and doesn't fall back to host matching
+		if url == "" {
+			terraformCred.host = host
+		}
+		handler.credentials = append(handler.credentials, terraformCred)
 	}
 	return &handler
 }
@@ -56,15 +79,15 @@ func (h *TerraformRegistryHandler) HandleRequest(request *http.Request, context
 	}
 
 	// Fall back to static credentials
-	host := request.URL.Hostname()
-	token, ok := h.credentials[host]
+	for _, cred := range h.credentials {
+		if !helpers.UrlMatchesRequest(request, cred.url, true) && !helpers.CheckHost(request, cred.host) {
+			continue
+		}
 
-	if !ok {
+		logging.RequestLogf(context, "* authenticating terraform registry request (host: %s)", request.URL.Hostname())
+		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
 }

internal/handlers/terraform_registry_test.go

@@ -74,4 +74,40 @@ func TestTerraformRegistryHandler(t *testing.T) {
 
 		assert.Equal(t, "", request.Header.Get("Authorization"), "should be empty")
 	})
+
+	t.Run("multiple credentials on same host with different URL paths", func(t *testing.T) {
+		credentials := config.Credentials{
+			config.Credential{"type": "terraform_registry", "url": "https://terraform.example.com/org1", "token": "token-org1"},
+			config.Credential{"type": "terraform_registry", "url": "https://terraform.example.com/org2", "token": "token-org2"},
+		}
+		handler := NewTerraformRegistryHandler(credentials)
+
+		// Request to org1 path should use org1 token
+		req1 := handleRequestAndClose(handler, httptest.NewRequest("GET", "https://terraform.example.com/org1/v1/providers/foo", nil), nil)
+		assert.Equal(t, "Bearer token-org1", req1.Header.Get("Authorization"), "should use org1 token")
+
+		// Request to org2 path should use org2 token
+		req2 := handleRequestAndClose(handler, httptest.NewRequest("GET", "https://terraform.example.com/org2/v1/providers/bar", nil), nil)
+		assert.Equal(t, "Bearer token-org2", req2.Header.Get("Authorization"), "should use org2 token")
+
+		// Request to unmatched path should not be authenticated
+		req3 := handleRequestAndClose(handler, httptest.NewRequest("GET", "https://terraform.example.com/org3/v1/providers/baz", nil), nil)
+		assert.Equal(t, "", req3.Header.Get("Authorization"), "should not be authenticated")
+	})
+
+	t.Run("skips credentials with empty token", func(t *testing.T) {
+		credentials := config.Credentials{
+			config.Credential{"type": "terraform_registry", "host": "terraform.example.org", "token": ""},
+		}
+		handler := NewTerraformRegistryHandler(credentials)
+		assert.Equal(t, 0, len(handler.credentials), "should skip credential with empty token")
+	})
+
+	t.Run("skips credentials with empty host and url", func(t *testing.T) {
+		credentials := config.Credentials{
+			config.Credential{"type": "terraform_registry", "token": "some-token"},
+		}
+		handler := NewTerraformRegistryHandler(credentials)
+		assert.Equal(t, 0, len(handler.credentials), "should skip credential with empty host and url")
+	})
 }