feat(sdk): Add obligations support. (#2759)

### Proposed Changes

1.) Expose `Obligations` function to TDF/NanoTDF
2.) Expose new `NanoTDFReader` with `LoadNanoTDF`,`Init` functions
3.) Refactor `bulk` to allow consumers to retrieve Obligations before
decrypting content
4.) Expose `PrepareBulkDecrypt` to retrieve Obligations before
decrypting the content
5.) Expose `ErrRewrapForbidden` for consumption by downstream consumers
of the SDK.

Manual E2E tests:
- Ran benchmark and rt tests for bulk
- Nano/ZTDF no access to resource regardless of obligations (returns
empty slice of FQNs)
- Nano/ZTDF has access to resource, but does not fulfill obligations
(returns required obligations)
- Nano/ZTDF has access to resource and fulfills obligations -> Decrypt
successfully.
- Multi-KAS decrypt successful when able to fulfill obligations, failure
with returned obligations when not

**Caveats**
>[!NOTE]
>Bulk decrypt is designed in a way that will only allow
>a global set of fulfillable obligations to be applied, not for
>an individual TDF.



### Checklist

- [ ] I have added or updated unit tests
- [ ] I have added or updated integration tests (if appropriate)
- [ ] I have added or updated documentation

### Testing Instructions

Author: c-r33d

sdk/bulk.go

@@ -12,9 +12,10 @@ import (
 
 // BulkTDF: Reader is TDF Content. Writer writes encrypted data. Error is the error that occurs if decrypting fails.
 type BulkTDF struct {
-	Reader io.ReadSeeker
-	Writer io.Writer
-	Error  error
+	Reader               io.ReadSeeker
+	Writer               io.Writer
+	Error                error
+	TriggeredObligations Obligations
 }
 
 type BulkDecryptRequest struct {
@@ -26,6 +27,15 @@ type BulkDecryptRequest struct {
 	ignoreAllowList       bool
 }
 
+// BulkDecryptPrepared holds the prepared state for bulk decryption
+// The PolicyTDF is a map of created policy IDs to their corresponding BulkTDF
+// The policy IDs are generated during the prepareDecryptors function
+type BulkDecryptPrepared struct {
+	PolicyTDF     map[string]*BulkTDF
+	tdfDecryptors map[string]decryptor
+	allRewrapResp map[string][]kaoResult
+}
+
 // BulkErrors List of Errors that Failed during Bulk Decryption
 type BulkErrors []error
 
@@ -116,17 +126,9 @@ func (s SDK) createDecryptor(tdf *BulkTDF, req *BulkDecryptRequest) (decryptor,
 	return nil, fmt.Errorf("unknown tdf type: %s", req.TDFType)
 }
 
-// BulkDecrypt Decrypts a list of BulkTDF and if a partial failure of TDFs unable to be decrypted, BulkErrors would be returned.
-func (s SDK) BulkDecrypt(ctx context.Context, opts ...BulkDecryptOption) error {
-	bulkReq, createError := createBulkRewrapRequest(opts...)
-	if createError != nil {
-		return fmt.Errorf("failed to create bulk rewrap request: %w", createError)
-	}
-	kasRewrapRequests := make(map[string][]*kas.UnsignedRewrapRequest_WithPolicyRequest)
-	tdfDecryptors := make(map[string]decryptor)
-	policyTDF := make(map[string]*BulkTDF)
-
-	if !bulkReq.ignoreAllowList && len(bulkReq.kasAllowlist) == 0 { //nolint:nestif // if kasAllowlist is not set, we get it from the registry
+// setupKasAllowlist configures the KAS allowlist for the bulk request
+func (s SDK) setupKasAllowlist(ctx context.Context, bulkReq *BulkDecryptRequest) error {
+	if !bulkReq.ignoreAllowList && len(bulkReq.kasAllowlist) == 0 { //nolint:nestif // not complex
 		if s.KeyAccessServerRegistry != nil {
 			platformEndpoint, err := s.PlatformConfiguration.platformEndpoint()
 			if err != nil {
@@ -145,10 +147,18 @@ func (s SDK) BulkDecrypt(ctx context.Context, opts ...BulkDecryptOption) error {
 			return errors.New("no KAS allowlist provided and no KeyAccessServerRegistry available")
 		}
 	}
+	return nil
+}
+
+// prepareDecryptors creates decryptors and rewrap requests for all TDFs
+func (s SDK) prepareDecryptors(ctx context.Context, bulkReq *BulkDecryptRequest) (map[string][]*kas.UnsignedRewrapRequest_WithPolicyRequest, map[string]decryptor, map[string]*BulkTDF) {
+	kasRewrapRequests := make(map[string][]*kas.UnsignedRewrapRequest_WithPolicyRequest)
+	tdfDecryptors := make(map[string]decryptor)
+	policyTDF := make(map[string]*BulkTDF)
 
 	for i, tdf := range bulkReq.TDFs {
 		policyID := fmt.Sprintf("policy-%d", i)
-		decryptor, err := s.createDecryptor(tdf, bulkReq) //nolint:contextcheck // dont want to change signature of LoadTDF
+		decryptor, err := s.createDecryptor(tdf, bulkReq) //nolint:contextcheck // context is not used in createDecryptor
 		if err != nil {
 			tdf.Error = err
 			continue
@@ -167,9 +177,15 @@ func (s SDK) BulkDecrypt(ctx context.Context, opts ...BulkDecryptOption) error {
 		}
 	}
 
-	kasClient := newKASClient(s.conn.Client, s.conn.Options, s.tokenSource, s.kasSessionKey)
+	return kasRewrapRequests, tdfDecryptors, policyTDF
+}
+
+// performRewraps executes all rewrap requests with KAS servers
+func (s SDK) performRewraps(ctx context.Context, bulkReq *BulkDecryptRequest, kasRewrapRequests map[string][]*kas.UnsignedRewrapRequest_WithPolicyRequest, fulfillableObligations []string) (map[string][]kaoResult, error) {
+	kasClient := newKASClient(s.conn.Client, s.conn.Options, s.tokenSource, s.kasSessionKey, fulfillableObligations)
 	allRewrapResp := make(map[string][]kaoResult)
 	var err error
+
 	for kasurl, rewrapRequests := range kasRewrapRequests {
 		if bulkReq.ignoreAllowList {
 			s.Logger().Warn("kasAllowlist is ignored, kas url is allowed", slog.String("kas_url", kasurl))
@@ -186,6 +202,7 @@ func (s SDK) BulkDecrypt(ctx context.Context, opts ...BulkDecryptOption) error {
 			}
 			continue
 		}
+
 		var rewrapResp map[string][]kaoResult
 		switch bulkReq.TDFType {
 		case Nano:
@@ -198,19 +215,73 @@ func (s SDK) BulkDecrypt(ctx context.Context, opts ...BulkDecryptOption) error {
 			allRewrapResp[id] = append(allRewrapResp[id], res...)
 		}
 	}
+
 	if err != nil {
-		return fmt.Errorf("bulk rewrap failed: %w", err)
+		return nil, fmt.Errorf("bulk rewrap failed: %w", err)
+	}
+
+	return allRewrapResp, nil
+}
+
+// PrepareBulkDecrypt does everything except decrypt from the Bulk Decrypt
+// ! Currently you cannot specify fulfillable obligations on an individual TDF basis
+func (s SDK) PrepareBulkDecrypt(ctx context.Context, opts ...BulkDecryptOption) (*BulkDecryptPrepared, error) {
+	bulkReq, createError := createBulkRewrapRequest(opts...)
+	if createError != nil {
+		return nil, fmt.Errorf("failed to create bulk rewrap request: %w", createError)
+	}
+
+	// Setup KAS allowlist
+	if err := s.setupKasAllowlist(ctx, bulkReq); err != nil {
+		return nil, err
+	}
+
+	// Prepare decryptors and rewrap requests
+	kasRewrapRequests, tdfDecryptors, policyTDF := s.prepareDecryptors(ctx, bulkReq)
+
+	// Use the default fulfillable obligations unless a decryptor is available to provide its own
+	fulfillableObligations := s.fulfillableObligationFQNs
+	if len(tdfDecryptors) > 0 {
+		for _, d := range tdfDecryptors {
+			fulfillableObligations = getFulfillableObligations(d, s.logger)
+			break
+		}
+	}
+
+	// Perform rewraps
+	allRewrapResp, err := s.performRewraps(ctx, bulkReq, kasRewrapRequests, fulfillableObligations)
+	if err != nil {
+		return nil, err
 	}
 
-	var errList []error
 	for id, tdf := range policyTDF {
-		kaoRes, ok := allRewrapResp[id]
+		policyRes, ok := allRewrapResp[id]
+		if !ok {
+			tdf.Error = errors.New("rewrap did not create a response for this TDF")
+			continue
+		}
+		tdf.TriggeredObligations = Obligations{FQNs: dedupRequiredObligations(policyRes)}
+	}
+
+	return &BulkDecryptPrepared{
+		PolicyTDF:     policyTDF,
+		tdfDecryptors: tdfDecryptors,
+		allRewrapResp: allRewrapResp,
+	}, nil
+}
+
+// Allow the bulk decryption to occur
+func (bp *BulkDecryptPrepared) BulkDecrypt(ctx context.Context) error {
+	var errList []error
+	var err error
+	for id, tdf := range bp.PolicyTDF {
+		kaoRes, ok := bp.allRewrapResp[id]
 		if !ok {
 			tdf.Error = errors.New("rewrap did not create a response for this TDF")
 			errList = append(errList, tdf.Error)
 			continue
 		}
-		decryptor := tdfDecryptors[id]
+		decryptor := bp.tdfDecryptors[id]
 		if _, err = decryptor.Decrypt(ctx, kaoRes); err != nil {
 			tdf.Error = err
 			errList = append(errList, tdf.Error)
@@ -225,9 +296,36 @@ func (s SDK) BulkDecrypt(ctx context.Context, opts ...BulkDecryptOption) error {
 	return nil
 }
 
+// BulkDecrypt Decrypts a list of BulkTDF and if a partial failure of TDFs unable to be decrypted, BulkErrors would be returned.
+func (s SDK) BulkDecrypt(ctx context.Context, opts ...BulkDecryptOption) error {
+	prepared, err := s.PrepareBulkDecrypt(ctx, opts...)
+	if err != nil {
+		return err
+	}
+
+	return prepared.BulkDecrypt(ctx)
+}
+
 func (b *BulkDecryptRequest) appendTDFs(tdfs ...*BulkTDF) {
 	b.TDFs = append(
 		b.TDFs,
 		tdfs...,
 	)
 }
+
+func getFulfillableObligations(decryptor decryptor, logger *slog.Logger) []string {
+	if decryptor == nil {
+		logger.Warn("decryptor is nil, cannot populate obligations")
+		return make([]string, 0)
+	}
+
+	switch d := decryptor.(type) {
+	case *tdf3DecryptHandler:
+		return d.reader.config.fulfillableObligationFQNs
+	case *NanoTDFDecryptHandler:
+		return d.config.fulfillableObligationFQNs
+	default:
+		logger.Warn("unknown decryptor type, cannot populate obligations", slog.String("type", fmt.Sprintf("%T", d)))
+		return make([]string, 0)
+	}
+}

sdk/granter_test.go

@@ -30,6 +30,7 @@ const (
 	specifiedKas        = "https://attr.kas.com/"
 	evenMoreSpecificKas = "https://value.kas.com/"
 	lessSpecificKas     = "https://namespace.kas.com/"
+	obligationKas       = "https://obligation.kas.com/"
 	fakePem             = mockRSAPublicKey1
 )
 
@@ -75,6 +76,21 @@ var (
 	mpc, _ = NewAttributeValueFQN("https://virtru.com/attr/mapped/value/c")
 	mpd, _ = NewAttributeValueFQN("https://virtru.com/attr/mapped/value/d")
 	mpu, _ = NewAttributeValueFQN("https://virtru.com/attr/mapped/value/unspecified")
+
+	// Attributes for testing obligations
+
+	OBLIGATIONATTR, _   = NewAttributeNameFQN("https://virtru.com/attr/obligation_test")
+	oa1, _              = NewAttributeValueFQN("https://virtru.com/attr/obligation_test/value/value1")
+	oa2, _              = NewAttributeValueFQN("https://virtru.com/attr/obligation_test/value/value2")
+	oa3, _              = NewAttributeValueFQN("https://virtru.com/attr/obligation_test/value/value3")
+	obligationWatermark = "https://virtru.com/obl/obligation_test/value/watermark"
+	obligationGeofence  = "https://virtru.com/obl/obligation_test/value/geofence"
+	obligationRedact    = "https://virtru.com/obl/obligation_test/value/redact"
+	obligationMap       = map[string]string{
+		oa1.key: obligationWatermark,
+		oa2.key: obligationGeofence,
+		oa3.key: obligationRedact,
+	}
 )
 
 func spongeCase(s string) string {
@@ -211,6 +227,14 @@ func mockAttributeFor(fqn AttributeNameFQN) *policy.Attribute {
 			Rule:      policy.AttributeRuleTypeEnum_ATTRIBUTE_RULE_TYPE_ENUM_ANY_OF,
 			Fqn:       fqn.String(),
 		}
+	case OBLIGATIONATTR.key:
+		return &policy.Attribute{
+			Id:        "OBL",
+			Namespace: &nsOne,
+			Name:      "obligation",
+			Rule:      policy.AttributeRuleTypeEnum_ATTRIBUTE_RULE_TYPE_ENUM_ANY_OF,
+			Fqn:       fqn.String(),
+		}
 	}
 	return nil
 }
@@ -452,6 +476,18 @@ func mockValueFor(fqn AttributeValueFQN) *policy.Value {
 			p.Grants = make([]*policy.KeyAccessServer, 1)
 			p.Grants[0] = mockGrant(evenMoreSpecificKas, "r1")
 		}
+	case OBLIGATIONATTR.key:
+		switch strings.ToLower(fqn.Value()) {
+		case "value1":
+			p.KasKeys = make([]*policy.SimpleKasKey, 1)
+			p.KasKeys[0] = mockSimpleKasKey(obligationKas, "r3")
+		case "value2":
+			p.KasKeys = make([]*policy.SimpleKasKey, 1)
+			p.KasKeys[0] = mockSimpleKasKey(obligationKas, "r3")
+		case "value3":
+			p.KasKeys = make([]*policy.SimpleKasKey, 1)
+			p.KasKeys[0] = mockSimpleKasKey("https://d.kas/", "e1")
+		}
 	}
 	return &p
 }