feat(sdk): Add a KAS allowlist

sdk/bulk.go

@@ -17,8 +17,10 @@ type BulkTDF struct {
 }
 
 type BulkDecryptRequest struct {
-	TDFs    []*BulkTDF
-	TDFType TdfType
+	TDFs                  []*BulkTDF
+	TDF3DecryptOptions    []TDFReaderOption     // Options for TDF3 Decryptor
+	NanoTDFDecryptOptions []NanoTDFReaderOption // Options for Nano TDF Decryptor
+	TDFType               TdfType
 }
 
 // BulkErrors List of Errors that Failed during Bulk Decryption
@@ -55,6 +57,18 @@ func WithTDFType(tdfType TdfType) BulkDecryptOption {
 	}
 }
 
+func WithTDF3DecryptOptions(options ...TDFReaderOption) BulkDecryptOption {
+	return func(request *BulkDecryptRequest) {
+		request.TDF3DecryptOptions = append(request.TDF3DecryptOptions, options...)
+	}
+}
+
+func WithNanoTDFDecryptOptions(options ...NanoTDFReaderOption) BulkDecryptOption {
+	return func(request *BulkDecryptRequest) {
+		request.NanoTDFDecryptOptions = append(request.NanoTDFDecryptOptions, options...)
+	}
+}
+
 func createBulkRewrapRequest(options ...BulkDecryptOption) *BulkDecryptRequest {
 	req := &BulkDecryptRequest{}
 	for _, opt := range options {
@@ -63,16 +77,15 @@ func createBulkRewrapRequest(options ...BulkDecryptOption) *BulkDecryptRequest {
 	return req
 }
 
-func (s SDK) createDecryptor(tdf *BulkTDF, tdfType TdfType) (decryptor, error) {
-	switch tdfType {
+func (s SDK) createDecryptor(tdf *BulkTDF, req *BulkDecryptRequest) (decryptor, error) {
+	switch req.TDFType {
 	case Nano:
-		decryptor := createNanoTDFDecryptHandler(tdf.Reader, tdf.Writer)
-		return decryptor, nil
+		return createNanoTDFDecryptHandler(tdf.Reader, tdf.Writer, req.NanoTDFDecryptOptions...)
 	case Standard:
-		return s.createTDF3DecryptHandler(tdf.Writer, tdf.Reader)
+		return s.createTDF3DecryptHandler(tdf.Writer, tdf.Reader, req.TDF3DecryptOptions...)
 	case Invalid:
 	}
-	return nil, fmt.Errorf("unknown tdf type: %s", tdfType)
+	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.
@@ -82,9 +95,45 @@ func (s SDK) BulkDecrypt(ctx context.Context, opts ...BulkDecryptOption) error {
 	tdfDecryptors := make(map[string]decryptor)
 	policyTDF := make(map[string]*BulkTDF)
 
+	switch bulkReq.TDFType {
+	case Nano:
+		dummy := &NanoTDFReaderConfig{}
+		for _, opt := range bulkReq.NanoTDFDecryptOptions {
+			_ = opt(dummy)
+		}
+		if !dummy.ignoreAllowList && dummy.kasAllowlist == nil {
+			// if no kasAllowlist is set, we get the allowlist from the registry
+			allowlist, err := allowListFromKASRegistry(ctx, s.KeyAccessServerRegistry, s.conn.Target())
+			if err != nil {
+				return fmt.Errorf("failed to get allowlist from registry: %w", err)
+			}
+			bulkReq.NanoTDFDecryptOptions = append(
+				bulkReq.NanoTDFDecryptOptions,
+				withNanoKasAllowlist(allowlist),
+			)
+		}
+	case Standard:
+		dummy := &TDFReaderConfig{}
+		for _, opt := range bulkReq.TDF3DecryptOptions {
+			_ = opt(dummy)
+		}
+		if !dummy.ignoreAllowList && dummy.kasAllowlist == nil {
+			// if no kasAllowlist is set, we get the allowlist from the registry
+			allowlist, err := allowListFromKASRegistry(ctx, s.KeyAccessServerRegistry, s.conn.Target())
+			if err != nil {
+				return fmt.Errorf("failed to get allowlist from registry: %w", err)
+			}
+			bulkReq.TDF3DecryptOptions = append(
+				bulkReq.TDF3DecryptOptions,
+				withKasAllowlist(allowlist),
+			)
+		}
+	case Invalid:
+	}
+
 	for i, tdf := range bulkReq.TDFs {
 		policyID := fmt.Sprintf("policy-%d", i)
-		decryptor, err := s.createDecryptor(tdf, bulkReq.TDFType)
+		decryptor, err := s.createDecryptor(tdf, bulkReq) //nolint:contextcheck // dont want to change signature of LoadTDF
 		if err != nil {
 			tdf.Error = err
 			continue

sdk/nanotdf.go

@@ -907,13 +907,20 @@ type NanoTDFDecryptHandler struct {
 
 	header    NanoTDFHeader
 	headerBuf []byte
+
+	config *NanoTDFReaderConfig
 }
 
-func createNanoTDFDecryptHandler(reader io.ReadSeeker, writer io.Writer) *NanoTDFDecryptHandler {
+func createNanoTDFDecryptHandler(reader io.ReadSeeker, writer io.Writer, opts ...NanoTDFReaderOption) (*NanoTDFDecryptHandler, error) {
+	nanoTdfReaderConfig, err := newNanoTDFReaderConfig(opts...)
+	if err != nil {
+		return nil, fmt.Errorf("newNanoTDFReaderConfig failed: %w", err)
+	}
 	return &NanoTDFDecryptHandler{
 		reader: reader,
 		writer: writer,
-	}
+		config: nanoTdfReaderConfig,
+	}, nil
 }
 
 func (n *NanoTDFDecryptHandler) getRawHeader() []byte {
@@ -942,6 +949,12 @@ func (n *NanoTDFDecryptHandler) CreateRewrapRequest(_ context.Context) (map[stri
 		return nil, err
 	}
 
+	if n.config.ignoreAllowList {
+		slog.Warn(fmt.Sprintf("KasAllowlist is ignored, kas url %s is allowed", kasURL))
+	} else if !n.config.kasAllowlist.IsAllowed(kasURL) {
+		return nil, fmt.Errorf("KasAllowlist: kas url %s is not allowed", kasURL)
+	}
+
 	req := &kas.UnsignedRewrapRequest_WithPolicyRequest{
 		KeyAccessObjects: []*kas.UnsignedRewrapRequest_WithKeyAccessObject{
 			{
@@ -1017,13 +1030,30 @@ func (n *NanoTDFDecryptHandler) Decrypt(_ context.Context, result []kaoResult) (
 }
 
 // ReadNanoTDF - read the nano tdf and return the decrypted data from it
-func (s SDK) ReadNanoTDF(writer io.Writer, reader io.ReadSeeker) (int, error) {
-	return s.ReadNanoTDFContext(context.Background(), writer, reader)
+func (s SDK) ReadNanoTDF(writer io.Writer, reader io.ReadSeeker, opts ...NanoTDFReaderOption) (int, error) {
+	return s.ReadNanoTDFContext(context.Background(), writer, reader, opts...)
 }
 
 // ReadNanoTDFContext - allows cancelling the reader
-func (s SDK) ReadNanoTDFContext(ctx context.Context, writer io.Writer, reader io.ReadSeeker) (int, error) {
-	handler := createNanoTDFDecryptHandler(reader, writer)
+func (s SDK) ReadNanoTDFContext(ctx context.Context, writer io.Writer, reader io.ReadSeeker, opts ...NanoTDFReaderOption) (int, error) {
+	handler, err := createNanoTDFDecryptHandler(reader, writer, opts...)
+	if err != nil {
+		return 0, fmt.Errorf("createNanoTDFDecryptHandler failed: %w", err)
+	}
+
+	if handler.config.kasAllowlist == nil && !handler.config.ignoreAllowList {
+		if s.KeyAccessServerRegistry != nil {
+			// retrieve the registered kases if not provided
+			allowList, err := allowListFromKASRegistry(ctx, s.KeyAccessServerRegistry, s.conn.Target())
+			if err != nil {
+				return 0, fmt.Errorf("allowListFromKASRegistry failed: %w", err)
+			}
+			handler.config.kasAllowlist = allowList
+		} else {
+			slog.Warn("No KAS allowlist provided and no KeyAccessServerRegistry available, ignoring allowlist")
+			handler.config.ignoreAllowList = true
+		}
+	}
 
 	symmetricKey, err := s.getNanoRewrapKey(ctx, handler)
 	if err != nil {

sdk/nanotdf_config.go

@@ -110,3 +110,48 @@ func WithECDSAPolicyBinding() NanoTDFOption {
 		return nil
 	}
 }
+
+type NanoTDFReaderConfig struct {
+	kasAllowlist    AllowList
+	ignoreAllowList bool
+}
+
+func newNanoTDFReaderConfig(opt ...NanoTDFReaderOption) (*NanoTDFReaderConfig, error) {
+	c := &NanoTDFReaderConfig{}
+
+	for _, o := range opt {
+		err := o(c)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return c, nil
+}
+
+type NanoTDFReaderOption func(*NanoTDFReaderConfig) error
+
+func WithNanoKasAllowlist(kasList []string) NanoTDFReaderOption {
+	return func(c *NanoTDFReaderConfig) error {
+		allowlist, err := newAllowList(kasList)
+		if err != nil {
+			return fmt.Errorf("failed to create kas allowlist: %w", err)
+		}
+		c.kasAllowlist = allowlist
+		return nil
+	}
+}
+
+func withNanoKasAllowlist(allowlist AllowList) NanoTDFReaderOption {
+	return func(c *NanoTDFReaderConfig) error {
+		c.kasAllowlist = allowlist
+		return nil
+	}
+}
+
+func WithNanoIgnoreAllowlist(ignore bool) NanoTDFReaderOption {
+	return func(c *NanoTDFReaderConfig) error {
+		c.ignoreAllowList = ignore
+		return nil
+	}
+}

sdk/nanotdf_config_test.go

@@ -2,6 +2,9 @@ package sdk
 
 import (
 	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 // TestNanoTDFConfig1 - Create a new config, verify that the config contains valid PEMs for the key pair
@@ -53,3 +56,75 @@ func TestNanoTDFConfig2(t *testing.T) {
 		t.Fatalf("expect %s, got %s", kasURL, readKasURL)
 	}
 }
+
+func TestNewNanoTDFReaderConfig(t *testing.T) {
+	t.Run("Valid options", func(t *testing.T) {
+		config, err := newNanoTDFReaderConfig(
+			WithNanoKasAllowlist([]string{"https://example.com:443", "https://another.com"}),
+			WithNanoIgnoreAllowlist(true),
+		)
+		require.NoError(t, err, "Expected no error when creating NanoTDFReaderConfig with valid options")
+		assert.NotNil(t, config, "Expected NanoTDFReaderConfig to be created")
+		assert.True(t, config.ignoreAllowList, "Expected ignoreAllowList to be true")
+		assert.True(t, config.kasAllowlist.IsAllowed("https://example.com:443"), "Expected KAS URL to be allowed")
+		assert.True(t, config.kasAllowlist.IsAllowed("https://another.com"), "Expected KAS URL to be allowed")
+	})
+
+	t.Run("Invalid KAS URL in allowlist", func(t *testing.T) {
+		config, err := newNanoTDFReaderConfig(
+			WithNanoKasAllowlist([]string{"invalid-url"}),
+		)
+		require.Error(t, err, "Expected an error when creating NanoTDFReaderConfig with invalid KAS URL")
+		assert.Nil(t, config, "Expected NanoTDFReaderConfig to be nil")
+	})
+}
+
+func TestWithNanoKasAllowlist(t *testing.T) {
+	t.Run("Valid KAS URLs", func(t *testing.T) {
+		config := &NanoTDFReaderConfig{}
+		err := WithNanoKasAllowlist([]string{"https://example.com:443", "https://another.com"})(config)
+		require.NoError(t, err, "Expected no error when adding valid KAS URLs to allowlist")
+		assert.True(t, config.kasAllowlist.IsAllowed("https://example.com:443"), "Expected KAS URL to be allowed")
+		assert.True(t, config.kasAllowlist.IsAllowed("https://another.com"), "Expected KAS URL to be allowed")
+	})
+
+	t.Run("Invalid KAS URL", func(t *testing.T) {
+		config := &NanoTDFReaderConfig{}
+		err := WithNanoKasAllowlist([]string{"invalid-url"})(config)
+		require.Error(t, err, "Expected an error when adding invalid KAS URL to allowlist")
+	})
+}
+
+func TestWithNanoIgnoreAllowlist(t *testing.T) {
+	t.Run("Set ignoreAllowList to true", func(t *testing.T) {
+		config := &NanoTDFReaderConfig{}
+		err := WithNanoIgnoreAllowlist(true)(config)
+		require.NoError(t, err, "Expected no error when setting ignoreAllowList to true")
+		assert.True(t, config.ignoreAllowList, "Expected ignoreAllowList to be true")
+	})
+
+	t.Run("Set ignoreAllowList to false", func(t *testing.T) {
+		config := &NanoTDFReaderConfig{}
+		err := WithNanoIgnoreAllowlist(false)(config)
+		require.NoError(t, err, "Expected no error when setting ignoreAllowList to false")
+		assert.False(t, config.ignoreAllowList, "Expected ignoreAllowList to be false")
+	})
+}
+
+func TestWithNanoKasAllowlist_with(t *testing.T) {
+	t.Run("Valid AllowList", func(t *testing.T) {
+		allowlist := AllowList{"example.com:443": true}
+		config := &NanoTDFReaderConfig{}
+		err := withNanoKasAllowlist(allowlist)(config)
+		require.NoError(t, err, "Expected no error when setting valid AllowList")
+		assert.True(t, config.kasAllowlist.IsAllowed("https://example.com:443"), "Expected KAS URL to be allowed")
+	})
+
+	t.Run("Empty AllowList", func(t *testing.T) {
+		allowlist := AllowList{}
+		config := &NanoTDFReaderConfig{}
+		err := withNanoKasAllowlist(allowlist)(config)
+		require.NoError(t, err, "Expected no error when setting empty AllowList")
+		assert.False(t, config.kasAllowlist.IsAllowed("https://example.com:443"), "Expected KAS URL to not be allowed")
+	})
+}