From b35894af7e65956402e9a597762f47c05739da86 Mon Sep 17 00:00:00 2001
From: btoews <benjamin.toews@gmail.com>
Date: Mon, 20 Nov 2023 08:14:05 -0700
Subject: [PATCH] flyio: add NoAdminFeatures caveat

---
 caveat.go             |  1 +
 flyio/caveats.go      | 64 +++++++++++++++++++++++++++++++++++++++++++
 flyio/caveats_test.go | 58 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 123 insertions(+)

diff --git a/caveat.go b/caveat.go
index ce07ec8..d8b7e74 100644
--- a/caveat.go
+++ b/caveat.go
@@ -35,6 +35,7 @@ const (
 	CavAuthConfineGoogleHD
 	CavAuthConfineGitHubOrg
 	CavAuthMaxValidity
+	CavNoAdminFeatures
 
 	// Globally-recognized user-registerable caveat types may be requested via
 	// pull requests to this repository. Add a meaningful name of the caveat
diff --git a/flyio/caveats.go b/flyio/caveats.go
index 893b0a6..1eb0e01 100644
--- a/flyio/caveats.go
+++ b/flyio/caveats.go
@@ -18,6 +18,7 @@ const (
 	CavMachineFeatureSet = macaroon.CavFlyioMachineFeatureSet
 	CavFromMachineSource = macaroon.CavFlyioFromMachineSource
 	CavClusters          = macaroon.CavFlyioClusters
+	CavNoAdminFeatures   = macaroon.CavNoAdminFeatures
 )
 
 type FromMachine struct {
@@ -233,3 +234,66 @@ func (c *Clusters) Prohibits(a macaroon.Access) error {
 
 	return c.Clusters.Prohibits(f.Cluster, f.Action)
 }
+
+var (
+	MemberFeatures = map[string]resset.Action{
+		"wg":           resset.ActionAll,
+		"domain":       resset.ActionAll,
+		"site":         resset.ActionAll,
+		"builder":      resset.ActionAll,
+		"addon":        resset.ActionAll,
+		"checks":       resset.ActionAll,
+		"litefs-cloud": resset.ActionAll,
+
+		"membership": resset.ActionRead,
+		"billing":    resset.ActionRead,
+
+		"deletion":         resset.ActionNone,
+		"document_signing": resset.ActionNone,
+	}
+)
+
+// NoAdminFeatures is a shorthand for specifying that the token isn't allowed to
+// access admin-only features. Same as:
+//
+//	resset.IfPresent{
+//	  Ifs: macaroon.NewCaveatSet(&FeatureSet{
+//	    "memberFeatureOne": resset.ActionAll,
+//	    "memberFeatureTwo": resset.ActionAll,
+//	    "memberFeatureNNN": resset.ActionAll,
+//	  }),
+//	  Else: resset.ActionAll
+//	}
+type NoAdminFeatures struct{}
+
+func init()                                                { macaroon.RegisterCaveatType(&NoAdminFeatures{}) }
+func (c *NoAdminFeatures) CaveatType() macaroon.CaveatType { return CavNoAdminFeatures }
+func (c *NoAdminFeatures) Name() string                    { return "NoAdminFeatures" }
+
+func (c *NoAdminFeatures) Prohibits(a macaroon.Access) error {
+	f, isFlyioAccess := a.(*Access)
+	if !isFlyioAccess {
+		return macaroon.ErrInvalidAccess
+	}
+	if f.Feature == nil {
+		return nil
+	}
+	if *f.Feature == "" {
+		return fmt.Errorf("%w admin org features", resset.ErrUnauthorizedForResource)
+	}
+
+	memberPermission, ok := MemberFeatures[*f.Feature]
+	if !ok {
+		return fmt.Errorf("%w %s", resset.ErrUnauthorizedForResource, *f.Feature)
+	}
+	if !f.Action.IsSubsetOf(memberPermission) {
+		return fmt.Errorf(
+			"%w %s access to %s",
+			resset.ErrUnauthorizedForAction,
+			f.Action.Remove(memberPermission),
+			*f.Feature,
+		)
+	}
+
+	return nil
+}
diff --git a/flyio/caveats_test.go b/flyio/caveats_test.go
index 88f76f4..404ebce 100644
--- a/flyio/caveats_test.go
+++ b/flyio/caveats_test.go
@@ -21,6 +21,7 @@ func TestCaveatSerialization(t *testing.T) {
 		&MachineFeatureSet{Features: resset.New(resset.ActionRead, "123")},
 		&FromMachine{ID: "asdf"},
 		&Clusters{Clusters: resset.New(resset.ActionRead, "123")},
+		&NoAdminFeatures{},
 	)
 
 	b, err := json.Marshal(cs)
@@ -37,3 +38,60 @@ func TestCaveatSerialization(t *testing.T) {
 	assert.NoError(t, err)
 	assert.Equal(t, cs, cs2)
 }
+
+func TestNoAdminFeatures(t *testing.T) {
+	cs := macaroon.NewCaveatSet(&NoAdminFeatures{})
+
+	yes := func(access *Access) {
+		t.Helper()
+		assert.NoError(t, cs.Validate(access))
+	}
+
+	no := func(access *Access, target error) {
+		t.Helper()
+		err := cs.Validate(access)
+		assert.Error(t, err)
+		assert.IsError(t, err, target)
+	}
+
+	yes(&Access{
+		DeprecatedOrgID: uptr(1),
+		Action:          resset.ActionAll,
+		Feature:         ptr("wg"),
+	})
+
+	yes(&Access{
+		DeprecatedOrgID: uptr(1),
+		Action:          resset.ActionRead,
+		Feature:         ptr("membership"),
+	})
+
+	yes(&Access{
+		DeprecatedOrgID: uptr(1),
+		Action:          resset.ActionAll,
+	})
+
+	no(&Access{
+		DeprecatedOrgID: uptr(1),
+		Action:          resset.ActionWrite,
+		Feature:         ptr("membership"),
+	}, resset.ErrUnauthorizedForAction)
+
+	no(&Access{
+		DeprecatedOrgID: uptr(1),
+		Action:          resset.ActionRead,
+		Feature:         ptr("unknown"),
+	}, resset.ErrUnauthorizedForResource)
+
+	no(&Access{
+		DeprecatedOrgID: uptr(1),
+		Action:          resset.ActionNone,
+		Feature:         ptr(""),
+	}, resset.ErrUnauthorizedForResource)
+
+	no(&Access{
+		DeprecatedOrgID: uptr(1),
+		Action:          resset.ActionNone,
+		Feature:         ptr(""),
+	}, resset.ErrUnauthorizedForResource)
+}