limit on api side

f0ssel · f0ssel · commit 32ff0009f477 · 2024-07-23T18:14:21.000Z
diff --git a/coderd/httpmw/provisionerdaemon.go b/coderd/httpmw/provisionerdaemon.go
@@ -51,6 +51,7 @@ func ExtractProvisionerDaemonAuthenticated(opts ExtractProvisionerAuthConfig) fu
 				return
 			}
 
+			psk := r.Header.Get(codersdk.ProvisionerDaemonPSK)
 			key := r.Header.Get(codersdk.ProvisionerDaemonKey)
 			if key == "" {
 				if opts.PSK == "" {
@@ -63,6 +64,12 @@ func ExtractProvisionerDaemonAuthenticated(opts ExtractProvisionerAuthConfig) fu
 				fallbackToPSK(ctx, opts.PSK, next, w, r, handleOptional)
 				return
 			}
+			if psk != "" {
+				handleOptional(http.StatusBadRequest, codersdk.Response{
+					Message: "provisioner daemon key and psk provided, but only one is allowed",
+				})
+				return
+			}
 
 			id, keyValue, err := provisionerkey.Parse(key)
 			if err != nil {
diff --git a/codersdk/provisionerdaemons.go b/codersdk/provisionerdaemons.go
@@ -226,9 +226,6 @@ func (c *Client) ServeProvisionerDaemon(ctx context.Context, req ServeProvisione
 
 	headers.Set(BuildVersionHeader, buildinfo.Version())
 
-	if req.ProvisionerKey != "" && req.PreSharedKey != "" {
-		return nil, xerrors.Errorf("cannot provide both a provisioner key and a pre-shared key")
-	}
 	if req.ProvisionerKey != "" {
 		headers.Set(ProvisionerDaemonKey, req.ProvisionerKey)
 	}
diff --git a/enterprise/coderd/provisionerdaemons_test.go b/enterprise/coderd/provisionerdaemons_test.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io"
 	"net/http"
+	"strings"
 	"testing"
 
 	"github.com/google/uuid"
@@ -614,28 +615,50 @@ func TestProvisionerDaemonServe(t *testing.T) {
 				name:                      "WrongKey",
 				multiOrgFeatureEnabled:    true,
 				multiOrgExperimentEnabled: true,
+				insertParams:              insertParams,
 				requestProvisionerKey:     "provisionersftw",
 				errStatusCode:             http.StatusUnauthorized,
 			},
 			{
-				name:                      "IdOKKeyWrong",
+				name:                      "IdOKKeyValueWrong",
 				multiOrgFeatureEnabled:    true,
 				multiOrgExperimentEnabled: true,
+				insertParams:              insertParams,
 				requestProvisionerKey:     insertParams.ID.String() + ":" + "wrong",
 				errStatusCode:             http.StatusUnauthorized,
 			},
 			{
-				name:                      "IdWrongKeyOK",
+				name:                      "IdWrongKeyValueOK",
 				multiOrgFeatureEnabled:    true,
 				multiOrgExperimentEnabled: true,
+				insertParams:              insertParams,
 				requestProvisionerKey:     uuid.NewString() + ":" + token,
 				errStatusCode:             http.StatusUnauthorized,
 			},
 			{
-				name:                      "TokenOnly",
+				name:                      "KeyValueOnly",
+				multiOrgFeatureEnabled:    true,
+				multiOrgExperimentEnabled: true,
+				insertParams:              insertParams,
+				requestProvisionerKey:     strings.Split(token, ":")[1],
+				errStatusCode:             http.StatusUnauthorized,
+			},
+			{
+				name:                      "KeyAndPSK",
 				multiOrgFeatureEnabled:    true,
 				multiOrgExperimentEnabled: true,
+				psk:                       "provisionersftw",
+				insertParams:              insertParams,
 				requestProvisionerKey:     token,
+				requestPSK:                "provisionersftw",
+				errStatusCode:             http.StatusUnauthorized,
+			},
+			{
+				name:                      "None",
+				multiOrgFeatureEnabled:    true,
+				multiOrgExperimentEnabled: true,
+				psk:                       "provisionersftw",
+				insertParams:              insertParams,
 				errStatusCode:             http.StatusUnauthorized,
 			},
 		}

Original file line number	Diff line number	Diff line change
`@@ -226,9 +226,6 @@ func (c *Client) ServeProvisionerDaemon(ctx context.Context, req ServeProvisione`
`226`	`226`
`227`	`227`	`headers.Set(BuildVersionHeader, buildinfo.Version())`
`228`	`228`
`229`		`- if req.ProvisionerKey != "" && req.PreSharedKey != "" {`
`230`		`- return nil, xerrors.Errorf("cannot provide both a provisioner key and a pre-shared key")`
`231`		`- }`
`232`	`229`	`if req.ProvisionerKey != "" {`
`233`	`230`	`headers.Set(ProvisionerDaemonKey, req.ProvisionerKey)`
`234`	`231`	`}`