Add feature for external daemons

coder · kylecarbs · Nov 16, 2022 · Oct 25, 2022 · Nov 7, 2022 · Nov 7, 2022
commit cf1221b1e8c2d111a44db89cf0fcbd213b067a56
diff --git a/coderd/coderdtest/coderdtest.go b/coderd/coderdtest/coderdtest.go
@@ -66,6 +66,7 @@ import (
 	"github.com/coder/coder/cryptorand"
 	"github.com/coder/coder/provisioner/echo"
 	"github.com/coder/coder/provisionerd"
+	provisionerdproto "github.com/coder/coder/provisionerd/proto"
 	"github.com/coder/coder/provisionersdk"
 	"github.com/coder/coder/provisionersdk/proto"
 	"github.com/coder/coder/tailnet"
@@ -337,6 +338,41 @@ func NewProvisionerDaemon(t *testing.T, coderAPI *coderd.API) io.Closer {
 	return closer
 }
 
+func NewExternalProvisionerDaemon(t *testing.T, client *codersdk.Client, org uuid.UUID, tags map[string]string) io.Closer {
+	echoClient, echoServer := provisionersdk.TransportPipe()
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	t.Cleanup(func() {
+		_ = echoClient.Close()
+		_ = echoServer.Close()
+		cancelFunc()
+	})
+	fs := afero.NewMemMapFs()
+	go func() {
+		err := echo.Serve(ctx, fs, &provisionersdk.ServeOptions{
+			Listener: echoServer,
+		})
+		assert.NoError(t, err)
+	}()
+
+	closer := provisionerd.New(func(ctx context.Context) (provisionerdproto.DRPCProvisionerDaemonClient, error) {
+		return client.ServeProvisionerDaemon(ctx, org, []codersdk.ProvisionerType{codersdk.ProvisionerTypeEcho}, tags)
+	}, &provisionerd.Options{
+		Filesystem:          fs,
+		Logger:              slogtest.Make(t, nil).Named("provisionerd").Leveled(slog.LevelDebug),
+		PollInterval:        50 * time.Millisecond,
+		UpdateInterval:      250 * time.Millisecond,
+		ForceCancelInterval: time.Second,
+		Provisioners: provisionerd.Provisioners{
+			string(database.ProvisionerTypeEcho): proto.NewDRPCProvisionerClient(provisionersdk.Conn(echoClient)),
+		},
+		WorkDirectory: t.TempDir(),
+	})
+	t.Cleanup(func() {
+		_ = closer.Close()
+	})
+	return closer
+}
+
 var FirstUserParams = codersdk.CreateFirstUserRequest{
 	Email:            "testuser@coder.com",
 	Username:         "testuser",

diff --git a/codersdk/features.go b/codersdk/features.go
@@ -15,14 +15,15 @@ const (
 )
 
 const (
-	FeatureUserLimit        = "user_limit"
-	FeatureAuditLog         = "audit_log"
-	FeatureBrowserOnly      = "browser_only"
-	FeatureSCIM             = "scim"
-	FeatureWorkspaceQuota   = "workspace_quota"
-	FeatureTemplateRBAC     = "template_rbac"
-	FeatureHighAvailability = "high_availability"
-	FeatureMultipleGitAuth  = "multiple_git_auth"
+	FeatureUserLimit                  = "user_limit"
+	FeatureAuditLog                   = "audit_log"
+	FeatureBrowserOnly                = "browser_only"
+	FeatureSCIM                       = "scim"
+	FeatureWorkspaceQuota             = "workspace_quota"
+	FeatureTemplateRBAC               = "template_rbac"
+	FeatureHighAvailability           = "high_availability"
+	FeatureMultipleGitAuth            = "multiple_git_auth"
+	FeatureExternalProvisionerDaemons = "external_provisioner_daemons"
 )
 
 var FeatureNames = []string{
@@ -34,6 +35,7 @@ var FeatureNames = []string{
 	FeatureTemplateRBAC,
 	FeatureHighAvailability,
 	FeatureMultipleGitAuth,
+	FeatureExternalProvisionerDaemons,
 }
 
 type Feature struct {

diff --git a/enterprise/coderd/coderdenttest/coderdenttest.go b/enterprise/coderd/coderdenttest/coderdenttest.go
@@ -100,20 +100,21 @@ func NewWithAPI(t *testing.T, options *Options) (*codersdk.Client, io.Closer, *c
 }
 
 type LicenseOptions struct {
-	AccountType      string
-	AccountID        string
-	Trial            bool
-	AllFeatures      bool
-	GraceAt          time.Time
-	ExpiresAt        time.Time
-	UserLimit        int64
-	AuditLog         bool
-	BrowserOnly      bool
-	SCIM             bool
-	WorkspaceQuota   bool
-	TemplateRBAC     bool
-	HighAvailability bool
-	MultipleGitAuth  bool
+	AccountType                string
+	AccountID                  string
+	Trial                      bool
+	AllFeatures                bool
+	GraceAt                    time.Time
+	ExpiresAt                  time.Time
+	UserLimit                  int64
+	AuditLog                   bool
+	BrowserOnly                bool
+	SCIM                       bool
+	WorkspaceQuota             bool
+	TemplateRBAC               bool
+	HighAvailability           bool
+	MultipleGitAuth            bool
+	ExternalProvisionerDaemons bool
 }
 
 // AddLicense generates a new license with the options provided and inserts it.
@@ -164,6 +165,11 @@ func GenerateLicense(t *testing.T, options LicenseOptions) string {
 		multipleGitAuth = 1
 	}
 
+	externalProvisionerDaemons := int64(0)
+	if options.ExternalProvisionerDaemons {
+		externalProvisionerDaemons = 1
+	}
+
 	c := &license.Claims{
 		RegisteredClaims: jwt.RegisteredClaims{
 			Issuer:    "test@testing.test",
@@ -178,14 +184,15 @@ func GenerateLicense(t *testing.T, options LicenseOptions) string {
 		Version:        license.CurrentVersion,
 		AllFeatures:    options.AllFeatures,
 		Features: license.Features{
-			UserLimit:        options.UserLimit,
-			AuditLog:         auditLog,
-			BrowserOnly:      browserOnly,
-			SCIM:             scim,
-			WorkspaceQuota:   workspaceQuota,
-			HighAvailability: highAvailability,
-			TemplateRBAC:     rbacEnabled,
-			MultipleGitAuth:  multipleGitAuth,
+			UserLimit:                  options.UserLimit,
+			AuditLog:                   auditLog,
+			BrowserOnly:                browserOnly,
+			SCIM:                       scim,
+			WorkspaceQuota:             workspaceQuota,
+			HighAvailability:           highAvailability,
+			TemplateRBAC:               rbacEnabled,
+			MultipleGitAuth:            multipleGitAuth,
+			ExternalProvisionerDaemons: externalProvisionerDaemons,
 		},
 	}
 	tok := jwt.NewWithClaims(jwt.SigningMethodEdDSA, c)

diff --git a/enterprise/coderd/license/license.go b/enterprise/coderd/license/license.go
@@ -123,6 +123,12 @@ func Entitlements(
 				Enabled:     true,
 			}
 		}
+		if claims.Features.ExternalProvisionerDaemons > 0 {
+			entitlements.Features[codersdk.FeatureExternalProvisionerDaemons] = codersdk.Feature{
+				Entitlement: entitlement,
+				Enabled:     true,
+			}
+		}
 		if claims.AllFeatures {
 			allFeatures = true
 		}
@@ -244,14 +250,15 @@ var (
 )
 
 type Features struct {
-	UserLimit        int64 `json:"user_limit"`
-	AuditLog         int64 `json:"audit_log"`
-	BrowserOnly      int64 `json:"browser_only"`
-	SCIM             int64 `json:"scim"`
-	WorkspaceQuota   int64 `json:"workspace_quota"`
-	TemplateRBAC     int64 `json:"template_rbac"`
-	HighAvailability int64 `json:"high_availability"`
-	MultipleGitAuth  int64 `json:"multiple_git_auth"`
+	UserLimit                  int64 `json:"user_limit"`
+	AuditLog                   int64 `json:"audit_log"`
+	BrowserOnly                int64 `json:"browser_only"`
+	SCIM                       int64 `json:"scim"`
+	WorkspaceQuota             int64 `json:"workspace_quota"`
+	TemplateRBAC               int64 `json:"template_rbac"`
+	HighAvailability           int64 `json:"high_availability"`
+	MultipleGitAuth            int64 `json:"multiple_git_auth"`
+	ExternalProvisionerDaemons int64 `json:"external_provisioner_daemons"`
 }
 
 type Claims struct {

diff --git a/enterprise/coderd/license/license_test.go b/enterprise/coderd/license/license_test.go
@@ -20,13 +20,14 @@ import (
 func TestEntitlements(t *testing.T) {
 	t.Parallel()
 	all := map[string]bool{
-		codersdk.FeatureAuditLog:         true,
-		codersdk.FeatureBrowserOnly:      true,
-		codersdk.FeatureSCIM:             true,
-		codersdk.FeatureWorkspaceQuota:   true,
-		codersdk.FeatureHighAvailability: true,
-		codersdk.FeatureTemplateRBAC:     true,
-		codersdk.FeatureMultipleGitAuth:  true,
+		codersdk.FeatureAuditLog:                   true,
+		codersdk.FeatureBrowserOnly:                true,
+		codersdk.FeatureSCIM:                       true,
+		codersdk.FeatureWorkspaceQuota:             true,
+		codersdk.FeatureHighAvailability:           true,
+		codersdk.FeatureTemplateRBAC:               true,
+		codersdk.FeatureMultipleGitAuth:            true,
+		codersdk.FeatureExternalProvisionerDaemons: true,
 	}
 
 	t.Run("Defaults", func(t *testing.T) {
@@ -62,14 +63,15 @@ func TestEntitlements(t *testing.T) {
 		db := databasefake.New()
 		db.InsertLicense(context.Background(), database.InsertLicenseParams{
 			JWT: coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
-				UserLimit:        100,
-				AuditLog:         true,
-				BrowserOnly:      true,
-				SCIM:             true,
-				WorkspaceQuota:   true,
-				HighAvailability: true,
-				TemplateRBAC:     true,
-				MultipleGitAuth:  true,
+				UserLimit:                  100,
+				AuditLog:                   true,
+				BrowserOnly:                true,
+				SCIM:                       true,
+				WorkspaceQuota:             true,
+				HighAvailability:           true,
+				TemplateRBAC:               true,
+				MultipleGitAuth:            true,
+				ExternalProvisionerDaemons: true,
 			}),
 			Exp: time.Now().Add(time.Hour),
 		})
@@ -86,15 +88,16 @@ func TestEntitlements(t *testing.T) {
 		db := databasefake.New()
 		db.InsertLicense(context.Background(), database.InsertLicenseParams{
 			JWT: coderdenttest.GenerateLicense(t, coderdenttest.LicenseOptions{
-				UserLimit:        100,
-				AuditLog:         true,
-				BrowserOnly:      true,
-				SCIM:             true,
-				WorkspaceQuota:   true,
-				HighAvailability: true,
-				TemplateRBAC:     true,
-				GraceAt:          time.Now().Add(-time.Hour),
-				ExpiresAt:        time.Now().Add(time.Hour),
+				UserLimit:                  100,
+				AuditLog:                   true,
+				BrowserOnly:                true,
+				SCIM:                       true,
+				WorkspaceQuota:             true,
+				HighAvailability:           true,
+				TemplateRBAC:               true,
+				ExternalProvisionerDaemons: true,
+				GraceAt:                    time.Now().Add(-time.Hour),
+				ExpiresAt:                  time.Now().Add(time.Hour),
 			}),
 			Exp: time.Now().Add(time.Hour),
 		})

diff --git a/enterprise/coderd/provisionerdaemons.go b/enterprise/coderd/provisionerdaemons.go
@@ -29,6 +29,21 @@ import (
 	"github.com/coder/coder/provisionerd/proto"
 )
 
+func (api *API) provisionerDaemonsEnabledMW(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		api.entitlementsMu.RLock()
+		epd := api.entitlements.Features[codersdk.FeatureExternalProvisionerDaemons].Enabled
+		api.entitlementsMu.RUnlock()
+
+		if !epd {
+			httpapi.Write(r.Context(), rw, http.Status)
+			return
+		}
+
+		next.ServeHTTP(rw, r)
+	})
+}
+
 func (api *API) provisionerDaemons(rw http.ResponseWriter, r *http.Request) {
 	ctx := r.Context()
 	org := httpmw.OrganizationParam(r)
@@ -64,6 +79,7 @@ func (api *API) provisionerDaemons(rw http.ResponseWriter, r *http.Request) {
 
 // Serves the provisioner daemon protobuf API over a WebSocket.
 func (api *API) provisionerDaemonServe(rw http.ResponseWriter, r *http.Request) {
+
 	tags := map[string]string{}
 	if r.URL.Query().Has("tag") {
 		for _, tag := range r.URL.Query()["tag"] {

diff --git a/enterprise/coderd/provisionerdaemons_test.go b/enterprise/coderd/provisionerdaemons_test.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"net/http"
 	"testing"
-	"time"
 
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
@@ -14,8 +13,6 @@ import (
 	"github.com/coder/coder/codersdk"
 	"github.com/coder/coder/enterprise/coderd/coderdenttest"
 	"github.com/coder/coder/provisioner/echo"
-	"github.com/coder/coder/provisionerd"
-	provisionerdproto "github.com/coder/coder/provisionerd/proto"
 	"github.com/coder/coder/provisionersdk/proto"
 )
 
@@ -52,14 +49,10 @@ func TestProvisionerDaemonServe(t *testing.T) {
 		t.Parallel()
 		client := coderdenttest.New(t, nil)
 		user := coderdtest.CreateFirstUser(t, client)
-		srv := provisionerd.New(func(ctx context.Context) (provisionerdproto.DRPCProvisionerDaemonClient, error) {
-			return client.ServeProvisionerDaemon(context.Background(), user.OrganizationID, []codersdk.ProvisionerType{
-				codersdk.ProvisionerTypeEcho,
-			}, map[string]string{
-				provisionerdserver.TagScope: provisionerdserver.ScopeUser,
-			})
-		}, nil)
-		defer srv.Close()
+		closer := coderdtest.NewExternalProvisionerDaemon(t, client, user.OrganizationID, map[string]string{
+			provisionerdserver.TagScope: provisionerdserver.ScopeUser,
+		})
+		defer closer.Close()
 
 		authToken := uuid.NewString()
 		data, err := echo.Tar(&echo.Responses{
@@ -100,7 +93,7 @@ func TestProvisionerDaemonServe(t *testing.T) {
 		file, err := client.Upload(context.Background(), codersdk.ContentTypeTar, data)
 		require.NoError(t, err)
 
-		_, err = client.CreateTemplateVersion(context.Background(), user.OrganizationID, codersdk.CreateTemplateVersionRequest{
+		version, err := client.CreateTemplateVersion(context.Background(), user.OrganizationID, codersdk.CreateTemplateVersionRequest{
 			Name:          "example",
 			StorageMethod: codersdk.ProvisionerStorageMethodFile,
 			FileID:        file.ID,
@@ -110,9 +103,16 @@ func TestProvisionerDaemonServe(t *testing.T) {
 			},
 		})
 		require.NoError(t, err)
-		// coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
-
-		time.Sleep(time.Second)
+		coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
+		template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
+		another := coderdtest.CreateAnotherUser(t, client, user.OrganizationID)
+		_ = closer.Close()
+		closer = coderdtest.NewExternalProvisionerDaemon(t, another, user.OrganizationID, map[string]string{
+			provisionerdserver.TagScope: provisionerdserver.ScopeUser,
+		})
+		defer closer.Close()
+		workspace := coderdtest.CreateWorkspace(t, another, user.OrganizationID, template.ID)
+		coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
 	})
 }