diff --git a/pkg/ccl/serverccl/BUILD.bazel b/pkg/ccl/serverccl/BUILD.bazel index d85efcca6b72..640c1ddcfa66 100644 --- a/pkg/ccl/serverccl/BUILD.bazel +++ b/pkg/ccl/serverccl/BUILD.bazel @@ -23,6 +23,7 @@ go_test( "//pkg/ccl", "//pkg/ccl/kvccl", "//pkg/ccl/utilccl", + "//pkg/ccl/utilccl/licenseccl", "//pkg/roachpb:with-mocks", "//pkg/security", "//pkg/security/securitytest", @@ -35,6 +36,7 @@ go_test( "//pkg/testutils/sqlutils", "//pkg/testutils/testcluster", "//pkg/util", + "//pkg/util/envutil", "//pkg/util/httputil", "//pkg/util/leaktest", "//pkg/util/log", diff --git a/pkg/ccl/serverccl/server_sql_test.go b/pkg/ccl/serverccl/server_sql_test.go index d7ae0c21b414..3e47517381f2 100644 --- a/pkg/ccl/serverccl/server_sql_test.go +++ b/pkg/ccl/serverccl/server_sql_test.go @@ -16,12 +16,15 @@ import ( "testing" "github.com/cockroachdb/cockroach/pkg/base" + "github.com/cockroachdb/cockroach/pkg/ccl/utilccl" + "github.com/cockroachdb/cockroach/pkg/ccl/utilccl/licenseccl" "github.com/cockroachdb/cockroach/pkg/roachpb" "github.com/cockroachdb/cockroach/pkg/security" "github.com/cockroachdb/cockroach/pkg/sql" "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode" "github.com/cockroachdb/cockroach/pkg/testutils/serverutils" "github.com/cockroachdb/cockroach/pkg/testutils/sqlutils" + "github.com/cockroachdb/cockroach/pkg/util/envutil" "github.com/cockroachdb/cockroach/pkg/util/httputil" "github.com/cockroachdb/cockroach/pkg/util/leaktest" "github.com/cockroachdb/cockroach/pkg/util/log" @@ -84,6 +87,29 @@ func TestTenantCannotSetClusterSetting(t *testing.T) { require.Equal(t, pq.ErrorCode(pgcode.InsufficientPrivilege.String()), pqErr.Code, "err %v has unexpected code", err) } +func TestTenantCanUseEnterpriseFeatures(t *testing.T) { + defer leaktest.AfterTest(t)() + defer log.Scope(t).Close(t) + + license, _ := (&licenseccl.License{ + Type: licenseccl.License_Enterprise, + }).Encode() + + defer utilccl.TestingDisableEnterprise()() + defer envutil.TestSetEnv(t, "COCKROACH_TENANT_LICENSE", license)() + + tc := serverutils.StartNewTestCluster(t, 1, base.TestClusterArgs{}) + defer tc.Stopper().Stop(context.Background()) + + _, db := serverutils.StartTenant(t, tc.Server(0), base.TestTenantArgs{TenantID: serverutils.TestTenantID(), AllowSettingClusterSettings: false}) + defer db.Close() + + _, err := db.Exec(`BACKUP INTO 'userfile:///backup'`) + require.NoError(t, err) + _, err = db.Exec(`BACKUP INTO LATEST IN 'userfile:///backup'`) + require.NoError(t, err) +} + func TestTenantUnauthenticatedAccess(t *testing.T) { defer leaktest.AfterTest(t)() defer log.Scope(t).Close(t) diff --git a/pkg/ccl/utilccl/BUILD.bazel b/pkg/ccl/utilccl/BUILD.bazel index a026c38bc6d9..f505497080fe 100644 --- a/pkg/ccl/utilccl/BUILD.bazel +++ b/pkg/ccl/utilccl/BUILD.bazel @@ -13,6 +13,7 @@ go_library( "//pkg/base", "//pkg/ccl/utilccl/licenseccl", "//pkg/kv/kvclient/kvcoord:with-mocks", + "//pkg/server", "//pkg/settings", "//pkg/settings/cluster", "//pkg/sql/catalog/colinfo", @@ -20,6 +21,7 @@ go_library( "//pkg/sql/pgwire/pgcode", "//pkg/sql/pgwire/pgerror", "//pkg/sql/types", + "//pkg/util/envutil", "//pkg/util/grpcutil", "//pkg/util/timeutil", "//pkg/util/uuid", @@ -40,6 +42,7 @@ go_test( "//pkg/ccl/utilccl/licenseccl", "//pkg/settings/cluster", "//pkg/testutils", + "//pkg/util/envutil", "//pkg/util/timeutil", "//pkg/util/uuid", "@com_github_stretchr_testify//require", diff --git a/pkg/ccl/utilccl/license_check.go b/pkg/ccl/utilccl/license_check.go index 8d4e101367f8..3a9e698171b6 100644 --- a/pkg/ccl/utilccl/license_check.go +++ b/pkg/ccl/utilccl/license_check.go @@ -17,10 +17,12 @@ import ( "github.com/cockroachdb/cockroach/pkg/base" "github.com/cockroachdb/cockroach/pkg/ccl/utilccl/licenseccl" + "github.com/cockroachdb/cockroach/pkg/server" "github.com/cockroachdb/cockroach/pkg/settings" "github.com/cockroachdb/cockroach/pkg/settings/cluster" "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgcode" "github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgerror" + "github.com/cockroachdb/cockroach/pkg/util/envutil" "github.com/cockroachdb/cockroach/pkg/util/timeutil" "github.com/cockroachdb/cockroach/pkg/util/uuid" "github.com/cockroachdb/errors" @@ -45,16 +47,13 @@ var enterpriseLicense = func() *settings.StringSetting { return s }() -// testingEnterprise determines whether the cluster is enabled -// or disabled for the purposes of testing. -// It should be loaded and stored using atomic as it can race with an -// in progress kv reader during TestingDisableEnterprise / -// TestingEnableEnterprise. -var testingEnterprise int32 +// enterpriseStatus determines whether the cluster is enabled +// for enterprise features or if enterprise status depends on the license. +var enterpriseStatus int32 = deferToLicense const ( - testingEnterpriseDisabled = 0 - testingEnterpriseEnabled = 1 + deferToLicense = 0 + enterpriseEnabled = 1 ) // errEnterpriseRequired is returned by check() when the caller does @@ -68,22 +67,38 @@ type licenseCacheKey string // TestingEnableEnterprise allows overriding the license check in tests. func TestingEnableEnterprise() func() { - before := atomic.LoadInt32(&testingEnterprise) - atomic.StoreInt32(&testingEnterprise, testingEnterpriseEnabled) + before := atomic.LoadInt32(&enterpriseStatus) + atomic.StoreInt32(&enterpriseStatus, enterpriseEnabled) return func() { - atomic.StoreInt32(&testingEnterprise, before) + atomic.StoreInt32(&enterpriseStatus, before) } } // TestingDisableEnterprise allows re-enabling the license check in tests. func TestingDisableEnterprise() func() { - before := atomic.LoadInt32(&testingEnterprise) - atomic.StoreInt32(&testingEnterprise, testingEnterpriseDisabled) + before := atomic.LoadInt32(&enterpriseStatus) + atomic.StoreInt32(&enterpriseStatus, deferToLicense) return func() { - atomic.StoreInt32(&testingEnterprise, before) + atomic.StoreInt32(&enterpriseStatus, before) } } +// ApplyTenantLicense verifies the COCKROACH_TENANT_LICENSE environment variable +// and enables enterprise features for the process. This is a bit of a hack and +// should be replaced once it is possible to read the host cluster's +// enterprise.license setting. +func ApplyTenantLicense() error { + license, ok := envutil.EnvString("COCKROACH_TENANT_LICENSE", 0) + if !ok { + return nil + } + if _, err := decode(license); err != nil { + return errors.Wrap(err, "COCKROACH_TENANT_LICENSE encoding is invalid") + } + atomic.StoreInt32(&enterpriseStatus, enterpriseEnabled) + return nil +} + // CheckEnterpriseEnabled returns a non-nil error if the requested enterprise // feature is not enabled, including information or a link explaining how to // enable it. @@ -108,6 +123,7 @@ func init() { base.CheckEnterpriseEnabled = CheckEnterpriseEnabled base.LicenseType = getLicenseType base.TimeToEnterpriseLicenseExpiry = TimeToEnterpriseLicenseExpiry + server.ApplyTenantLicense = ApplyTenantLicense } // TimeToEnterpriseLicenseExpiry returns a Duration from `asOf` until the current @@ -128,7 +144,7 @@ func TimeToEnterpriseLicenseExpiry( func checkEnterpriseEnabledAt( st *cluster.Settings, at time.Time, cluster uuid.UUID, org, feature string, withDetails bool, ) error { - if atomic.LoadInt32(&testingEnterprise) == testingEnterpriseEnabled { + if atomic.LoadInt32(&enterpriseStatus) == enterpriseEnabled { return nil } license, err := getLicense(st) diff --git a/pkg/ccl/utilccl/license_check_test.go b/pkg/ccl/utilccl/license_check_test.go index ae85598f2f9b..6d91fb36e7c7 100644 --- a/pkg/ccl/utilccl/license_check_test.go +++ b/pkg/ccl/utilccl/license_check_test.go @@ -16,6 +16,7 @@ import ( "github.com/cockroachdb/cockroach/pkg/ccl/utilccl/licenseccl" "github.com/cockroachdb/cockroach/pkg/settings/cluster" "github.com/cockroachdb/cockroach/pkg/testutils" + "github.com/cockroachdb/cockroach/pkg/util/envutil" "github.com/cockroachdb/cockroach/pkg/util/timeutil" "github.com/cockroachdb/cockroach/pkg/util/uuid" "github.com/stretchr/testify/require" @@ -188,3 +189,40 @@ func TestTimeToEnterpriseLicenseExpiry(t *testing.T) { }) } } + +func TestApplyTenantLicenseWithLicense(t *testing.T) { + license, _ := (&licenseccl.License{ + Type: licenseccl.License_Enterprise, + }).Encode() + + defer TestingDisableEnterprise()() + defer envutil.TestSetEnv(t, "COCKROACH_TENANT_LICENSE", license)() + + settings := cluster.MakeClusterSettings() + + require.Error(t, CheckEnterpriseEnabled(settings, uuid.MakeV4(), "", "")) + require.False(t, IsEnterpriseEnabled(settings, uuid.MakeV4(), "", "")) + require.NoError(t, ApplyTenantLicense()) + require.NoError(t, CheckEnterpriseEnabled(settings, uuid.MakeV4(), "", "")) + require.True(t, IsEnterpriseEnabled(settings, uuid.MakeV4(), "", "")) +} + +func TestApplyTenantLicenseWithoutLicense(t *testing.T) { + defer TestingDisableEnterprise()() + + settings := cluster.MakeClusterSettings() + _, ok := envutil.EnvString("COCKROACH_TENANT_LICENSE", 0) + envutil.ClearEnvCache() + require.False(t, ok) + + require.Error(t, CheckEnterpriseEnabled(settings, uuid.MakeV4(), "", "")) + require.False(t, IsEnterpriseEnabled(settings, uuid.MakeV4(), "", "")) + require.NoError(t, ApplyTenantLicense()) + require.Error(t, CheckEnterpriseEnabled(settings, uuid.MakeV4(), "", "")) + require.False(t, IsEnterpriseEnabled(settings, uuid.MakeV4(), "", "")) +} + +func TestApplyTenantLicenseWithInvalidLicense(t *testing.T) { + defer envutil.TestSetEnv(t, "COCKROACH_TENANT_LICENSE", "THIS IS NOT A VALID LICENSE")() + require.Error(t, ApplyTenantLicense()) +} diff --git a/pkg/server/tenant.go b/pkg/server/tenant.go index 4fa76f510067..cba39c7a229d 100644 --- a/pkg/server/tenant.go +++ b/pkg/server/tenant.go @@ -57,6 +57,11 @@ func StartTenant( baseCfg BaseConfig, sqlCfg SQLConfig, ) (sqlServer *SQLServer, pgAddr string, httpAddr string, _ error) { + err := ApplyTenantLicense() + if err != nil { + return nil, "", "", err + } + args, err := makeTenantSQLServerArgs(stopper, kvClusterName, baseCfg, sqlCfg) if err != nil { return nil, "", "", err @@ -405,6 +410,11 @@ var NewTenantSideCostController = func( return noopTenantSideCostController{}, nil } +// ApplyTenantLicense is a hook for CCL code which enables enterprise features +// for the tenant process if the COCKROACH_TENANT_LICENSE environment variable +// is set. +var ApplyTenantLicense = func() error { return nil /* no-op */ } + // noopTenantSideCostController is a no-op implementation of // TenantSideCostController. type noopTenantSideCostController struct{} diff --git a/pkg/util/envutil/BUILD.bazel b/pkg/util/envutil/BUILD.bazel index ac03c340b785..8e34a3d2f413 100644 --- a/pkg/util/envutil/BUILD.bazel +++ b/pkg/util/envutil/BUILD.bazel @@ -17,4 +17,5 @@ go_test( size = "small", srcs = ["env_test.go"], embed = [":envutil"], + deps = ["@com_github_stretchr_testify//require"], ) diff --git a/pkg/util/envutil/env.go b/pkg/util/envutil/env.go index 1e78227a8971..87e8c5fa622e 100644 --- a/pkg/util/envutil/env.go +++ b/pkg/util/envutil/env.go @@ -19,6 +19,7 @@ import ( "sort" "strconv" "strings" + "testing" "time" "github.com/cockroachdb/cockroach/pkg/util/humanizeutil" @@ -361,3 +362,26 @@ func EnvOrDefaultDuration(name string, value time.Duration) time.Duration { } return value } + +// TestSetEnv sets an environment variable and the cleanup function +// resets it to the original value. +func TestSetEnv(t *testing.T, name string, value string) func() { + ClearEnvCache() + before, exists := os.LookupEnv(name) + + if err := os.Setenv(name, value); err != nil { + t.Fatal(err) + } + return func() { + if exists { + if err := os.Setenv(name, before); err != nil { + t.Fatal(err) + } + } else { + if err := os.Unsetenv(name); err != nil { + t.Fatal(err) + } + } + ClearEnvCache() + } +} diff --git a/pkg/util/envutil/env_test.go b/pkg/util/envutil/env_test.go index 9e71424f9c4a..264425cae322 100644 --- a/pkg/util/envutil/env_test.go +++ b/pkg/util/envutil/env_test.go @@ -13,6 +13,8 @@ package envutil import ( "os" "testing" + + "github.com/stretchr/testify/require" ) func TestEnvOrDefault(t *testing.T) { @@ -27,3 +29,43 @@ func TestEnvOrDefault(t *testing.T) { t.Errorf("expected %d, got %d", def, act) } } + +func TestTestSetEnvExists(t *testing.T) { + key := "COCKROACH_ENVUTIL_TESTSETTING" + require.NoError(t, os.Setenv(key, "before")) + + ClearEnvCache() + value, ok := EnvString(key, 0) + require.True(t, ok) + require.Equal(t, value, "before") + + cleanup := TestSetEnv(t, key, "testvalue") + value, ok = EnvString(key, 0) + require.True(t, ok) + require.Equal(t, value, "testvalue") + + cleanup() + + value, ok = EnvString(key, 0) + require.True(t, ok) + require.Equal(t, value, "before") +} + +func TestTestSetEnvDoesNotExist(t *testing.T) { + key := "COCKROACH_ENVUTIL_TESTSETTING" + require.NoError(t, os.Unsetenv(key)) + + ClearEnvCache() + _, ok := EnvString(key, 0) + require.False(t, ok) + + cleanup := TestSetEnv(t, key, "foo") + value, ok := EnvString(key, 0) + require.True(t, ok) + require.Equal(t, value, "foo") + + cleanup() + + _, ok = EnvString(key, 0) + require.False(t, ok) +}