fix: use AlwaysEnable for licenses with all features (coder#9808)

spikecurtis · web-flow · commit e9077f3bd2f0 · 2023-09-22T05:02:28.000Z
* fix: use AlwaysEnable for licenses with all features

Signed-off-by: Spike Curtis &lt;spike@coder.com&gt;

* use dbtime.Now() intead of time.Now()

Signed-off-by: Spike Curtis &lt;spike@coder.com&gt;

---------

Signed-off-by: Spike Curtis &lt;spike@coder.com&gt;
diff --git a/enterprise/coderd/license/license.go b/enterprise/coderd/license/license.go
@@ -61,6 +61,7 @@ func Entitlements(
 	}
 
 	allFeatures := false
+	allFeaturesEntitlement := codersdk.EntitlementNotEntitled
 
 	// Here we loop through licenses to detect enabled features.
 	for _, l := range licenses {
@@ -117,14 +118,15 @@ func Entitlements(
 				}
 			default:
 				entitlements.Features[featureName] = codersdk.Feature{
-					Entitlement: entitlement,
+					Entitlement: maxEntitlement(entitlements.Features[featureName].Entitlement, entitlement),
 					Enabled:     enablements[featureName] || featureName.AlwaysEnable(),
 				}
 			}
 		}
 
 		if claims.AllFeatures {
 			allFeatures = true
+			allFeaturesEntitlement = maxEntitlement(allFeaturesEntitlement, entitlement)
 		}
 		entitlements.RequireTelemetry = entitlements.RequireTelemetry || claims.RequireTelemetry
 	}
@@ -136,7 +138,8 @@ func Entitlements(
 				continue
 			}
 			feature := entitlements.Features[featureName]
-			feature.Entitlement = codersdk.EntitlementEntitled
+			feature.Entitlement = maxEntitlement(feature.Entitlement, allFeaturesEntitlement)
+			feature.Enabled = enablements[featureName] || featureName.AlwaysEnable()
 			entitlements.Features[featureName] = feature
 		}
 	}
@@ -324,3 +327,14 @@ func keyFunc(keys map[string]ed25519.PublicKey) func(*jwt.Token) (interface{}, e
 		return k, nil
 	}
 }
+
+// maxEntitlement is the "greater" entitlement between the given values
+func maxEntitlement(e1, e2 codersdk.Entitlement) codersdk.Entitlement {
+	if e1 == codersdk.EntitlementEntitled || e2 == codersdk.EntitlementEntitled {
+		return codersdk.EntitlementEntitled
+	}
+	if e1 == codersdk.EntitlementGracePeriod || e2 == codersdk.EntitlementGracePeriod {
+		return codersdk.EntitlementGracePeriod
+	}
+	return codersdk.EntitlementNotEntitled
+}
diff --git a/enterprise/coderd/license/license_test.go b/enterprise/coderd/license/license_test.go
@@ -378,6 +378,53 @@ func TestEntitlements(t *testing.T) {
 		}
 	})
 
+	t.Run("AllFeaturesAlwaysEnable", func(t *testing.T) {
+		t.Parallel()
+		db := dbfake.New()
+		db.InsertLicense(context.Background(), database.InsertLicenseParams{
+			Exp: dbtime.Now().Add(time.Hour),
+			JWT: coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
+				AllFeatures: true,
+			}),
+		})
+		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, empty)
+		require.NoError(t, err)
+		require.True(t, entitlements.HasLicense)
+		require.False(t, entitlements.Trial)
+		for _, featureName := range codersdk.FeatureNames {
+			if featureName == codersdk.FeatureUserLimit {
+				continue
+			}
+			feature := entitlements.Features[featureName]
+			require.Equal(t, featureName.AlwaysEnable(), feature.Enabled)
+			require.Equal(t, codersdk.EntitlementEntitled, feature.Entitlement)
+		}
+	})
+
+	t.Run("AllFeaturesGrace", func(t *testing.T) {
+		t.Parallel()
+		db := dbfake.New()
+		db.InsertLicense(context.Background(), database.InsertLicenseParams{
+			Exp: dbtime.Now().Add(time.Hour),
+			JWT: coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
+				AllFeatures: true,
+				GraceAt:     dbtime.Now().Add(-time.Hour),
+				ExpiresAt:   dbtime.Now().Add(time.Hour),
+			}),
+		})
+		entitlements, err := license.Entitlements(context.Background(), db, slog.Logger{}, 1, 1, coderdenttest.Keys, all)
+		require.NoError(t, err)
+		require.True(t, entitlements.HasLicense)
+		require.False(t, entitlements.Trial)
+		for _, featureName := range codersdk.FeatureNames {
+			if featureName == codersdk.FeatureUserLimit {
+				continue
+			}
+			require.True(t, entitlements.Features[featureName].Enabled)
+			require.Equal(t, codersdk.EntitlementGracePeriod, entitlements.Features[featureName].Entitlement)
+		}
+	})
+
 	t.Run("MultipleReplicasNoLicense", func(t *testing.T) {
 		t.Parallel()
 		db := dbfake.New()

Original file line number	Diff line number	Diff line change
`@@ -61,6 +61,7 @@ func Entitlements(`
`61`	`61`	`}`
`62`	`62`
`63`	`63`	`allFeatures := false`
	`64`	`+ allFeaturesEntitlement := codersdk.EntitlementNotEntitled`
`64`	`65`
`65`	`66`	`// Here we loop through licenses to detect enabled features.`
`66`	`67`	`for _, l := range licenses {`
`@@ -117,14 +118,15 @@ func Entitlements(`
`117`	`118`	`}`
`118`	`119`	`default:`
`119`	`120`	`entitlements.Features[featureName] = codersdk.Feature{`
`120`		`- Entitlement: entitlement,`
	`121`	`+ Entitlement: maxEntitlement(entitlements.Features[featureName].Entitlement, entitlement),`
`121`	`122`	`Enabled: enablements[featureName] \|\| featureName.AlwaysEnable(),`
`122`	`123`	`}`
`123`	`124`	`}`
`124`	`125`	`}`
`125`	`126`
`126`	`127`	`if claims.AllFeatures {`
`127`	`128`	`allFeatures = true`
	`129`	`+ allFeaturesEntitlement = maxEntitlement(allFeaturesEntitlement, entitlement)`
`128`	`130`	`}`
`129`	`131`	`entitlements.RequireTelemetry = entitlements.RequireTelemetry \|\| claims.RequireTelemetry`
`130`	`132`	`}`
`@@ -136,7 +138,8 @@ func Entitlements(`
`136`	`138`	`continue`
`137`	`139`	`}`
`138`	`140`	`feature := entitlements.Features[featureName]`
`139`		`- feature.Entitlement = codersdk.EntitlementEntitled`
	`141`	`+ feature.Entitlement = maxEntitlement(feature.Entitlement, allFeaturesEntitlement)`
	`142`	`+ feature.Enabled = enablements[featureName] \|\| featureName.AlwaysEnable()`
`140`	`143`	`entitlements.Features[featureName] = feature`
`141`	`144`	`}`
`142`	`145`	`}`
`@@ -324,3 +327,14 @@ func keyFunc(keys map[string]ed25519.PublicKey) func(*jwt.Token) (interface{}, e`
`324`	`327`	`return k, nil`
`325`	`328`	`}`
`326`	`329`	`}`
	`330`	`+`
	`331`	`+// maxEntitlement is the "greater" entitlement between the given values`
	`332`	`+func maxEntitlement(e1, e2 codersdk.Entitlement) codersdk.Entitlement {`
	`333`	`+ if e1 == codersdk.EntitlementEntitled \|\| e2 == codersdk.EntitlementEntitled {`
	`334`	`+ return codersdk.EntitlementEntitled`
	`335`	`+ }`
	`336`	`+ if e1 == codersdk.EntitlementGracePeriod \|\| e2 == codersdk.EntitlementGracePeriod {`
	`337`	`+ return codersdk.EntitlementGracePeriod`
	`338`	`+ }`
	`339`	`+ return codersdk.EntitlementNotEntitled`
	`340`	`+}`