feat/enterpriseportal: implement SubscriptionLicenseChecksService (#64400)

Implements the RPC defined in
https://github.com/sourcegraph/sourcegraph/pull/64396. Follow-up PRs
will implement migrations for in-instance checks and for through-dotcom
checks.

One major change is that we now bypass the check for subscriptions that
are denoted as associated with `INTERNAL` instances.

Most of the diff is generated mocks.

Part of https://linear.app/sourcegraph/issue/CORE-227

## Test plan

- [x] Unit tests
- [x] E2E tests (`sg test enterprise-portal-e2e`)


![image](https://github.com/user-attachments/assets/56fde7dd-95a0-4d98-bb4c-943b1f155e33)

Author: bobheadxi

client/web/src/enterprise/site-admin/dotcom/productSubscriptions/enterpriseportalgen/subscriptions_pb.ts

@@ -582,12 +582,20 @@ export enum EnterpriseSubscriptionLicenseCondition_Status {
    * @generated from enum value: STATUS_REVOKED = 2;
    */
   REVOKED = 2,
+
+  /**
+   * License usage from a Sourcegraph instance was detected.
+   *
+   * @generated from enum value: STATUS_INSTANCE_USAGE_DETECTED = 3;
+   */
+  INSTANCE_USAGE_DETECTED = 3,
 }
 // Retrieve enum metadata with: proto3.getEnumType(EnterpriseSubscriptionLicenseCondition_Status)
 proto3.util.setEnumType(EnterpriseSubscriptionLicenseCondition_Status, "enterpriseportal.subscriptions.v1.EnterpriseSubscriptionLicenseCondition.Status", [
   { no: 0, name: "STATUS_UNSPECIFIED" },
   { no: 1, name: "STATUS_CREATED" },
   { no: 2, name: "STATUS_REVOKED" },
+  { no: 3, name: "STATUS_INSTANCE_USAGE_DETECTED" },
 ]);
 
 /**

cmd/enterprise-portal/e2e/BUILD.bazel

@@ -6,8 +6,11 @@ go_test(
     deps = [
         "//internal/grpc/defaults",
         "//internal/grpc/grpcoauth",
+        "//internal/license",
         "//lib/enterpriseportal/codyaccess/v1:codyaccess",
+        "//lib/enterpriseportal/subscriptionlicensechecks/v1:subscriptionlicensechecks",
         "//lib/enterpriseportal/subscriptions/v1:subscriptions",
+        "@com_github_hexops_autogold_v2//:autogold",
         "@com_github_sourcegraph_log//logtest",
         "@com_github_sourcegraph_sourcegraph_accounts_sdk_go//:sourcegraph-accounts-sdk-go",
         "@com_github_sourcegraph_sourcegraph_accounts_sdk_go//scopes",

cmd/enterprise-portal/e2e/e2e_test.go

@@ -14,6 +14,7 @@ import (
 	"google.golang.org/protobuf/types/known/fieldmaskpb"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
+	"github.com/hexops/autogold/v2"
 	"github.com/sourcegraph/log/logtest"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -22,14 +23,17 @@ import (
 	"github.com/sourcegraph/sourcegraph-accounts-sdk-go/scopes"
 	"github.com/sourcegraph/sourcegraph/internal/grpc/defaults"
 	"github.com/sourcegraph/sourcegraph/internal/grpc/grpcoauth"
+	"github.com/sourcegraph/sourcegraph/internal/license"
 
 	codyaccessv1 "github.com/sourcegraph/sourcegraph/lib/enterpriseportal/codyaccess/v1"
+	subscriptionlicensechecks "github.com/sourcegraph/sourcegraph/lib/enterpriseportal/subscriptionlicensechecks/v1"
 	subscriptionsv1 "github.com/sourcegraph/sourcegraph/lib/enterpriseportal/subscriptions/v1"
 )
 
 type Clients struct {
 	Subscriptions subscriptionsv1.SubscriptionsServiceClient
 	CodyAccess    codyaccessv1.CodyAccessServiceClient
+	LicenseChecks subscriptionlicensechecks.SubscriptionLicenseChecksServiceClient
 }
 
 func newE2EClients(t *testing.T) *Clients {
@@ -80,9 +84,15 @@ func newE2EClients(t *testing.T) *Clients {
 	require.NoError(t, err)
 	t.Cleanup(func() { _ = client.Close() })
 
+	clientWithoutCreds, err := grpc.NewClient("dns:///"+addr.Host,
+		defaults.DialOptions(logtest.Scoped(t).Scoped("grpc"))...)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = clientWithoutCreds.Close() })
+
 	return &Clients{
 		Subscriptions: subscriptionsv1.NewSubscriptionsServiceClient(client),
 		CodyAccess:    codyaccessv1.NewCodyAccessServiceClient(client),
+		LicenseChecks: subscriptionlicensechecks.NewSubscriptionLicenseChecksServiceClient(clientWithoutCreds),
 	}
 }
 
@@ -160,6 +170,7 @@ func runLifecycleTest(t *testing.T, ctx context.Context, clients *Clients, runID
 	})
 
 	var createdLicenseID string
+	var createdLicenseKey string
 	t.Run("Create license", func(t *testing.T) {
 		got, err := clients.Subscriptions.CreateEnterpriseSubscriptionLicense(ctx, &subscriptionsv1.CreateEnterpriseSubscriptionLicenseRequest{
 			License: &subscriptionsv1.EnterpriseSubscriptionLicense{
@@ -178,6 +189,7 @@ func runLifecycleTest(t *testing.T, ctx context.Context, clients *Clients, runID
 		})
 		require.NoError(t, err)
 		createdLicenseID = got.GetLicense().GetId()
+		createdLicenseKey = got.GetLicense().GetKey().GetLicenseKey()
 		prettyPrint(t, got)
 	})
 
@@ -224,6 +236,34 @@ func runLifecycleTest(t *testing.T, ctx context.Context, clients *Clients, runID
 		prettyPrint(t, got)
 	})
 
+	t.Run("Check license", func(t *testing.T) {
+		got, err := clients.LicenseChecks.CheckLicenseKey(ctx, &subscriptionlicensechecks.CheckLicenseKeyRequest{
+			InstanceId: "test-instance-id",
+			LicenseKey: createdLicenseKey,
+		})
+		require.NoError(t, err)
+		assert.True(t, got.GetValid())
+
+		t.Run("back-compat with license key token", func(t *testing.T) {
+			got, err := clients.LicenseChecks.CheckLicenseKey(ctx, &subscriptionlicensechecks.CheckLicenseKeyRequest{
+				InstanceId: "test-instance-id",
+				LicenseKey: license.GenerateLicenseKeyBasedAccessToken(createdLicenseKey),
+			})
+			require.NoError(t, err)
+			assert.True(t, got.GetValid())
+		})
+
+		t.Run("with wrong site ID", func(t *testing.T) {
+			got, err := clients.LicenseChecks.CheckLicenseKey(ctx, &subscriptionlicensechecks.CheckLicenseKeyRequest{
+				InstanceId: "wrong-instance-id",
+				LicenseKey: createdLicenseKey,
+			})
+			require.NoError(t, err)
+			assert.False(t, got.GetValid())
+			autogold.Expect("license has already been used by another instance").Equal(t, got.GetReason())
+		})
+	})
+
 	t.Run("Revoke license", func(t *testing.T) {
 		got, err := clients.Subscriptions.RevokeEnterpriseSubscriptionLicense(ctx, &subscriptionsv1.RevokeEnterpriseSubscriptionLicenseRequest{
 			LicenseId: createdLicenseID,
@@ -242,6 +282,16 @@ func runLifecycleTest(t *testing.T, ctx context.Context, clients *Clients, runID
 		})
 	})
 
+	t.Run("Check revoked license", func(t *testing.T) {
+		got, err := clients.LicenseChecks.CheckLicenseKey(ctx, &subscriptionlicensechecks.CheckLicenseKeyRequest{
+			InstanceId: "test-instance-id",
+			LicenseKey: createdLicenseKey,
+		})
+		require.NoError(t, err)
+		assert.False(t, got.GetValid())
+		autogold.Expect("license has been revoked").Equal(t, got.GetReason())
+	})
+
 	t.Run("Archive subscription", func(t *testing.T) {
 		got, err := clients.Subscriptions.ArchiveEnterpriseSubscription(ctx, &subscriptionsv1.ArchiveEnterpriseSubscriptionRequest{
 			SubscriptionId: createdSubscriptionID,

cmd/enterprise-portal/internal/database/subscriptions/licenses.go

@@ -77,6 +77,11 @@ type SubscriptionLicense struct {
 	// Value shapes correspond to API types appropriate for each
 	// 'EnterpriseSubscriptionLicenseType'.
 	LicenseData json.RawMessage `gorm:"type:jsonb"`
+
+	// DetectedInstanceID is the identifier of the Sourcegraph instance that has
+	// been automatically detected via onlince license checks (subscriptionlicensechecks).
+	// It should only be used internally for reference.
+	DetectedInstanceID *string
 }
 
 // subscriptionLicenseWithConditionsColumns must match scanSubscriptionLicense()
@@ -93,6 +98,8 @@ func subscriptionLicenseWithConditionsColumns() []string {
 		"license_type",
 		"license_data",
 
+		"detected_instance_id",
+
 		subscriptionLicenseConditionJSONBAgg(),
 	}
 }
@@ -113,6 +120,7 @@ func scanSubscriptionLicenseWithConditions(row pgx.Row) (*LicenseWithConditions,
 		&l.ExpireAt,
 		&l.LicenseType,
 		&l.LicenseData,
+		&l.DetectedInstanceID,
 		&l.Conditions, // see subscriptionLicenseConditionJSONBAgg docstring
 	)
 	return &l, err
@@ -132,10 +140,19 @@ func NewLicensesStore(db *pgxpool.Pool) *LicensesStore {
 }
 
 type ListLicensesOpts struct {
-	SubscriptionID          string
-	LicenseType             subscriptionsv1.EnterpriseSubscriptionLicenseType
-	LicenseKeySubstring     string
+	SubscriptionID string
+	LicenseType    subscriptionsv1.EnterpriseSubscriptionLicenseType
+	// LicenseKey is an exact match on the signed key.
+	LicenseKey string
+	// LicenseKeySubstring is a substring match on the signed key.
+	LicenseKeySubstring string
+
 	SalesforceOpportunityID string
+
+	// LicenseKeyHash should be removed once subscriptionlicensechecks no longer
+	// supports the old key hash format
+	LicenseKeyHash []byte
+
 	// PageSize is the maximum number of licenses to return.
 	PageSize int
 }
@@ -155,6 +172,11 @@ func (opts ListLicensesOpts) toQueryConditions() (where, limitClause string, _ p
 
 	switch opts.LicenseType {
 	case subscriptionsv1.EnterpriseSubscriptionLicenseType_ENTERPRISE_SUBSCRIPTION_LICENSE_TYPE_KEY:
+		if opts.LicenseKey != "" {
+			whereConds = append(whereConds,
+				"license_data->>'SignedKey' = @licenseKey")
+			namedArgs["licenseKey"] = opts.LicenseKey
+		}
 		if opts.LicenseKeySubstring != "" {
 			whereConds = append(whereConds,
 				"license_data->>'SignedKey' LIKE  '%' || @licenseKeySubstring || '%'")
@@ -165,6 +187,11 @@ func (opts ListLicensesOpts) toQueryConditions() (where, limitClause string, _ p
 				"license_data->'Info'->>'sf_opp_id' = @salesforceOpportunityID")
 			namedArgs["salesforceOpportunityID"] = opts.SalesforceOpportunityID
 		}
+		if opts.LicenseKeyHash != nil {
+			whereConds = append(whereConds,
+				"DIGEST(license_data->>'SignedKey','sha256') = @licenseKeyHash")
+			namedArgs["licenseKeyHash"] = opts.LicenseKeyHash
+		}
 	}
 
 	where = strings.Join(whereConds, " AND ")
@@ -409,7 +436,7 @@ func (s *LicensesStore) Revoke(ctx context.Context, licenseID string, opts Revok
 		return nil, errors.Wrap(err, "begin transaction")
 	}
 	defer func() {
-		if rollbackErr := tx.Rollback(context.Background()); rollbackErr != nil {
+		if rollbackErr := tx.Rollback(context.WithoutCancel(ctx)); rollbackErr != nil {
 			err = errors.Append(err, rollbackErr)
 		}
 	}()
@@ -442,3 +469,59 @@ WHERE id = @licenseID
 
 	return s.Get(ctx, licenseID)
 }
+
+type SetDetectedInstanceOpts struct {
+	// InstanceID is the ID of the instance that was detected to be using this
+	// license.
+	InstanceID string
+	// Message to associate with the detection event.
+	Message string
+	// If nil, the detection time will be set to the current time.
+	Time *utctime.Time
+}
+
+// SetDetectedInstance sets the instance ID that was detected to be using this
+// license.
+func (s *LicensesStore) SetDetectedInstance(ctx context.Context, licenseID string, opts SetDetectedInstanceOpts) error {
+	if opts.Time == nil {
+		opts.Time = pointers.Ptr(utctime.Now())
+	}
+
+	tx, err := s.db.Begin(ctx)
+	if err != nil {
+		return errors.Wrap(err, "begin transaction")
+	}
+	defer func() {
+		if rollbackErr := tx.Rollback(context.WithoutCancel(ctx)); rollbackErr != nil {
+			err = errors.Append(err, rollbackErr)
+		}
+	}()
+
+	if _, err := tx.Exec(ctx, `
+UPDATE enterprise_portal_subscription_licenses
+SET detected_instance_id = @instanceID
+WHERE id = @licenseID
+`, pgx.NamedArgs{
+		"instanceID": opts.InstanceID,
+		"licenseID":  licenseID,
+	}); err != nil {
+		if errors.Is(err, pgx.ErrNoRows) {
+			return ErrSubscriptionLicenseNotFound
+		}
+		return errors.Wrap(err, "update detected instance for license")
+	}
+
+	if err := newLicenseConditionsStore(tx).createLicenseCondition(ctx, licenseID, createLicenseConditionOpts{
+		Status:         subscriptionsv1.EnterpriseSubscriptionLicenseCondition_STATUS_INSTANCE_USAGE_DETECTED,
+		Message:        opts.Message,
+		TransitionTime: *opts.Time,
+	}); err != nil {
+		return errors.Wrap(err, "create license condition")
+	}
+
+	if err := tx.Commit(ctx); err != nil {
+		return errors.Wrap(err, "commit transaction")
+	}
+
+	return nil
+}

cmd/enterprise-portal/internal/database/subscriptions/licenses_test.go

@@ -242,6 +242,20 @@ func TestLicensesStore(t *testing.T) {
 			})
 		})
 
+		t.Run("list by license key hash token", func(t *testing.T) {
+			hash, err := license.ExtractLicenseKeyBasedAccessTokenContents(
+				license.GenerateLicenseKeyBasedAccessToken(signedKeyExample),
+			)
+			require.NoError(t, err)
+			listedLicenses, err := licenses.List(ctx, subscriptions.ListLicensesOpts{
+				LicenseType:    subscriptionsv1.EnterpriseSubscriptionLicenseType_ENTERPRISE_SUBSCRIPTION_LICENSE_TYPE_KEY,
+				LicenseKeyHash: []byte(hash),
+			})
+			require.NoError(t, err)
+			require.Len(t, listedLicenses, 1)
+			assert.Equal(t, subscriptionID1, listedLicenses[0].SubscriptionID)
+		})
+
 		t.Run("List by salesforce opportunity ID", func(t *testing.T) {
 			listedLicenses, err := licenses.List(ctx, subscriptions.ListLicensesOpts{
 				LicenseType:             subscriptionsv1.EnterpriseSubscriptionLicenseType_ENTERPRISE_SUBSCRIPTION_LICENSE_TYPE_KEY,
@@ -270,20 +284,41 @@ func TestLicensesStore(t *testing.T) {
 		}
 	})
 
+	t.Run("SetDetectedInstance", func(t *testing.T) {
+		require.NoError(t, licenses.SetDetectedInstance(ctx, createdLicenses[0].ID, subscriptions.SetDetectedInstanceOpts{
+			InstanceID: "instance-id",
+			Message:    t.Name(),
+		}))
+		got, err := licenses.Get(ctx, createdLicenses[0].ID)
+		require.NoError(t, err)
+		assert.Equal(t, "instance-id", *got.DetectedInstanceID)
+	})
+
 	t.Run("Revoke", func(t *testing.T) {
 		for idx, license := range createdLicenses {
-			revokeTime := utctime.FromTime(time.Now().Add(-time.Second))
+			revokeTime := utctime.FromTime(time.Now())
 			got, err := licenses.Revoke(ctx, license.ID, subscriptions.RevokeLicenseOpts{
 				Message: fmt.Sprintf("%s %d", t.Name(), idx),
 				Time:    pointers.Ptr(revokeTime),
 			})
 			require.NoError(t, err)
 			assert.Equal(t, revokeTime.AsTime(), got.RevokedAt.AsTime())
-			require.Len(t, got.Conditions, 2)
-			// Most recent condition is sorted first, and should be the revocation
-			assert.Equal(t, "STATUS_REVOKED", got.Conditions[0].Status)
-			assert.Equal(t, revokeTime.AsTime(), got.Conditions[0].TransitionTime.AsTime())
-			assert.Equal(t, "STATUS_CREATED", got.Conditions[1].Status)
+			if idx > 0 {
+				require.Len(t, got.Conditions, 2)
+				// Most recent condition is sorted first, and should be the revocation
+				assert.Equal(t, "STATUS_REVOKED", got.Conditions[0].Status)
+				assert.Equal(t, revokeTime.AsTime(), got.Conditions[0].TransitionTime.AsTime())
+				assert.Equal(t, "STATUS_CREATED", got.Conditions[1].Status)
+			} else {
+				require.Len(t, got.Conditions, 3)
+				// Most recent condition is sorted first, and should be the revocation
+				assert.Equal(t, "STATUS_REVOKED", got.Conditions[0].Status)
+				assert.Equal(t, revokeTime.AsTime(), got.Conditions[0].TransitionTime.AsTime())
+				// Then, the condition from SetDetectedInstance test
+				assert.Equal(t, "STATUS_INSTANCE_USAGE_DETECTED", got.Conditions[1].Status)
+				// Finally, the subscription creation event
+				assert.Equal(t, "STATUS_CREATED", got.Conditions[2].Status)
+			}
 		}
 	})
 }