coder · f0ssel · Jul 18, 2024 · Jul 17, 2024 · Jul 18, 2024 · Jul 18, 2024
diff --git a/codersdk/organizations.go b/codersdk/organizations.go
@@ -290,6 +290,24 @@ func (c *Client) ProvisionerDaemons(ctx context.Context) ([]ProvisionerDaemon, e
 	return daemons, json.NewDecoder(res.Body).Decode(&daemons)
 }
 
+func (c *Client) OrganizationProvisionerDaemons(ctx context.Context, organizationID uuid.UUID) ([]ProvisionerDaemon, error) {
+	res, err := c.Request(ctx, http.MethodGet,
+		fmt.Sprintf("/api/v2/organizations/%s/provisionerdaemons", organizationID.String()),
+		nil,
+	)
+	if err != nil {
+		return nil, xerrors.Errorf("execute request: %w", err)
+	}
+	defer res.Body.Close()
+
+	if res.StatusCode != http.StatusOK {
+		return nil, ReadBodyAsError(res)
+	}
+
+	var daemons []ProvisionerDaemon
+	return daemons, json.NewDecoder(res.Body).Decode(&daemons)
+}
+
 // CreateTemplateVersion processes source-code and optionally associates the version with a template.
 // Executing without a template is useful for validating source-code.
 func (c *Client) CreateTemplateVersion(ctx context.Context, organizationID uuid.UUID, req CreateTemplateVersionRequest) (TemplateVersion, error) {

diff --git a/docs/cli/provisionerd_start.md b/docs/cli/provisionerd_start.md
diff --git a/enterprise/cli/provisionerdaemons.go b/enterprise/cli/provisionerdaemons.go
@@ -4,7 +4,9 @@ package cli
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"net/http"
 	"os"
 	"regexp"
 	"time"
@@ -74,6 +76,7 @@ func (r *RootCmd) provisionerDaemonStart() *serpent.Command {
 		prometheusEnable  bool
 		prometheusAddress string
 	)
+	orgContext := agpl.NewOrganizationContext()
 	client := new(codersdk.Client)
 	cmd := &serpent.Command{
 		Use:   "start",
@@ -93,6 +96,35 @@ func (r *RootCmd) provisionerDaemonStart() *serpent.Command {
 			interruptCtx, interruptCancel := inv.SignalNotifyContext(ctx, agpl.InterruptSignals...)
 			defer interruptCancel()
 
+			// This can fail to get the current organization
+			// if the client is not authenticated as a user,
+			// like when only PSK is provided.
+			// This will be cleaner once PSK is replaced
+			// with org scoped authentication tokens.
+			org, err := orgContext.Selected(inv, client)
+			if err != nil {
+				var cErr *codersdk.Error
+				if !errors.As(err, &cErr) || cErr.StatusCode() != http.StatusUnauthorized {
+					return xerrors.Errorf("current organization: %w", err)
+				}
+
+				if preSharedKey == "" {
+					return xerrors.New("must provide a pre-shared key when not authenticated as a user")
+				}
+
+				org = codersdk.Organization{ID: uuid.Nil}
+				if orgContext.FlagSelect != "" {
+					// If we are using PSK, we can't fetch the organization
+					// to validate org name so we need the user to provide
+					// a valid organization ID.
+					orgID, err := uuid.Parse(orgContext.FlagSelect)
+					if err != nil {
+						return xerrors.New("must provide an org ID when not authenticated as a user and organization is specified")
+					}
+					org = codersdk.Organization{ID: orgID}
+				}
+			}
+
 			tags, err := agpl.ParseProvisionerTags(rawTags)
 			if err != nil {
 				return err
@@ -196,16 +228,16 @@ func (r *RootCmd) provisionerDaemonStart() *serpent.Command {
 			connector := provisionerd.LocalProvisioners{
 				string(database.ProvisionerTypeTerraform): proto.NewDRPCProvisionerClient(terraformClient),
 			}
-			id := uuid.New()
 			srv := provisionerd.New(func(ctx context.Context) (provisionerdproto.DRPCProvisionerDaemonClient, error) {
 				return client.ServeProvisionerDaemon(ctx, codersdk.ServeProvisionerDaemonRequest{
-					ID:   id,
+					ID:   uuid.New(),
 					Name: name,
 					Provisioners: []codersdk.ProvisionerType{
 						codersdk.ProvisionerTypeTerraform,
 					},
 					Tags:         tags,
 					PreSharedKey: preSharedKey,
+					Organization: org.ID,
 				})
 			}, &provisionerd.Options{
 				Logger:         logger,
@@ -346,6 +378,7 @@ func (r *RootCmd) provisionerDaemonStart() *serpent.Command {
 			Default:     "127.0.0.1:2112",
 		},
 	}
+	orgContext.AttachOptions(cmd)
 
 	return cmd
 }
diff --git a/enterprise/cli/provisionerdaemons_test.go b/enterprise/cli/provisionerdaemons_test.go
@@ -27,36 +27,134 @@ import (
 func TestProvisionerDaemon_PSK(t *testing.T) {
 	t.Parallel()
 
-	client, _ := coderdenttest.New(t, &coderdenttest.Options{
-		ProvisionerDaemonPSK: "provisionersftw",
-		LicenseOptions: &coderdenttest.LicenseOptions{
-			Features: license.Features{
-				codersdk.FeatureExternalProvisionerDaemons: 1,
+	t.Run("OK", func(t *testing.T) {
+		t.Parallel()
+
+		client, _ := coderdenttest.New(t, &coderdenttest.Options{
+			ProvisionerDaemonPSK: "provisionersftw",
+			LicenseOptions: &coderdenttest.LicenseOptions{
+				Features: license.Features{
+					codersdk.FeatureExternalProvisionerDaemons: 1,
+				},
 			},
-		},
+		})
+		inv, conf := newCLI(t, "provisionerd", "start", "--psk=provisionersftw", "--name=matt-daemon")
+		err := conf.URL().Write(client.URL.String())
+		require.NoError(t, err)
+		pty := ptytest.New(t).Attach(inv)
+		ctx, cancel := context.WithTimeout(inv.Context(), testutil.WaitLong)
+		defer cancel()
+		clitest.Start(t, inv)
+		pty.ExpectNoMatchBefore(ctx, "check entitlement", "starting provisioner daemon")
+		pty.ExpectMatchContext(ctx, "matt-daemon")
+
+		var daemons []codersdk.ProvisionerDaemon
+		require.Eventually(t, func() bool {
+			daemons, err = client.ProvisionerDaemons(ctx)
+			if err != nil {
+				return false
+			}
+			return len(daemons) == 1
+		}, testutil.WaitShort, testutil.IntervalSlow)
+		require.Equal(t, "matt-daemon", daemons[0].Name)
+		require.Equal(t, provisionersdk.ScopeOrganization, daemons[0].Tags[provisionersdk.TagScope])
+		require.Equal(t, buildinfo.Version(), daemons[0].Version)
+		require.Equal(t, proto.CurrentVersion.String(), daemons[0].APIVersion)
+	})
+
+	t.Run("AnotherOrg", func(t *testing.T) {
+		t.Parallel()
+		client, _ := coderdenttest.New(t, &coderdenttest.Options{
+			ProvisionerDaemonPSK: "provisionersftw",
+			LicenseOptions: &coderdenttest.LicenseOptions{
+				Features: license.Features{
+					codersdk.FeatureExternalProvisionerDaemons: 1,
+				},
+			},
+		})
+		anotherOrg := coderdtest.CreateOrganization(t, client, coderdtest.CreateOrganizationOptions{})
+		inv, conf := newCLI(t, "provisionerd", "start", "--psk=provisionersftw", "--name", "org-daemon", "--org", anotherOrg.ID.String())
+		err := conf.URL().Write(client.URL.String())
+		require.NoError(t, err)
+		pty := ptytest.New(t).Attach(inv)
+		ctx, cancel := context.WithTimeout(inv.Context(), testutil.WaitLong)
+		defer cancel()
+		clitest.Start(t, inv)
+		pty.ExpectMatchContext(ctx, "starting provisioner daemon")
+
+		var daemons []codersdk.ProvisionerDaemon
+		require.Eventually(t, func() bool {
+			daemons, err = client.OrganizationProvisionerDaemons(ctx, anotherOrg.ID)
+			if err != nil {
+				return false
+			}
+			return len(daemons) == 1
+		}, testutil.WaitShort, testutil.IntervalSlow)
+		assert.Equal(t, "org-daemon", daemons[0].Name)
+		assert.Equal(t, provisionersdk.ScopeOrganization, daemons[0].Tags[provisionersdk.TagScope])
+		assert.Equal(t, buildinfo.Version(), daemons[0].Version)
+		assert.Equal(t, proto.CurrentVersion.String(), daemons[0].APIVersion)
+	})
+
+	t.Run("AnotherOrgByNameWithUser", func(t *testing.T) {
+		t.Parallel()
+		client, _ := coderdenttest.New(t, &coderdenttest.Options{
+			ProvisionerDaemonPSK: "provisionersftw",
+			LicenseOptions: &coderdenttest.LicenseOptions{
+				Features: license.Features{
+					codersdk.FeatureExternalProvisionerDaemons: 1,
+				},
+			},
+		})
+		anotherOrg := coderdtest.CreateOrganization(t, client, coderdtest.CreateOrganizationOptions{})
+		anotherClient, _ := coderdtest.CreateAnotherUser(t, client, anotherOrg.ID, rbac.RoleTemplateAdmin())
+		inv, conf := newCLI(t, "provisionerd", "start", "--psk=provisionersftw", "--name", "org-daemon", "--org", anotherOrg.Name)
+		clitest.SetupConfig(t, anotherClient, conf)
+		pty := ptytest.New(t).Attach(inv)
+		ctx, cancel := context.WithTimeout(inv.Context(), testutil.WaitLong)
+		defer cancel()
+		clitest.Start(t, inv)
+		pty.ExpectMatchContext(ctx, "starting provisioner daemon")
+	})
+
+	t.Run("AnotherOrgByNameNoUser", func(t *testing.T) {
+		t.Parallel()
+		client, _ := coderdenttest.New(t, &coderdenttest.Options{
+			ProvisionerDaemonPSK: "provisionersftw",
+			LicenseOptions: &coderdenttest.LicenseOptions{
+				Features: license.Features{
+					codersdk.FeatureExternalProvisionerDaemons: 1,
+				},
+			},
+		})
+		anotherOrg := coderdtest.CreateOrganization(t, client, coderdtest.CreateOrganizationOptions{})
+		inv, conf := newCLI(t, "provisionerd", "start", "--psk=provisionersftw", "--name", "org-daemon", "--org", anotherOrg.Name)
+		err := conf.URL().Write(client.URL.String())
+		require.NoError(t, err)
+		ctx, cancel := context.WithTimeout(inv.Context(), testutil.WaitLong)
+		defer cancel()
+		err = inv.WithContext(ctx).Run()
+		require.ErrorContains(t, err, "must provide an org ID when not authenticated as a user and organization is specified")
+	})
+
+	t.Run("NoUserNoPSK", func(t *testing.T) {
+		t.Parallel()
+		client, _ := coderdenttest.New(t, &coderdenttest.Options{
+			ProvisionerDaemonPSK: "provisionersftw",
+			LicenseOptions: &coderdenttest.LicenseOptions{
+				Features: license.Features{
+					codersdk.FeatureExternalProvisionerDaemons: 1,
+				},
+			},
+		})
+		inv, conf := newCLI(t, "provisionerd", "start", "--name", "org-daemon")
+		err := conf.URL().Write(client.URL.String())
+		require.NoError(t, err)
+		ctx, cancel := context.WithTimeout(inv.Context(), testutil.WaitLong)
+		defer cancel()
+		err = inv.WithContext(ctx).Run()
+		require.ErrorContains(t, err, "must provide a pre-shared key when not authenticated as a user")
 	})
-	inv, conf := newCLI(t, "provisionerd", "start", "--psk=provisionersftw", "--name=matt-daemon")
-	err := conf.URL().Write(client.URL.String())
-	require.NoError(t, err)
-	pty := ptytest.New(t).Attach(inv)
-	ctx, cancel := context.WithTimeout(inv.Context(), testutil.WaitLong)
-	defer cancel()
-	clitest.Start(t, inv)
-	pty.ExpectNoMatchBefore(ctx, "check entitlement", "starting provisioner daemon")
-	pty.ExpectMatchContext(ctx, "matt-daemon")
-
-	var daemons []codersdk.ProvisionerDaemon
-	require.Eventually(t, func() bool {
-		daemons, err = client.ProvisionerDaemons(ctx)
-		if err != nil {
-			return false
-		}
-		return len(daemons) == 1
-	}, testutil.WaitShort, testutil.IntervalSlow)
-	require.Equal(t, "matt-daemon", daemons[0].Name)
-	require.Equal(t, provisionersdk.ScopeOrganization, daemons[0].Tags[provisionersdk.TagScope])
-	require.Equal(t, buildinfo.Version(), daemons[0].Version)
-	require.Equal(t, proto.CurrentVersion.String(), daemons[0].APIVersion)
 }
 
 func TestProvisionerDaemon_SessionToken(t *testing.T) {
@@ -166,6 +264,42 @@ func TestProvisionerDaemon_SessionToken(t *testing.T) {
 		assert.Equal(t, proto.CurrentVersion.String(), daemons[0].APIVersion)
 	})
 
+	t.Run("ScopeUserAnotherOrg", func(t *testing.T) {
+		t.Parallel()
+		client, _ := coderdenttest.New(t, &coderdenttest.Options{
+			ProvisionerDaemonPSK: "provisionersftw",
+			LicenseOptions: &coderdenttest.LicenseOptions{
+				Features: license.Features{
+					codersdk.FeatureExternalProvisionerDaemons: 1,
+				},
+			},
+		})
+		anotherOrg := coderdtest.CreateOrganization(t, client, coderdtest.CreateOrganizationOptions{})
+		anotherClient, anotherUser := coderdtest.CreateAnotherUser(t, client, anotherOrg.ID, rbac.RoleTemplateAdmin())
+		inv, conf := newCLI(t, "provisionerd", "start", "--tag", "scope=user", "--name", "org-daemon", "--org", anotherOrg.ID.String())
+		clitest.SetupConfig(t, anotherClient, conf)
+		pty := ptytest.New(t).Attach(inv)
+		ctx, cancel := context.WithTimeout(inv.Context(), testutil.WaitLong)
+		defer cancel()
+		clitest.Start(t, inv)
+		pty.ExpectMatchContext(ctx, "starting provisioner daemon")
+
+		var daemons []codersdk.ProvisionerDaemon
+		var err error
+		require.Eventually(t, func() bool {
+			daemons, err = client.OrganizationProvisionerDaemons(ctx, anotherOrg.ID)
+			if err != nil {
+				return false
+			}
+			return len(daemons) == 1
+		}, testutil.WaitShort, testutil.IntervalSlow)
+		assert.Equal(t, "org-daemon", daemons[0].Name)
+		assert.Equal(t, provisionersdk.ScopeUser, daemons[0].Tags[provisionersdk.TagScope])
+		assert.Equal(t, anotherUser.ID.String(), daemons[0].Tags[provisionersdk.TagOwner])
+		assert.Equal(t, buildinfo.Version(), daemons[0].Version)
+		assert.Equal(t, proto.CurrentVersion.String(), daemons[0].APIVersion)
+	})
+
 	t.Run("PrometheusEnabled", func(t *testing.T) {
 		t.Parallel()
 

diff --git a/enterprise/cli/testdata/coder_provisionerd_start_--help.golden b/enterprise/cli/testdata/coder_provisionerd_start_--help.golden
@@ -6,6 +6,9 @@ USAGE:
   Run a provisioner daemon
 
 OPTIONS:
+  -O, --org string, $CODER_ORGANIZATION
+          Select which organization (uuid or name) to use.
+
   -c, --cache-dir string, $CODER_CACHE_DIRECTORY (default: [cache dir])
           Directory to store cached data.