MB-62186: Introduce key management for encryption-at-rest

* Add support for receiving encryption-at-rest key notifications from
cbauth

* cbauth will push key information for each "key data type". Query will
receive key information for the  "logs", "other", and per-bucket keys.

* Implement a node-level key manager that maintains a cache/store of
key information for fast access

* Update the key store only when received key info from cbauth differs
from the cached info for a given data type.

* On Query service startup, "prime" the key store with key info for
all known data types ("logs", "other", and all existing buckets)
before starting listeners.

* Priming means proactively loading the key store with key info
from cbauth, so that manager's key stor eis ready before any Query
requests are handled.

* Requests to the key store for unloaded key types can trigger calls to
cbauth to fetch key information.

To prevent this potential flood of fetch requests, prime the key store
on startup.

* During priming, Query waits up to 5 seconds for each key fetch
from cbauth to avoid indefinite blocking. If fetching fails or times
out, priming for that type is skipped.

* Priming failures will not result in the Query service startup failing
and exiting. Or that key info is permanently missing from Query's key
manager.

Query service can function and serve requests even when priming fails.
This is because, missing key info entries are loaded into the key store
by future refresh callbacks or by lazy loading the key information when
the key information is requested from the key store for the first time.

* Key management is only implemented for EE version

Note:

* For now, The code to create the encryption manager and register
the encryption callbacks with cbauth is currently commented out.

And instead, dummy callbacks are registered with cbauth.

This is because currently, ns-server does not push key information
to the Query service. And registering the actual callbacks will result
in errors.

Change-Id: Ib1bf0664272f2ef3a502917023fc2b7fcbfd3ce8
Reviewed-on: https://review.couchbase.org/c/query/+/242165
Reviewed-by: Sitaram Vemulapalli <sitaram.vemulapalli@couchbase.com>
Reviewed-by: Bingjie Miao <bingjie.miao@couchbase.com>
Tested-by: Dhanya Gowrish <dhanya.gowrish@couchbase.com>

Author: dhanyagowrish

datastore/couchbase/couchbase.go

@@ -162,15 +162,16 @@ func SetDeploymentModel(deploymentModel string) {
 
 // store is the root for the couchbase datastore
 type store struct {
-	client         cb.Client // instance of primitives/couchbase client
-	gcClient       *gcagent.Client
-	namespaceCache map[string]*namespace // map of pool-names and IDs
-	CbAuthInit     bool                  // whether cbAuth is initialized
-	inferencer     datastore.Inferencer  // what we use to infer schemas
-	statUpdater    datastore.StatUpdater // what we use to update statistics
-	connectionUrl  string                // where to contact ns_server
-	connSecConfig  *datastore.ConnectionSecurityConfig
-	nslock         sync.RWMutex
+	client             cb.Client // instance of primitives/couchbase client
+	gcClient           *gcagent.Client
+	namespaceCache     map[string]*namespace // map of pool-names and IDs
+	CbAuthInit         bool                  // whether cbAuth is initialized
+	inferencer         datastore.Inferencer  // what we use to infer schemas
+	statUpdater        datastore.StatUpdater // what we use to update statistics
+	connectionUrl      string                // where to contact ns_server
+	connSecConfig      *datastore.ConnectionSecurityConfig
+	nslock             sync.RWMutex
+	encryptionProvider datastore.EncryptionProvider
 }
 
 func (s *store) Id() string {
@@ -1072,6 +1073,8 @@ func NewDatastore(u string) (s datastore.Datastore, e errors.Error) {
 		return nil, er
 	}
 
+	store.encryptionProvider = datastore.NoopEncryptionProviderInstance
+
 	// initialize the default pool.
 	// TODO can couchbase server contain more than one pool ?
 
@@ -3908,3 +3911,11 @@ func (s *store) CheckSystemCollection(bucketName, requestId string, forceIndex b
 
 	return empty, nil
 }
+
+func (s *store) EncryptionProvider() (datastore.EncryptionProvider, errors.Error) {
+	return s.encryptionProvider, nil
+}
+
+func (s *store) SetEncryptionProvider(encProvider datastore.EncryptionProvider) {
+	s.encryptionProvider = encProvider
+}

datastore/datastore.go

@@ -110,6 +110,9 @@ type Datastore interface {
 	RollbackTransaction(stmtAtomicity bool, context QueryContext, sname string) errors.Error
 	SetSavepoint(stmtAtomicity bool, context QueryContext, sname string) errors.Error
 	TransactionDeltaKeyScan(keyspace string, conn *IndexConnection) // Keys of Delta keyspace
+
+	EncryptionProvider() (EncryptionProvider, errors.Error)
+	SetEncryptionProvider(encryptionProvider EncryptionProvider)
 }
 
 type Systemstore interface {

datastore/encryption_provider.go

@@ -0,0 +1,26 @@
+//  Copyright 2026-Present Couchbase, Inc.
+//
+//  Use of this software is governed by the Business Source License included
+//  in the file licenses/BSL-Couchbase.txt.  As of the Change Date specified
+//  in that file, in accordance with the Business Source License, use of this
+//  software will be governed by the Apache License, Version 2.0, included in
+//  the file licenses/APL2.txt.
+
+package datastore
+
+import (
+	"github.com/couchbase/query/encryption"
+	"github.com/couchbase/query/errors"
+)
+
+var NoopEncryptionProviderInstance = &NoopEncryptionProvider{}
+
+type EncryptionProvider interface {
+	GetActiveKey(dt encryption.KeyDataType) (*encryption.EaRKey, errors.Error)
+}
+
+type NoopEncryptionProvider struct{}
+
+func (NoopEncryptionProvider) GetActiveKey(dt encryption.KeyDataType) (*encryption.EaRKey, errors.Error) {
+	return nil, nil
+}

datastore/file/file.go

@@ -281,6 +281,14 @@ func (s *store) TransactionDeltaKeyScan(keyspace string, conn *datastore.IndexCo
 	defer conn.Sender().Close()
 }
 
+func (s *store) EncryptionProvider() (datastore.EncryptionProvider, errors.Error) {
+	return datastore.NoopEncryptionProviderInstance, nil
+}
+
+func (s *store) SetEncryptionProvider(datastore.EncryptionProvider) {
+	return
+}
+
 // NewStore creates a new file-based store for the given filepath.
 func NewDatastore(path string) (s datastore.Datastore, e errors.Error) {
 	path, er := filepath.Abs(path)

datastore/mock/mock.go

@@ -267,6 +267,14 @@ func (s *store) TransactionDeltaKeyScan(keyspace string, conn *datastore.IndexCo
 	defer conn.Sender().Close()
 }
 
+func (s *store) EncryptionProvider() (datastore.EncryptionProvider, errors.Error) {
+	return datastore.NoopEncryptionProviderInstance, nil
+}
+
+func (s *store) SetEncryptionProvider(datastore.EncryptionProvider) {
+	return
+}
+
 // namespace represents a mock-based Namespace.
 type namespace struct {
 	store         *store

datastore/system/system.go

@@ -381,6 +381,14 @@ func (s *store) TransactionDeltaKeyScan(keyspace string, conn *datastore.IndexCo
 	defer conn.Sender().Close()
 }
 
+func (s *store) EncryptionProvider() (datastore.EncryptionProvider, errors.Error) {
+	return s.actualStore.EncryptionProvider()
+}
+
+func (s *store) SetEncryptionProvider(datastore.EncryptionProvider) {
+	return
+}
+
 func NewDatastore(actualStore datastore.Datastore, acctStore accounting.AccountingStore,
 	enterprise bool) (datastore.Systemstore, errors.Error) {
 	s := &store{actualStore: actualStore, acctStore: acctStore, enterprise: enterprise}

encryption/keymgmt/encryption_mgmt_ce.go

@@ -0,0 +1,53 @@
+//  Copyright 2026-Present Couchbase, Inc.
+//
+//  Use of this software is governed by the Business Source License included
+//  in the file licenses/BSL-Couchbase.txt.  As of the Change Date specified
+//  in that file, in accordance with the Business Source License, use of this
+//  software will be governed by the Apache License, Version 2.0, included in
+//  the file licenses/APL2.txt.
+
+//go:build !enterprise
+
+package keymgmt
+
+import (
+	"github.com/couchbase/cbauth"
+	"github.com/couchbase/query/encryption"
+	"github.com/couchbase/query/errors"
+)
+
+func NewEncryptionManager() EncryptionManager {
+	return &NoopEncryptionManager{}
+}
+
+type NoopEncryptionManager struct{}
+
+func (this *NoopEncryptionManager) GetActiveKey(dt encryption.KeyDataType) (*encryption.EaRKey, errors.Error) {
+	return nil, nil
+}
+
+func (this *NoopEncryptionManager) PrimeKeys(keyDataTypes []encryption.KeyDataType) errors.Error {
+	return nil
+}
+
+func (this *NoopEncryptionManager) UpdateKeys(dataType cbauth.KeyDataType, newInfo *cbauth.EncrKeysInfo, prime bool) errors.Error {
+	return nil
+}
+
+func (this *NoopEncryptionManager) RegisterCbauthEncryptionCallbacks() {
+}
+
+func (this *NoopEncryptionManager) GetInUseKeysCallback(dt cbauth.KeyDataType) ([]string, error) {
+	return nil, nil
+}
+
+func (this *NoopEncryptionManager) DropKeysCallback(dt cbauth.KeyDataType, KeyIdsToDrop []string) {
+}
+
+func (this *NoopEncryptionManager) SynchronizeKeyFilesCallback(dt cbauth.KeyDataType) error {
+	return nil
+}
+
+func (this *NoopEncryptionManager) RefreshKeysCallback(dt cbauth.KeyDataType) error {
+	return nil
+}

encryption/keymgmt/encryption_mgmt_ee.go

@@ -0,0 +1,15 @@
+//  Copyright 2026-Present Couchbase, Inc.
+//
+//  Use of this software is governed by the Business Source License included
+//  in the file licenses/BSL-Couchbase.txt.  As of the Change Date specified
+//  in that file, in accordance with the Business Source License, use of this
+//  software will be governed by the Apache License, Version 2.0, included in
+//  the file licenses/APL2.txt.
+
+//go:build enterprise
+
+package keymgmt
+
+func NewEncryptionManager() EncryptionManager {
+	return NewNodeEncryptionManager()
+}

encryption/keymgmt/mgmt.go

@@ -0,0 +1,26 @@
+//  Copyright 2026-Present Couchbase, Inc.
+//
+//  Use of this software is governed by the Business Source License included
+//  in the file licenses/BSL-Couchbase.txt.  As of the Change Date specified
+//  in that file, in accordance with the Business Source License, use of this
+//  software will be governed by the Apache License, Version 2.0, included in
+//  the file licenses/APL2.txt.
+
+package keymgmt
+
+import (
+	"github.com/couchbase/cbauth"
+	"github.com/couchbase/query/encryption"
+	"github.com/couchbase/query/errors"
+)
+
+type EncryptionManager interface {
+	GetActiveKey(dt encryption.KeyDataType) (*encryption.EaRKey, errors.Error)
+	PrimeKeys(keyDataTypes []encryption.KeyDataType) errors.Error
+	UpdateKeys(dataType cbauth.KeyDataType, newInfo *cbauth.EncrKeysInfo, prime bool) errors.Error
+	RegisterCbauthEncryptionCallbacks()
+	GetInUseKeysCallback(dt cbauth.KeyDataType) ([]string, error)
+	DropKeysCallback(dt cbauth.KeyDataType, KeyIdsToDrop []string)
+	SynchronizeKeyFilesCallback(dt cbauth.KeyDataType) error
+	RefreshKeysCallback(dt cbauth.KeyDataType) error
+}

encryption/keymgmt/node_encryption_mgr.go

@@ -0,0 +1,135 @@
+//  Copyright 2026-Present Couchbase, Inc.
+//
+//  Use of this software is governed by the Business Source License included
+//  in the file licenses/BSL-Couchbase.txt.  As of the Change Date specified
+//  in that file, in accordance with the Business Source License, use of this
+//  software will be governed by the Apache License, Version 2.0, included in
+//  the file licenses/APL2.txt.
+
+//go:build enterprise
+
+package keymgmt
+
+import (
+	"fmt"
+
+	"github.com/couchbase/cbauth"
+	"github.com/couchbase/query/encryption"
+	"github.com/couchbase/query/errors"
+	"github.com/couchbase/query/logging"
+)
+
+type NodeEncryptionManager struct {
+	keyStore *nodeKeyStore
+}
+
+func NewNodeEncryptionManager() *NodeEncryptionManager {
+	return &NodeEncryptionManager{
+		keyStore: newNodeKeyStore(),
+	}
+}
+
+// Performs initialization of the manager with info about the provided key datatypes
+// Typically invoked during service startup to initialize the manager with required key data.
+func (this *NodeEncryptionManager) PrimeKeys(keyDataTypes []encryption.KeyDataType) errors.Error {
+	logging.Infof("Priming encryption-at-rest manager with keys for all available key data types")
+	this.keyStore.PrimeKeys(keyDataTypes)
+	logging.Infof("Finished priming operation of encryption-at-rest manager with keys for all available key data types")
+	return nil
+}
+
+func (this *NodeEncryptionManager) UpdateKeys(dataType cbauth.KeyDataType, newInfo *cbauth.EncrKeysInfo, prime bool) errors.Error {
+	return this.keyStore.UpdateKeys(dataType, newInfo, prime)
+}
+
+func (this *NodeEncryptionManager) GetActiveKey(dt encryption.KeyDataType) (*encryption.EaRKey, errors.Error) {
+	return this.keyStore.GetActiveKey(dt)
+}
+
+func (this *NodeEncryptionManager) RegisterCbauthEncryptionCallbacks() {
+	cbauth.RegisterEncryptionKeysCallbacks(this.RefreshKeysCallback, this.GetInUseKeysCallback, this.DropKeysCallback,
+		this.SynchronizeKeyFilesCallback)
+}
+
+func (this *NodeEncryptionManager) RefreshKeysCallback(dt cbauth.KeyDataType) error {
+	// Key info will always be present in cbauth when RefreshKeysCallback is called.
+	// Thus call cbauth.GetEncryptionKeys() instead of GetEncryptionKeysBlocking().
+	newKeys, cbErr := cbauth.GetEncryptionKeys(dt)
+
+	// In RefreshKeysCallback, any error returned by GetEncryptionKeys() is a hard error including cbauth.ErrKeysNotAvailable.
+	// This is because key info should always be available when RefreshKeysCallback is called by cbauth.
+	if cbErr != nil {
+		logging.Errorf(
+			"Error refreshing encryption-at-rest configuration for key data type %s. Failed to fetch configuration from cbauth: %v",
+			cbauthTypeToDataType(dt).String(), cbErr)
+		return cbErr
+	}
+
+	err := this.keyStore.UpdateKeys(dt, newKeys, false)
+	if err != nil && err.Code() != errors.E_INVALID_ENCRYPTION_KEY_DATATYPE {
+		logging.Errorf(
+			"Error refreshing encryption-at-rest configuration for key data type %s. Failed to update local encryption manager: %v",
+			cbauthTypeToDataType(dt).String(), err)
+	}
+
+	return nil
+}
+
+func (this *NodeEncryptionManager) GetInUseKeysCallback(dt cbauth.KeyDataType) ([]string, error) {
+	// EAR TODO
+	return nil, nil
+}
+
+func (this *NodeEncryptionManager) DropKeysCallback(dt cbauth.KeyDataType, KeyIdsToDrop []string) {
+	// EAR TODO
+}
+
+func (this *NodeEncryptionManager) SynchronizeKeyFilesCallback(dt cbauth.KeyDataType) error {
+	// Query has no requirement for this as of now
+	return nil
+}
+
+func validateKeyDataType(dt cbauth.KeyDataType) (encryption.KeyDataType, errors.Error) {
+	kdt := encryption.KeyDataType{
+		TypeName:   dt.TypeName,
+		BucketUUID: dt.BucketUUID,
+	}
+
+	if dt.BucketUUID != "" && dt.TypeName != encryption.BUCKET_KEY_DATATYPE {
+		return encryption.KeyDataType{}, errors.NewEncryptionError(errors.E_INVALID_ENCRYPTION_KEY_DATATYPE,
+			fmt.Errorf("bucketUUID is only valid when typeName is service_bucket"), kdt.String())
+	}
+
+	if dt.TypeName == encryption.BUCKET_KEY_DATATYPE && dt.BucketUUID == "" {
+		return encryption.KeyDataType{}, errors.NewEncryptionError(errors.E_INVALID_ENCRYPTION_KEY_DATATYPE,
+			fmt.Errorf("bucketUUID is required when typeName is service_bucket"), kdt.String())
+	}
+
+	switch dt.TypeName {
+	case encryption.BUCKET_KEY_DATATYPE:
+		if dt.BucketUUID == "" {
+			return encryption.KeyDataType{}, errors.NewEncryptionError(errors.E_INVALID_ENCRYPTION_KEY_DATATYPE,
+				fmt.Errorf("bucketUUID is required when typeName is service_bucket"), kdt.String())
+		}
+	case encryption.LOG_KEY_DATATYPE, encryption.OTHER_KEY_DATATYPE:
+	default:
+		return encryption.KeyDataType{}, errors.NewEncryptionError(errors.E_INVALID_ENCRYPTION_KEY_DATATYPE,
+			fmt.Errorf("unsupported key data type"), kdt.String())
+	}
+
+	return kdt, nil
+}
+
+func dataTypeToCbauthType(dt encryption.KeyDataType) cbauth.KeyDataType {
+	return cbauth.KeyDataType{
+		TypeName:   dt.TypeName,
+		BucketUUID: dt.BucketUUID,
+	}
+}
+
+func cbauthTypeToDataType(dt cbauth.KeyDataType) encryption.KeyDataType {
+	return encryption.KeyDataType{
+		TypeName:   dt.TypeName,
+		BucketUUID: dt.BucketUUID,
+	}
+}