From e2fbe18d52ebaf24034dd48bef9af8c09ab142cf Mon Sep 17 00:00:00 2001
From: Taylor Raack <traack@fastly.com>
Date: Fri, 5 Jan 2024 08:34:26 -1000
Subject: [PATCH] Update product cache to use /entitlements endpoint

---
 pkg/api/product_cache.go      | 91 +++++++++++++++++++++++------------
 pkg/api/product_cache_test.go | 52 +++++++++-----------
 2 files changed, 82 insertions(+), 61 deletions(-)

diff --git a/pkg/api/product_cache.go b/pkg/api/product_cache.go
index 73ef221..971280a 100644
--- a/pkg/api/product_cache.go
+++ b/pkg/api/product_cache.go
@@ -27,12 +27,20 @@ const (
 // Products is the slice of available products supported by real-time stats.
 var Products = []string{ProductDefault, ProductOriginInspector, ProductDomainInspector}
 
-// Product models the response from the Fastly Product Entitlement API.
-type Product struct {
-	HasAccess bool `json:"has_access"`
-	Product   struct {
-		Name string `json:"id"`
-	} `json:"product"`
+type response struct {
+	Customers *[]customer `json:"customers,omitempty"`
+}
+
+type customer struct {
+	Contracts *[]contract `json:"contracts,omitempty"`
+}
+
+type contract struct {
+	Items *[]item `json:"items,omitempty"`
+}
+
+type item struct {
+	ProductID *string `json:"product_id,omitempty"`
 }
 
 // ProductCache fetches product information from the Fastly Product Entitlement API
@@ -59,41 +67,60 @@ func NewProductCache(client HTTPClient, token string, logger log.Logger) *Produc
 
 // Refresh requests data from the Fastly API and stores data in the cache.
 func (p *ProductCache) Refresh(ctx context.Context) error {
-	for _, product := range Products {
-		if product == ProductDefault {
-			continue
-		}
-		uri := fmt.Sprintf("https://api.fastly.com/entitled-products/%s", product)
 
-		req, err := http.NewRequestWithContext(ctx, "GET", uri, nil)
-		if err != nil {
-			return fmt.Errorf("error constructing API product request: %w", err)
-		}
+	req, err := http.NewRequestWithContext(ctx, "GET", "https://api.fastly.com/entitlements", nil)
+	if err != nil {
+		return fmt.Errorf("error constructing API product request: %w", err)
+	}
 
-		req.Header.Set("Fastly-Key", p.token)
-		req.Header.Set("Accept", "application/json")
-		resp, err := p.client.Do(req)
-		if err != nil {
-			return fmt.Errorf("error executing API product request: %w", err)
-		}
-		defer resp.Body.Close()
+	req.Header.Set("Fastly-Key", p.token)
+	req.Header.Set("Accept", "application/json")
+	resp, err := p.client.Do(req)
+	if err != nil {
+		return fmt.Errorf("error executing API product request: %w", err)
+	}
+	defer resp.Body.Close()
 
-		if resp.StatusCode != http.StatusOK {
-			return NewError(resp)
-		}
+	if resp.StatusCode != http.StatusOK {
+		return NewError(resp)
+	}
 
-		var response Product
+	var response response
 
-		if err := json.NewDecoder(resp.Body).Decode(&response); err != nil {
-			return fmt.Errorf("error decoding API product response: %w", err)
-		}
+	if err := json.NewDecoder(resp.Body).Decode(&response); err != nil {
+		return fmt.Errorf("error decoding API product response: %w", err)
+	}
 
-		level.Debug(p.logger).Log("product", response.Product.Name, "hasAccess", response.HasAccess)
+	activeProducts := make(map[string]interface{})
+
+	if response.Customers != nil {
+		for _, customer := range *response.Customers {
+			if customer.Contracts == nil {
+				continue
+			}
+			for _, contract := range *customer.Contracts {
+				if contract.Items == nil {
+					continue
+				}
+				for _, item := range *contract.Items {
+					if item.ProductID == nil {
+						continue
+					}
+					activeProducts[*item.ProductID] = true
+				}
+			}
+		}
+	}
 
+	for _, product := range Products {
+		if product == ProductDefault {
+			continue
+		}
+		_, hasAccess := activeProducts[product]
+		level.Debug(p.logger).Log("product", product, "hasAccess", hasAccess)
 		p.mtx.Lock()
-		p.products[response.Product.Name] = response.HasAccess
+		p.products[product] = hasAccess
 		p.mtx.Unlock()
-
 	}
 
 	return nil
diff --git a/pkg/api/product_cache_test.go b/pkg/api/product_cache_test.go
index 352ae2c..726e1a8 100644
--- a/pkg/api/product_cache_test.go
+++ b/pkg/api/product_cache_test.go
@@ -21,7 +21,7 @@ func TestProductCache(t *testing.T) {
 	}{
 		{
 			name:    "success",
-			client:  newSequentialResponseClient(productsResponseOne, productsResponseTwo),
+			client:  newSequentialResponseClient(productsResponse),
 			wantErr: nil,
 			wantProds: map[string]bool{
 				"origin_inspector": true,
@@ -55,34 +55,28 @@ func TestProductCache(t *testing.T) {
 	}
 }
 
-const productsResponseOne = `
+// only products that the customer is entitled to are returned from the /entitlements endpoint
+const productsResponse = `
 {
-  "product": {
-    "id": "origin_inspector",
-    "object": "product"
-  },
-  "has_access": true,
-  "access_level": "Origin_Inspector",
-  "has_permission_to_enable": false,
-  "has_permission_to_disable": true,
-  "_links": {
-    "self": ""
-  }
-}
-`
-
-const productsResponseTwo = `
-{
-  "product": {
-    "id": "domain_inspector",
-    "object": "product"
-  },
-  "has_access": false,
-  "access_level": "Domain_Inspector",
-  "has_permission_to_enable": false,
-  "has_permission_to_disable": true,
-  "_links": {
-    "self": ""
-  }
+  "customers": [
+	{
+	  "contracts": [
+		{
+		  "items": [
+			{
+			  "product_id": "origin_inspector"
+			},
+			{
+			  "other_key": "other"
+			}
+		  ]
+		},
+		{
+		  "other_key_2": [
+		  ]
+		}
+	  ]
+	}
+  ]
 }
 `