From 4e39f690f2c093962271d6cabae724733266ac9a Mon Sep 17 00:00:00 2001
From: Hitoshi Mitake <mitake.hitoshi@gmail.com>
Date: Mon, 21 Mar 2016 22:43:41 +0900
Subject: [PATCH] auth, etcdserver: add a method for recoverying from backend
 during apply snapshot

This commit adds a new method Recovery() to auth.AuthStore for
recoverying auth state from backend during apply snapshot. It follows
a manner of the lessor.
---
 auth/store.go        | 25 ++++++++++++++-----------
 etcdserver/server.go |  6 +++++-
 2 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/auth/store.go b/auth/store.go
index 591cc29f4..3e181eb54 100644
--- a/auth/store.go
+++ b/auth/store.go
@@ -19,10 +19,6 @@ import (
 	"github.com/coreos/etcd/storage/backend"
 )
 
-type backendGetter interface {
-	Backend() backend.Backend
-}
-
 var (
 	enableFlagKey  = []byte("authEnabled")
 	authBucketName = []byte("auth")
@@ -33,16 +29,19 @@ var (
 type AuthStore interface {
 	// AuthEnable() turns on the authentication feature
 	AuthEnable()
+
+	// Recover recovers the state of auth store from the given backend
+	Recover(b backend.Backend)
 }
 
 type authStore struct {
-	bgetter backendGetter
+	be backend.Backend
 }
 
 func (as *authStore) AuthEnable() {
 	value := []byte{1}
 
-	b := as.bgetter.Backend()
+	b := as.be
 	tx := b.BatchTx()
 	tx.Lock()
 	tx.UnsafePut(authBucketName, enableFlagKey, value)
@@ -52,15 +51,19 @@ func (as *authStore) AuthEnable() {
 	plog.Noticef("Authentication enabled")
 }
 
-func NewAuthStore(bgetter backendGetter) *authStore {
-	b := bgetter.Backend()
-	tx := b.BatchTx()
+func (as *authStore) Recover(be backend.Backend) {
+	as.be = be
+	// TODO(mitake): recovery process
+}
+
+func NewAuthStore(be backend.Backend) *authStore {
+	tx := be.BatchTx()
 	tx.Lock()
 	tx.UnsafeCreateBucket(authBucketName)
 	tx.Unlock()
-	b.ForceCommit()
+	be.ForceCommit()
 
 	return &authStore{
-		bgetter: bgetter,
+		be: be,
 	}
 }
diff --git a/etcdserver/server.go b/etcdserver/server.go
index fb8d75062..5299a0a15 100644
--- a/etcdserver/server.go
+++ b/etcdserver/server.go
@@ -374,7 +374,7 @@ func NewServer(cfg *ServerConfig) (*EtcdServer, error) {
 		srv.be = backend.NewDefaultBackend(path.Join(cfg.SnapDir(), databaseFilename))
 		srv.lessor = lease.NewLessor(srv.be)
 		srv.kv = dstorage.New(srv.be, srv.lessor, &srv.consistIndex)
-		srv.authStore = auth.NewAuthStore(srv)
+		srv.authStore = auth.NewAuthStore(srv.be)
 		if h := cfg.AutoCompactionRetention; h != 0 {
 			srv.compactor = compactor.NewPeriodic(h, srv.kv, srv)
 			srv.compactor.Run()
@@ -621,6 +621,10 @@ func (s *EtcdServer) applySnapshot(ep *etcdProgress, apply *apply) {
 		if s.lessor != nil {
 			s.lessor.Recover(newbe, s.kv)
 		}
+
+		if s.authStore != nil {
+			s.authStore.Recover(newbe)
+		}
 	}
 	if err := s.store.Recovery(apply.snapshot.Data); err != nil {
 		plog.Panicf("recovery store error: %v", err)