Skip to content

Commit

Permalink
Merge pull request #13162 from serathius/auth
Browse files Browse the repository at this point in the history
etcdserver: Move read/update methods on Auth bucket to one place
  • Loading branch information
ptabor authored Jul 3, 2021
2 parents a8e6016 + b2e08fb commit 33b2cdb
Show file tree
Hide file tree
Showing 3 changed files with 173 additions and 29 deletions.
36 changes: 7 additions & 29 deletions server/auth/store.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@ import (
"bytes"
"context"
"encoding/base64"
"encoding/binary"
"errors"
"sort"
"strings"
Expand Down Expand Up @@ -76,8 +75,6 @@ const (

tokenTypeSimple = "simple"
tokenTypeJWT = "jwt"

revBytesLen = 8
)

type AuthInfo struct {
Expand Down Expand Up @@ -237,7 +234,7 @@ func (as *authStore) AuthEnable() error {
return ErrRootRoleNotExist
}

tx.UnsafePut(buckets.Auth, buckets.AuthEnabledKeyName, authEnabled)
buckets.UnsafeSaveAuthEnabled(tx, true)

as.enabled = true
as.tokenProvider.enable()
Expand All @@ -259,7 +256,7 @@ func (as *authStore) AuthDisable() {
b := as.be
tx := b.BatchTx()
tx.Lock()
tx.UnsafePut(buckets.Auth, buckets.AuthEnabledKeyName, authDisabled)
buckets.UnsafeSaveAuthEnabled(tx, false)
as.commitRevision(tx)
tx.Unlock()
b.ForceCommit()
Expand Down Expand Up @@ -350,17 +347,11 @@ func (as *authStore) CheckPassword(username, password string) (uint64, error) {
}

func (as *authStore) Recover(be backend.Backend) {
enabled := false
as.be = be
tx := be.BatchTx()
tx.Lock()
_, vs := tx.UnsafeRange(buckets.Auth, buckets.AuthEnabledKeyName, nil, 0)
if len(vs) == 1 {
if bytes.Equal(vs[0], authEnabled) {
enabled = true
}
}

enabled := buckets.UnsafeReadAuthEnabled(tx)
as.setRevision(getRevision(tx))

tx.Unlock()
Expand Down Expand Up @@ -936,17 +927,11 @@ func NewAuthStore(lg *zap.Logger, be backend.Backend, tp TokenProvider, bcryptCo
tx := be.BatchTx()
tx.Lock()

tx.UnsafeCreateBucket(buckets.Auth)
buckets.UnsafeCreateAuthBucket(tx)
tx.UnsafeCreateBucket(buckets.AuthUsers)
tx.UnsafeCreateBucket(buckets.AuthRoles)

enabled := false
_, vs := tx.UnsafeRange(buckets.Auth, buckets.AuthEnabledKeyName, nil, 0)
if len(vs) == 1 {
if bytes.Equal(vs[0], authEnabled) {
enabled = true
}
}
enabled := buckets.UnsafeReadAuthEnabled(tx)

as := &authStore{
revision: getRevision(tx),
Expand Down Expand Up @@ -982,18 +967,11 @@ func hasRootRole(u *authpb.User) bool {

func (as *authStore) commitRevision(tx backend.BatchTx) {
atomic.AddUint64(&as.revision, 1)
revBytes := make([]byte, revBytesLen)
binary.BigEndian.PutUint64(revBytes, as.Revision())
tx.UnsafePut(buckets.Auth, buckets.AuthRevisionKeyName, revBytes)
buckets.UnsafeSaveAuthRevision(tx, as.Revision())
}

func getRevision(tx backend.BatchTx) uint64 {
_, vs := tx.UnsafeRange(buckets.Auth, buckets.AuthRevisionKeyName, nil, 0)
if len(vs) != 1 {
// this can happen in the initialization phase
return 0
}
return binary.BigEndian.Uint64(vs[0])
return buckets.UnsafeReadAuthRevision(tx)
}

func (as *authStore) setRevision(rev uint64) {
Expand Down
67 changes: 67 additions & 0 deletions server/mvcc/buckets/auth.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,67 @@
// Copyright 2021 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 buckets

import (
"bytes"
"encoding/binary"
"go.etcd.io/etcd/server/v3/mvcc/backend"
)

const (
revBytesLen = 8
)

var (
authEnabled = []byte{1}
authDisabled = []byte{0}
)

func UnsafeCreateAuthBucket(tx backend.BatchTx) {
tx.UnsafeCreateBucket(Auth)
}

func UnsafeSaveAuthEnabled(tx backend.BatchTx, enabled bool) {
if enabled {
tx.UnsafePut(Auth, AuthEnabledKeyName, authEnabled)
} else {
tx.UnsafePut(Auth, AuthEnabledKeyName, authDisabled)
}
}

func UnsafeReadAuthEnabled(tx backend.ReadTx) bool {
_, vs := tx.UnsafeRange(Auth, AuthEnabledKeyName, nil, 0)
if len(vs) == 1 {
if bytes.Equal(vs[0], authEnabled) {
return true
}
}
return false
}

func UnsafeSaveAuthRevision(tx backend.BatchTx, rev uint64) {
revBytes := make([]byte, revBytesLen)
binary.BigEndian.PutUint64(revBytes, rev)
tx.UnsafePut(Auth, AuthRevisionKeyName, revBytes)
}

func UnsafeReadAuthRevision(tx backend.ReadTx) uint64 {
_, vs := tx.UnsafeRange(Auth, AuthRevisionKeyName, nil, 0)
if len(vs) != 1 {
// this can happen in the initialization phase
return 0
}
return binary.BigEndian.Uint64(vs[0])
}
99 changes: 99 additions & 0 deletions server/mvcc/buckets/auth_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,99 @@
// Copyright 2021 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 buckets

import (
"fmt"
"math"
"testing"
"time"

"github.com/stretchr/testify/assert"
"go.etcd.io/etcd/server/v3/mvcc/backend"
betesting "go.etcd.io/etcd/server/v3/mvcc/backend/testing"
)

// TestAuthEnabled ensures that UnsafeSaveAuthEnabled&UnsafeReadAuthEnabled work well together.
func TestAuthEnabled(t *testing.T) {
tcs := []struct {
enabled bool
}{
{
enabled: true,
},
{
enabled: false,
},
}
for _, tc := range tcs {
t.Run(fmt.Sprint(tc.enabled), func(t *testing.T) {
be, tmpPath := betesting.NewTmpBackend(t, time.Microsecond, 10)
tx := be.BatchTx()
if tx == nil {
t.Fatal("batch tx is nil")
}
tx.Lock()
UnsafeCreateAuthBucket(tx)
UnsafeSaveAuthEnabled(tx, tc.enabled)
tx.Unlock()
be.ForceCommit()
be.Close()

b := backend.NewDefaultBackend(tmpPath)
defer b.Close()
v := UnsafeReadAuthEnabled(b.BatchTx())

assert.Equal(t, tc.enabled, v)
})
}
}

// TestAuthRevision ensures that UnsafeSaveAuthRevision&UnsafeReadAuthRevision work well together.
func TestAuthRevision(t *testing.T) {
tcs := []struct {
revision uint64
}{
{
revision: 0,
},
{
revision: 1,
},
{
revision: math.MaxUint64,
},
}
for _, tc := range tcs {
t.Run(fmt.Sprint(tc.revision), func(t *testing.T) {
be, tmpPath := betesting.NewTmpBackend(t, time.Microsecond, 10)
tx := be.BatchTx()
if tx == nil {
t.Fatal("batch tx is nil")
}
tx.Lock()
UnsafeCreateAuthBucket(tx)
UnsafeSaveAuthRevision(tx, tc.revision)
tx.Unlock()
be.ForceCommit()
be.Close()

b := backend.NewDefaultBackend(tmpPath)
defer b.Close()
v := UnsafeReadAuthRevision(b.BatchTx())

assert.Equal(t, tc.revision, v)
})
}
}

0 comments on commit 33b2cdb

Please sign in to comment.