Merge pull request #16005 from tjungblu/putauthshort

Early exit auth check on lease puts

server/etcdserver/apply/apply_auth.go

@@ -125,16 +125,27 @@ func (aa *authApplierV3) LeaseRevoke(lc *pb.LeaseRevokeRequest) (*pb.LeaseRevoke
 func (aa *authApplierV3) checkLeasePuts(leaseID lease.LeaseID) error {
 	l := aa.lessor.Lookup(leaseID)
 	if l != nil {
-		for _, key := range l.Keys() {
-			if err := aa.as.IsPutPermitted(&aa.authInfo, []byte(key)); err != nil {
-				return err
-			}
-		}
+		return aa.checkLeasePutsKeys(l)
 	}
 
 	return nil
 }
 
+func (aa *authApplierV3) checkLeasePutsKeys(l *lease.Lease) error {
+	// early return for most-common scenario of either disabled auth or admin user.
+	// IsAdminPermitted also checks whether auth is enabled
+	if err := aa.as.IsAdminPermitted(&aa.authInfo); err == nil {
+		return nil
+	}
+
+	for _, key := range l.Keys() {
+		if err := aa.as.IsPutPermitted(&aa.authInfo, []byte(key)); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 func (aa *authApplierV3) UserGet(r *pb.AuthUserGetRequest) (*pb.AuthUserGetResponse, error) {
 	err := aa.as.IsAdminPermitted(&aa.authInfo)
 	if err != nil && r.Name != aa.authInfo.Username {

server/etcdserver/apply/apply_auth_test.go

@@ -0,0 +1,124 @@
+// Copyright 2023 The etcd Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package apply
+
+import (
+	"encoding/base64"
+	"testing"
+	"time"
+
+	"golang.org/x/crypto/bcrypt"
+
+	betesting "go.etcd.io/etcd/server/v3/storage/backend/testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/zap/zaptest"
+
+	"go.etcd.io/etcd/server/v3/storage/schema"
+
+	"go.etcd.io/etcd/api/v3/authpb"
+	pb "go.etcd.io/etcd/api/v3/etcdserverpb"
+	"go.etcd.io/etcd/server/v3/auth"
+	"go.etcd.io/etcd/server/v3/lease"
+)
+
+func TestCheckLeasePutsKeys(t *testing.T) {
+	lg := zaptest.NewLogger(t)
+
+	b, _ := betesting.NewDefaultTmpBackend(t)
+	defer betesting.Close(t, b)
+
+	simpleTokenTTLDefault := 300 * time.Second
+	tokenTypeSimple := "simple"
+	dummyIndexWaiter := func(index uint64) <-chan struct{} {
+		ch := make(chan struct{}, 1)
+		go func() {
+			ch <- struct{}{}
+		}()
+		return ch
+	}
+
+	tp, _ := auth.NewTokenProvider(zaptest.NewLogger(t), tokenTypeSimple, dummyIndexWaiter, simpleTokenTTLDefault)
+	as := auth.NewAuthStore(lg, schema.NewAuthBackend(lg, b), tp, bcrypt.MinCost)
+
+	aa := authApplierV3{as: as}
+	assert.NoError(t, aa.checkLeasePutsKeys(lease.NewLease(lease.LeaseID(1), 3600)), "auth is disabled, should allow puts")
+	assert.NoError(t, enableAuthAndCreateRoot(aa.as), "error while enabling auth")
+	aa.authInfo = auth.AuthInfo{Username: "root"}
+	assert.NoError(t, aa.checkLeasePutsKeys(lease.NewLease(lease.LeaseID(1), 3600)), "auth is enabled, should allow puts for root")
+
+	l := lease.NewLease(lease.LeaseID(1), 3600)
+	l.SetLeaseItem(lease.LeaseItem{Key: "a"})
+	aa.authInfo = auth.AuthInfo{Username: "bob", Revision: 0}
+	assert.ErrorIs(t, aa.checkLeasePutsKeys(l), auth.ErrUserEmpty, "auth is enabled, should not allow bob, non existing at rev 0")
+	aa.authInfo = auth.AuthInfo{Username: "bob", Revision: 1}
+	assert.ErrorIs(t, aa.checkLeasePutsKeys(l), auth.ErrAuthOldRevision, "auth is enabled, old revision")
+
+	aa.authInfo = auth.AuthInfo{Username: "bob", Revision: aa.as.Revision()}
+	assert.ErrorIs(t, aa.checkLeasePutsKeys(l), auth.ErrPermissionDenied, "auth is enabled, bob does not have permissions, bob does not exist")
+	_, err := aa.as.UserAdd(&pb.AuthUserAddRequest{Name: "bob", Options: &authpb.UserAddOptions{NoPassword: true}})
+	assert.NoError(t, err, "bob should be added without error")
+	aa.authInfo = auth.AuthInfo{Username: "bob", Revision: aa.as.Revision()}
+	assert.ErrorIs(t, aa.checkLeasePutsKeys(l), auth.ErrPermissionDenied, "auth is enabled, bob exists yet does not have permissions")
+
+	// allow bob to access "a"
+	_, err = aa.as.RoleAdd(&pb.AuthRoleAddRequest{Name: "bobsrole"})
+	assert.NoError(t, err, "bobsrole should be added without error")
+	_, err = aa.as.RoleGrantPermission(&pb.AuthRoleGrantPermissionRequest{
+		Name: "bobsrole",
+		Perm: &authpb.Permission{
+			PermType: authpb.READWRITE,
+			Key:      []byte("a"),
+			RangeEnd: nil,
+		},
+	})
+	assert.NoError(t, err, "bobsrole should be granted permissions without error")
+	_, err = aa.as.UserGrantRole(&pb.AuthUserGrantRoleRequest{
+		User: "bob",
+		Role: "bobsrole",
+	})
+	assert.NoError(t, err, "bob should be granted bobsrole without error")
+
+	aa.authInfo = auth.AuthInfo{Username: "bob", Revision: aa.as.Revision()}
+	assert.NoError(t, aa.checkLeasePutsKeys(l), "bob should be able to access key 'a'")
+
+}
+
+func enableAuthAndCreateRoot(as auth.AuthStore) error {
+	_, err := as.UserAdd(&pb.AuthUserAddRequest{
+		Name:           "root",
+		HashedPassword: encodePassword("root"),
+		Options:        &authpb.UserAddOptions{NoPassword: false}})
+	if err != nil {
+		return err
+	}
+
+	_, err = as.RoleAdd(&pb.AuthRoleAddRequest{Name: "root"})
+	if err != nil {
+		return err
+	}
+
+	_, err = as.UserGrantRole(&pb.AuthUserGrantRoleRequest{User: "root", Role: "root"})
+	if err != nil {
+		return err
+	}
+
+	return as.AuthEnable()
+}
+
+func encodePassword(s string) string {
+	hashedPassword, _ := bcrypt.GenerateFromPassword([]byte(s), bcrypt.MinCost)
+	return base64.StdEncoding.EncodeToString(hashedPassword)
+}

server/lease/lease.go

@@ -39,6 +39,15 @@ type Lease struct {
 	revokec chan struct{}
 }
 
+func NewLease(id LeaseID, ttl int64) *Lease {
+	return &Lease{
+		ID:      id,
+		ttl:     ttl,
+		itemSet: make(map[LeaseItem]struct{}),
+		revokec: make(chan struct{}),
+	}
+}
+
 func (l *Lease) expired() bool {
 	return l.Remaining() <= 0
 }
@@ -56,6 +65,13 @@ func (l *Lease) TTL() int64 {
 	return l.ttl
 }
 
+// SetLeaseItem sets the given lease item, this func is thread-safe
+func (l *Lease) SetLeaseItem(item LeaseItem) {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+	l.itemSet[item] = struct{}{}
+}
+
 // getRemainingTTL returns the last checkpointed remaining TTL of the lease.
 func (l *Lease) getRemainingTTL() int64 {
 	if l.remainingTTL > 0 {

server/lease/lessor.go

@@ -279,12 +279,7 @@ func (le *lessor) Grant(id LeaseID, ttl int64) (*Lease, error) {
 
 	// TODO: when lessor is under high load, it should give out lease
 	// with longer TTL to reduce renew load.
-	l := &Lease{
-		ID:      id,
-		ttl:     ttl,
-		itemSet: make(map[LeaseItem]struct{}),
-		revokec: make(chan struct{}),
-	}
+	l := NewLease(id, ttl)
 
 	if l.ttl < le.minLeaseTTL {
 		l.ttl = le.minLeaseTTL