zip: Add tenant ID flag for debug zip

pkg/acceptance/cluster/certs.go

@@ -53,12 +53,12 @@ func GenerateCerts(ctx context.Context) func() {
 	// Root user.
 	maybePanic(security.CreateClientPair(
 		certsDir, filepath.Join(certsDir, security.EmbeddedCAKey),
-		2048, 48*time.Hour, false, security.RootUserName(), true /* generate pk8 key */))
+		2048, 48*time.Hour, false, security.RootUserName(), "" /* tenantID */, true /* generate pk8 key */))
 
 	// Test user.
 	maybePanic(security.CreateClientPair(
 		certsDir, filepath.Join(certsDir, security.EmbeddedCAKey),
-		1024, 48*time.Hour, false, security.TestUserName(), true /* generate pk8 key */))
+		1024, 48*time.Hour, false, security.TestUserName(), "" /* tenantID */, true /* generate pk8 key */))
 
 	// Certs for starting a cockroach server. Key size is from cli/cert.go:defaultKeySize.
 	maybePanic(security.CreateNodePair(

pkg/cli/cert.go

@@ -177,6 +177,7 @@ func runCreateClientCert(cmd *cobra.Command, args []string) error {
 			certCtx.certificateLifetime,
 			certCtx.overwriteFiles,
 			username,
+			certCtx.tenantScope,
 			certCtx.generatePKCS8Key),
 		"failed to generate client certificate and key")
 }

pkg/cli/client_url.go

@@ -16,7 +16,6 @@ import (
 	"path/filepath"
 
 	"github.com/cockroachdb/cockroach/pkg/cli/cliflags"
-	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/rpc"
 	"github.com/cockroachdb/cockroach/pkg/security"
 	"github.com/cockroachdb/cockroach/pkg/server/pgurl"
@@ -360,7 +359,7 @@ func (cliCtx *cliContext) makeClientConnURL() (*pgurl.URL, error) {
 		userName = security.RootUserName()
 	}
 
-	sCtx := rpc.MakeSecurityContext(cliCtx.Config, security.CommandTLSSettings{}, roachpb.SystemTenantID)
+	sCtx := rpc.MakeSecurityContext(cliCtx.Config, security.CommandTLSSettings{}, cliCtx.tenantID)
 	if err := sCtx.LoadSecurityOptions(purl, userName); err != nil {
 		return nil, err
 	}

pkg/cli/cliflags/flags.go

@@ -734,6 +734,14 @@ Note: that --external-io-disable-http or --external-io-disable-implicit-credenti
 		Description: `Certificate and key files are overwritten if they exist.`,
 	}
 
+	TenantScope = FlagInfo{
+		Name: "tenant-scope",
+		Description: `Assign a tenant scope to the certificate.
+This will allow for the certificate to only be used specifically for a particular
+tenant. This flag is optional, when omitted, the certificate is not tied
+for usage on a specific tenant.`,
+	}
+
 	GeneratePKCS8Key = FlagInfo{
 		Name:        "also-generate-pkcs8-key",
 		Description: `Also write the key in pkcs8 format to <certs-dir>/client.<username>.key.pk8.`,
@@ -1465,6 +1473,14 @@ Can be set to 1 to ensure only one node is polled for data at a time.
 `,
 	}
 
+	ZipTenant = FlagInfo{
+		Name: "tenant-id",
+		Description: `
+Specify the tenant ID of the server. This is required to be set while
+running debug zip against a SQL only server for a tenant.
+`,
+	}
+
 	StmtDiagDeleteAll = FlagInfo{
 		Name:        "all",
 		Description: `Delete all bundles.`,

pkg/cli/context.go

@@ -26,6 +26,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/cli/clisqlshell"
 	"github.com/cockroachdb/cockroach/pkg/cli/democluster"
 	"github.com/cockroachdb/cockroach/pkg/config/zonepb"
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/security"
 	"github.com/cockroachdb/cockroach/pkg/server"
 	"github.com/cockroachdb/cockroach/pkg/server/pgurl"
@@ -200,6 +201,10 @@ type cliContext struct {
 
 	// For `cockroach version --build-tag`.
 	showVersionUsingOnlyBuildTag bool
+
+	// tenantID indicates the tenant to run the CLI utility against.
+	// Default value is the system tenant.
+	tenantID roachpb.TenantID
 }
 
 // cliCtx captures the command-line parameters common to most CLI utilities.
@@ -233,6 +238,7 @@ func setCliContextDefaults() {
 	// TODO(knz): Deprecated in v21.1. Remove this.
 	cliCtx.deprecatedLogOverrides.reset()
 	cliCtx.showVersionUsingOnlyBuildTag = false
+	cliCtx.tenantID = roachpb.SystemTenantID
 }
 
 // sqlConnContext captures the connection configuration for all SQL
@@ -266,6 +272,10 @@ var certCtx struct {
 	// This configuration flag is only used for 'cert' commands
 	// that generate certificates.
 	certPrincipalMap []string
+	// tenantScope indicates the ID of the tenant that a certificate is being
+	// scoped to. By creating a tenant-scoped certicate, the usage of that certificate
+	// is restricted to a specific tenant.
+	tenantScope string
 }
 
 func setCertContextDefaults() {
@@ -278,6 +288,7 @@ func setCertContextDefaults() {
 	certCtx.overwriteFiles = false
 	certCtx.generatePKCS8Key = false
 	certCtx.certPrincipalMap = nil
+	certCtx.tenantScope = ""
 }
 
 var sqlExecCtx = clisqlexec.Context{
@@ -345,6 +356,10 @@ type zipContext struct {
 
 	// The log/heap/etc files to include.
 	files fileSelection
+
+	// tenantID of the server being connected to. This flag should
+	// be set while running debug zip against a tenant SQL server.
+	tenantID roachpb.TenantID
 }
 
 // setZipContextDefaults set the default values in zipCtx.  This
@@ -364,6 +379,7 @@ func setZipContextDefaults() {
 	now := timeutil.Now()
 	zipCtx.files.startTimestamp = timestampValue(now.Add(-48 * time.Hour))
 	zipCtx.files.endTimestamp = timestampValue(now.Add(24 * time.Hour))
+	zipCtx.tenantID = roachpb.SystemTenantID
 }
 
 // dumpCtx captures the command-line parameters of the `dump` command.

pkg/cli/democluster/demo_cluster.go

@@ -1073,6 +1073,7 @@ func (demoCtx *Context) generateCerts(certsDir string) (err error) {
 		demoCtx.DefaultCertLifetime,
 		false, /* overwrite */
 		security.RootUserName(),
+		"",    /* tenantID */
 		false, /* generatePKCS8Key */
 	); err != nil {
 		return err

pkg/cli/flags.go

@@ -271,6 +271,33 @@ func (f *keyTypeFilter) Set(v string) error {
 	return nil
 }
 
+// tenantIDSetter wraps the tenantID variable within zipContext
+// and verifies its value during execution.
+type tenantIDSetter struct {
+	tenantID *roachpb.TenantID
+}
+
+// String implements the pflag.Value interface.
+func (t tenantIDSetter) String() string { return t.tenantID.String() }
+
+// Type implements the pflag.Value interface.
+func (t tenantIDSetter) Type() string { return "<uint>" }
+
+// Set implements the pflag.Value interface.
+func (t tenantIDSetter) Set(v string) error {
+	if v == "" {
+		*t.tenantID = roachpb.SystemTenantID
+		return nil
+	}
+
+	tID, err := roachpb.ParseTenantID(v)
+	if err != nil {
+		return err
+	}
+	*t.tenantID = tID
+	return nil
+}
+
 const backgroundEnvVar = "COCKROACH_BACKGROUND_RESTART"
 
 // flagSetForCmd is a replacement for cmd.Flag() that properly merges
@@ -580,6 +607,7 @@ func init() {
 		stringFlag(f, &certCtx.caKey, cliflags.CAKey)
 		intFlag(f, &certCtx.keySize, cliflags.KeySize)
 		boolFlag(f, &certCtx.overwriteFiles, cliflags.OverwriteFiles)
+		stringFlag(f, &certCtx.tenantScope, cliflags.TenantScope)
 
 		if strings.HasSuffix(cmd.Name(), "-ca") {
 			// CA-only commands.
@@ -684,6 +712,7 @@ func init() {
 		boolFlag(f, &zipCtx.redactLogs, cliflags.ZipRedactLogs)
 		durationFlag(f, &zipCtx.cpuProfDuration, cliflags.ZipCPUProfileDuration)
 		intFlag(f, &zipCtx.concurrency, cliflags.ZipConcurrency)
+		varFlag(f, tenantIDSetter{&zipCtx.tenantID}, cliflags.ZipTenant)
 	}
 	// List-files + Zip commands.
 	for _, cmd := range []*cobra.Command{debugZipCmd, debugListFilesCmd} {

pkg/cli/start.go

@@ -1263,7 +1263,7 @@ func getClientGRPCConn(
 	stopper := stop.NewStopper(stop.WithTracer(tracer))
 	rpcContext := rpc.NewContext(ctx,
 		rpc.ContextOptions{
-			TenantID: roachpb.SystemTenantID,
+			TenantID: cfg.TenantID,
 			Config:   cfg.Config,
 			Clock:    clock,
 			Stopper:  stopper,

pkg/cli/testdata/zip/testzip_tenant

@@ -1,6 +1,6 @@
 zip
 ----
-debug zip --concurrency=1 --cpu-profile-duration=1s /dev/null
+debug zip --concurrency=1 --cpu-profile-duration=1s --tenant-id=10 /dev/null
 [cluster] establishing RPC connection to ...
 [cluster] retrieving the node status to get the SQL address... done
 [cluster] using SQL address: ...

pkg/cli/testutils.go

@@ -328,18 +328,30 @@ func isSQLCommand(args []string) (bool, error) {
 	return false, nil
 }
 
-func (c TestCLI) getRPCAddr() string {
-	if c.tenant != nil {
-		return c.tenant.RPCAddr()
+func (c TestCLI) getRPCAddr(tenantID roachpb.TenantID) (string, error) {
+	if tenantID == roachpb.SystemTenantID {
+		return c.ServingRPCAddr(), nil
 	}
-	return c.ServingRPCAddr()
+	if c.tenant == nil {
+		return "", errors.Errorf("cannot run CLI for tenant %d on system tenant", tenantID)
+	}
+	if c.tenant.RPCContext().TenantID != tenantID {
+		return "", errors.Errorf("cannot run CLI for tenant %d on tenant %d", tenantID, c.tenant.RPCContext().TenantID)
+	}
+	return c.tenant.RPCAddr(), nil
 }
 
-func (c TestCLI) getSQLAddr() string {
-	if c.tenant != nil {
-		return c.tenant.SQLAddr()
+func (c TestCLI) getSQLAddr(tenantID roachpb.TenantID) (string, error) {
+	if tenantID == roachpb.SystemTenantID {
+		return c.ServingSQLAddr(), nil
+	}
+	if c.tenant == nil {
+		return "", errors.Errorf("cannot run CLI for tenant %d on system tenant", tenantID)
 	}
-	return c.ServingSQLAddr()
+	if c.tenant.RPCContext().TenantID != tenantID {
+		return "", errors.Errorf("cannot run CLI for tenant %d on tenant %d", tenantID, c.tenant.RPCContext().TenantID)
+	}
+	return c.tenant.SQLAddr(), nil
 }
 
 // RunWithArgs add args according to TestCLI cfg.
@@ -349,11 +361,21 @@ func (c TestCLI) RunWithArgs(origArgs []string) {
 	if err := func() error {
 		args := append([]string(nil), origArgs[:1]...)
 		if c.TestServer != nil {
-			addr := c.getRPCAddr()
+			tenantID, err := getTenantID(origArgs)
+			if err != nil {
+				return err
+			}
+			addr, err := c.getRPCAddr(tenantID)
+			if err != nil {
+				return err
+			}
 			if isSQL, err := isSQLCommand(origArgs); err != nil {
 				return err
 			} else if isSQL {
-				addr = c.getSQLAddr()
+				addr, err = c.getSQLAddr(tenantID)
+				if err != nil {
+					return err
+				}
 			}
 			h, p, err := net.SplitHostPort(addr)
 			if err != nil {
@@ -526,3 +548,20 @@ func MatchCSV(csvStr string, matchColRow [][]string) (err error) {
 	}
 	return err
 }
+
+func getTenantID(args []string) (roachpb.TenantID, error) {
+	for _, arg := range args {
+		if strings.HasPrefix(arg, "--tenant-id") {
+			parts := strings.Split(arg, "=")
+			if len(parts) != 2 {
+				return roachpb.TenantID{}, errors.Errorf("invalid tenant-id argument %s", arg)
+			}
+			tenantID, err := roachpb.ParseTenantID(parts[1])
+			if err != nil {
+				return roachpb.TenantID{}, nil
+			}
+			return tenantID, nil
+		}
+	}
+	return roachpb.SystemTenantID, nil
+}

pkg/cli/zip.go

@@ -166,6 +166,7 @@ func runDebugZip(_ *cobra.Command, args []string) (retErr error) {
 	zr := zipCtx.newZipReporter("cluster")
 
 	s := zr.start("establishing RPC connection to %s", serverCfg.AdvertiseAddr)
+	serverCfg.TenantID = zipCtx.tenantID
 	conn, _, finish, err := getClientGRPCConn(ctx, serverCfg)
 	if err != nil {
 		return s.fail(err)
@@ -192,6 +193,7 @@ func runDebugZip(_ *cobra.Command, args []string) (retErr error) {
 	s = zr.start("using SQL address: %s", sqlAddr.AddressField)
 
 	cliCtx.clientConnHost, cliCtx.clientConnPort, err = net.SplitHostPort(sqlAddr.AddressField)
+	cliCtx.tenantID = zipCtx.tenantID
 	if err != nil {
 		return s.fail(err)
 	}

pkg/cli/zip_tenant_test.go

@@ -11,6 +11,7 @@
 package cli
 
 import (
+	"fmt"
 	"os"
 	"testing"
 
@@ -34,8 +35,9 @@ func TestTenantZip(t *testing.T) {
 	skip.UnderRace(t, "test too slow under race")
 	tenantDir, tenantDirCleanupFn := testutils.TempDir(t)
 	defer tenantDirCleanupFn()
+	tenantID := serverutils.TestTenantID()
 	tenantArgs := base.TestTenantArgs{
-		TenantID:             serverutils.TestTenantID(),
+		TenantID:             tenantID,
 		HeapProfileDirName:   tenantDir,
 		GoroutineDumpDirName: tenantDir,
 	}
@@ -53,7 +55,8 @@ func TestTenantZip(t *testing.T) {
 	})
 	defer c.Cleanup()
 
-	out, err := c.RunWithCapture("debug zip --concurrency=1 --cpu-profile-duration=1s " + os.DevNull)
+	zipCmd := fmt.Sprintf("debug zip --concurrency=1 --cpu-profile-duration=1s --tenant-id=%d %s", tenantID.ToUint64(), os.DevNull)
+	out, err := c.RunWithCapture(zipCmd)
 	if err != nil {
 		t.Fatal(err)
 	}

pkg/roachpb/tenant.go

@@ -14,6 +14,8 @@ import (
 	"context"
 	"math"
 	"strconv"
+
+	"github.com/cockroachdb/errors"
 )
 
 // SystemTenantID is the ID associated with the system's internal tenant in a
@@ -97,5 +99,14 @@ func TenantFromContext(ctx context.Context) (tenID TenantID, ok bool) {
 	return
 }
 
+// ParseTenantID parses a tenant ID contained a string.
+func ParseTenantID(tenantID string) (TenantID, error) {
+	tID, err := strconv.ParseUint(tenantID, 10, 64)
+	if err != nil {
+		return TenantID{}, errors.Wrapf(err, "invalid tenant ID %s, tenant ID should be an unsigned int greater than 0", tenantID)
+	}
+	return MakeTenantID(tID), nil
+}
+
 // Silence unused warning.
 var _ = TenantFromContext