Merge branch 'develop' into jsimnz/fix/groupby-field-result-ignore

Author: jsimnz

internal/kms/enc_store.go

@@ -57,20 +57,26 @@ func (s *ipldEncStorage) get(ctx context.Context, cidBytes []byte) (*coreblock.E
 	return coreblock.GetEncryptionBlockFromNode(nd)
 }
 
-func (s *ipldEncStorage) put(ctx context.Context, blockBytes []byte) ([]byte, error) {
+func (s *ipldEncStorage) putBlock(ctx context.Context, block coreblock.Encryption) ([]byte, error) {
 	lsys := cidlink.DefaultLinkSystem()
 	lsys.SetWriteStorage(blockstore.NewIPLDStore(s.encstore))
 
-	var encBlock coreblock.Encryption
-	err := encBlock.Unmarshal(blockBytes)
+	link, err := lsys.Store(linking.LinkContext{Ctx: ctx}, coreblock.GetLinkPrototype(), block.GenerateNode())
 	if err != nil {
 		return nil, err
 	}
 
-	link, err := lsys.Store(linking.LinkContext{Ctx: ctx}, coreblock.GetLinkPrototype(), encBlock.GenerateNode())
+	return []byte(link.String()), nil
+}
+
+func (s *ipldEncStorage) computeBlockLink(ctx context.Context, block coreblock.Encryption) ([]byte, error) {
+	lsys := cidlink.DefaultLinkSystem()
+	lsys.SetWriteStorage(blockstore.NewIPLDStore(s.encstore))
+
+	link, err := lsys.ComputeLink(coreblock.GetLinkPrototype(), block.GenerateNode())
 	if err != nil {
 		return nil, err
 	}
 
-	return []byte(link.String()), nil
+	return []byte(link.Binary()), nil
 }

internal/kms/errors.go

@@ -17,13 +17,38 @@ import (
 const (
 	errUnknownKMSType              string = "unknown KMS type"
 	errNoPeerSuppliedEncryptionKey string = "no peer supplied the encryption key"
+
+	errPeerErrorKeyReply           string = "encryption key peer reply carried error"
+	errEncryptionKeyUnmarshal      string = "unmarshal encryption key response"
+	errReplyLinksAndBlocksMismatch string = "encryption key peer reply had mismatched links and blocks"
+	errDecryptEncryptionKey        string = "decrypting encryption key"
+	errDecodingEncryptionKey       string = "decoding encryption key"
+	errEncryptionKeyCIDMismatch    string = "key CID mismatch"
+	errEncryptionKeyStore          string = "store encryption key"
+	errKeyCIDGeneration            string = "key cid generation"
 )
 
 var (
 	ErrUnknownKMSType              = errors.New(errUnknownKMSType)
 	ErrNoPeerSuppliedEncryptionKey = errors.New(errNoPeerSuppliedEncryptionKey)
+
+	ErrPeerErrorKeyReply           = errors.New(errPeerErrorKeyReply)
+	ErrEncryptionKeyUnmarshal      = errors.New(errEncryptionKeyUnmarshal)
+	ErrReplyLinksAndBlocksMismatch = errors.New(errReplyLinksAndBlocksMismatch)
+	ErrDecryptEncryptionKey        = errors.New(errDecryptEncryptionKey)
+	ErrDecodingEncryptionKey       = errors.New(errDecodingEncryptionKey)
+	ErrEncryptionKeyCIDMismatch    = errors.New(errEncryptionKeyCIDMismatch)
+	ErrEncryptionKeyStore          = errors.New(errEncryptionKeyStore)
+	ErrKeyCIDGeneration            = errors.New(errKeyCIDGeneration)
 )
 
 func NewErrUnknownKMSType(t ServiceType) error {
 	return errors.New(errUnknownKMSType, errors.NewKV("Type", t))
 }
+
+func NewErrReplyLinksAndBlocksMismatch(linkCount, blockCount int) error {
+	return errors.New(errReplyLinksAndBlocksMismatch,
+		errors.NewKV("LinkCount", linkCount),
+		errors.NewKV("BlockCount", blockCount),
+	)
+}

internal/kms/pubsub.go

@@ -22,7 +22,6 @@ import (
 	cidlink "github.com/ipld/go-ipld-prime/linking/cid"
 	grpcpeer "google.golang.org/grpc/peer"
 
-	"github.com/sourcenetwork/corelog"
 	"github.com/sourcenetwork/immutable"
 
 	"github.com/sourcenetwork/defradb/acp/dac"
@@ -31,6 +30,7 @@ import (
 	"github.com/sourcenetwork/defradb/client"
 	"github.com/sourcenetwork/defradb/crypto"
 	"github.com/sourcenetwork/defradb/errors"
+	coreblock "github.com/sourcenetwork/defradb/internal/core/block"
 	"github.com/sourcenetwork/defradb/internal/datastore"
 	acpDB "github.com/sourcenetwork/defradb/internal/db/acp"
 	"github.com/sourcenetwork/defradb/internal/encryption"
@@ -290,16 +290,14 @@ func (s *pubSubService) handleFetchEncryptionKeyResponses(
 				respChan = nextRespChan
 				continue
 			}
-			items, ok := s.tryHandleFetchEncryptionKeyResponse(resp, req, privateKey)
+			items, ok, err := s.tryHandleFetchEncryptionKeyResponse(resp, req, privateKey, false /* skipVerify */)
 			if !ok {
+				if err != nil {
+					log.ErrorContextE(s.ctx, "Failed handling of encryption key response", err)
+				}
 				continue
 			}
 
-			// TODO: If multi-key requests become common, aggregate partial
-			// responses from multiple peers until every requested link is found or
-			// the timeout fires. A single valid response may contain fewer blocks
-			// than the request asked for.
-			// https://github.com/sourcenetwork/defradb/issues/4947
 			result <- encryption.Result{Items: items}
 			return
 
@@ -364,37 +362,35 @@ func (s *pubSubService) tryHandleFetchEncryptionKeyResponse(
 	resp client.PubsubResponse,
 	req *fetchEncryptionKeyRequest,
 	privateKey *ecdh.PrivateKey,
-) ([]encryption.Item, bool) {
+	skipVerify bool,
+) ([]encryption.Item, bool, error) {
 	if resp.Err != nil {
-		log.ErrorContextE(s.ctx, "encryption key peer reply carried error", resp.Err)
-		return nil, false
+		return nil, false, errors.Join(ErrPeerErrorKeyReply, resp.Err)
 	}
 
 	var keyResp fetchEncryptionKeyReply
 	if err := cbor.Unmarshal(resp.Data, &keyResp); err != nil {
-		log.ErrorContextE(s.ctx, "Failed to unmarshal encryption key response", err)
-		return nil, false
+		return nil, false, errors.Join(ErrEncryptionKeyUnmarshal, err)
 	}
 
 	if len(keyResp.Blocks) == 0 {
 		// Peer didn't have the key; keep waiting for the one that does.
-		return nil, false
+		return nil, false, nil
 	}
 	if len(keyResp.Links) != len(keyResp.Blocks) {
-		log.ErrorContext(
-			s.ctx,
-			"encryption key peer reply had mismatched links and blocks",
-			corelog.Int("LinkCount", len(keyResp.Links)),
-			corelog.Int("BlockCount", len(keyResp.Blocks)),
-		)
-		return nil, false
+		return nil, false, NewErrReplyLinksAndBlocksMismatch(len(keyResp.Links), len(keyResp.Blocks))
 	}
 
 	senderID := keyResp.Sender
 	if senderID == "" {
 		senderID = resp.From
 	}
 
+	reqSet := make(map[string]struct{}, len(req.Links))
+	for _, l := range req.Links {
+		reqSet[string(l)] = struct{}{}
+	}
+
 	resultEncItems := make([]encryption.Item, 0, len(keyResp.Blocks))
 	for i, block := range keyResp.Blocks {
 		decryptedData, err := crypto.DecryptECIES(
@@ -405,25 +401,41 @@ func (s *pubSubService) tryHandleFetchEncryptionKeyResponse(
 			crypto.WithPubKeyPrepended(false),
 		)
 		if err != nil {
-			log.ErrorContextE(s.ctx, "Failed to decrypt encryption key", err)
-			return nil, false
+			return nil, false, errors.Join(ErrDecryptEncryptionKey, err)
+		}
+
+		var encBlock coreblock.Encryption
+		err = encBlock.Unmarshal(decryptedData)
+		if err != nil {
+			return nil, false, errors.Join(ErrDecodingEncryptionKey, err)
+		}
+
+		if !skipVerify {
+			link, err := s.encStore.computeBlockLink(s.ctx, encBlock)
+			if err != nil {
+				return nil, false, errors.Join(ErrKeyCIDGeneration, err)
+			}
+
+			if !bytes.Equal(keyResp.Links[i], link) {
+				return nil, false, ErrEncryptionKeyCIDMismatch
+			}
+
+			if _, ok := reqSet[string(link)]; !ok {
+				return nil, false, ErrEncryptionKeyCIDMismatch
+			}
 		}
 
-		if _, err := s.encStore.put(context.Background(), decryptedData); err != nil {
-			log.ErrorContextE(s.ctx, "Failed to store encryption key", err)
-			return nil, false
+		if _, err := s.encStore.putBlock(context.Background(), encBlock); err != nil {
+			return nil, false, errors.Join(ErrEncryptionKeyStore, err)
 		}
 
-		// todo: verify incoming response order block/CIDs match the request
-		// current implementation assumes trusted response ordering
-		// https://github.com/sourcenetwork/defradb/issues/4948
 		resultEncItems = append(resultEncItems, encryption.Item{
 			Link:  keyResp.Links[i],
 			Block: decryptedData,
 		})
 	}
 
-	return resultEncItems, true
+	return resultEncItems, true, nil
 }
 
 // makeAssociatedData creates the associated data for the encryption key request