Repository: hashicorp/vault
File: vault/plugin_reload_test.go (Go)

Code before:
	"github.com/hashicorp/vault/sdk/logical"
)

type goodPluginTestBackend struct{}

func (g goodPluginTestBackend) Initialize(ctx context.Context, request *logical.InitializationRequest) error {
	return nil
}

func (g goodPluginTestBackend) HandleRequest(ctx context.Context, request *logical.Request) (*logical.Response, error) {
	return nil, nil
}
func (g goodPluginTestBackend) SpecialPaths() *logical.Paths { return nil }
func (g goodPluginTestBackend) System() logical.SystemView   { return nil }
func (g goodPluginTestBackend) Logger() hclog.Logger         { return nil }
func (g goodPluginTestBackend) HandleExistenceCheck(ctx context.Context, request *logical.Request) (bool, bool, error) {
	return false, false, nil
}
func (g goodPluginTestBackend) Cleanup(ctx context.Context)                 {}
func (g goodPluginTestBackend) InvalidateKey(ctx context.Context, s string) {}
func (g goodPluginTestBackend) Setup(ctx context.Context, config *logical.BackendConfig) error {
	return nil
}
func (g goodPluginTestBackend) Type() logical.BackendType { return logical.TypeLogical }

type badPluginTestBackend struct{}

func (b badPluginTestBackend) Initialize(ctx context.Context, r *logical.InitializationRequest) error {
	// we need the plugin to succeed the first initialize or it can't mount
	se, err := r.Storage.Get(ctx, "tag")
	if err != nil {
		return fmt.Errorf("couldn't read storage: %s", err)
	}
	if se == nil {
		r.Storage.Put(ctx, &logical.StorageEntry{
			Key:   "tag",
			Value: []byte("boo"),
		})
		return nil
	}

	return errors.New("already initialized")
}

func (b badPluginTestBackend) HandleRequest(_ context.Context, _ *logical.Request) (*logical.Response, error) {
	return nil, nil
}
func (b badPluginTestBackend) SpecialPaths() *logical.Paths { return nil }
func (b badPluginTestBackend) System() logical.SystemView   { return nil }
func (b badPluginTestBackend) Logger() hclog.Logger         { return nil }
func (b badPluginTestBackend) HandleExistenceCheck(ctx context.Context, request *logical.Request) (bool, bool, error) {

Review comment (suggestion):
This kind of thing kind of exists in `builtin/plugin/[v5]` and `sdk/plugin/mock` but I couldn't get it to work satisfactorily so i implemented specifically what i needed here.

Code after:
	"github.com/hashicorp/vault/sdk/helper/consts"
	"github.com/hashicorp/vault/sdk/logical"
)

const reloadFailurePrefix = "cannot reload"

// A goodReloadPluginTestBackend does nothing but always succeeds at Initialization (how good!)
//
// The name is long to avoid collisions in the vault package.
type goodReloadPluginTestBackend struct{}

func (g goodReloadPluginTestBackend) Initialize(ctx context.Context, request *logical.InitializationRequest) error {
	return nil
}

func (g goodReloadPluginTestBackend) HandleRequest(ctx context.Context, request *logical.Request) (*logical.Response, error) {
	return nil, nil
}
func (g goodReloadPluginTestBackend) SpecialPaths() *logical.Paths { return nil }
func (g goodReloadPluginTestBackend) System() logical.SystemView   { return nil }
func (g goodReloadPluginTestBackend) Logger() hclog.Logger         { return nil }
func (g goodReloadPluginTestBackend) HandleExistenceCheck(ctx context.Context, request *logical.Request) (bool, bool, error) {
	return false, false, nil
}
func (g goodReloadPluginTestBackend) Cleanup(ctx context.Context)                 {}
func (g goodReloadPluginTestBackend) InvalidateKey(ctx context.Context, s string) {}
func (g goodReloadPluginTestBackend) Setup(ctx context.Context, config *logical.BackendConfig) error {
	return nil
}
func (g goodReloadPluginTestBackend) Type() logical.BackendType { return logical.TypeLogical }

// An incrementallyBadReloadPluginTestBackend is a plugin that will fail after a specified number of Initialize calls.
// The first time it is initialized (at a given path), the check value is put in storage and persisted. When the plugin
// has been initialized that many times, the next call will fail. Note that every plugin needs to succeed at initialization
// once, or it will just fail to mount.
type incrementallyBadReloadPluginTestBackend struct {
	check byte
}

func (i incrementallyBadReloadPluginTestBackend) Initialize(ctx context.Context, r *logical.InitializationRequest) error {
	countEntry, err := r.Storage.Get(ctx, "count")
	if err != nil {
		return fmt.Errorf("failed to initialize for the wrong reason: %s", err)
	}
	if countEntry == nil {
		r.Storage.Put(ctx, &logical.StorageEntry{
			Key:   "count",
			Value: []byte{1},
		})
		// also write check to storage, because this is our real value
		r.Storage.Put(ctx, &logical.StorageEntry{