feat(sdk): Add a KAS allowlist

sdk/bulk.go

@@ -84,7 +84,7 @@ func (s SDK) BulkDecrypt(ctx context.Context, opts ...BulkDecryptOption) error {
 
 	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.TDFType) //nolint:contextcheck // dont want to change signature of LoadTDF
 		if err != nil {
 			tdf.Error = err
 			continue

sdk/tdf.go

@@ -9,12 +9,14 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"log/slog"
 	"math"
 	"strconv"
 	"strings"
 
 	"github.com/Masterminds/semver/v3"
 	"github.com/opentdf/platform/protocol/go/kas"
+	"github.com/opentdf/platform/protocol/go/policy/kasregistry"
 
 	"github.com/google/uuid"
 	"github.com/opentdf/platform/lib/ocrypto"
@@ -664,6 +666,31 @@ func (s SDK) LoadTDF(reader io.ReadSeeker, opts ...TDFReaderOption) (*Reader, er
 		return nil, fmt.Errorf("newAssertionConfig failed: %w", err)
 	}
 
+	if config.kasAllowlist == nil && !config.ignoreAllowList { //nolint:nestif // need to handle both cases
+		if s.KeyAccessServerRegistry != nil {
+			// retrieve the registered kases if not provided
+			kases, err := s.KeyAccessServerRegistry.ListKeyAccessServers(context.Background(), &kasregistry.ListKeyAccessServersRequest{})
+			if err != nil {
+				return nil, fmt.Errorf("kasregistry.ListKeyAccessServers failed: %w", err)
+			}
+			config.kasAllowlist = AllowList{}
+			for _, kas := range kases.GetKeyAccessServers() {
+				err = config.kasAllowlist.Add(kas.GetUri())
+				if err != nil {
+					return nil, fmt.Errorf("kasAllowlist.Add failed: %w", err)
+				}
+			}
+			// grpc target does not have a scheme
+			err = config.kasAllowlist.Add("http://" + s.conn.Target())
+			if err != nil {
+				return nil, fmt.Errorf("kasAllowlist.Add failed: %w", err)
+			}
+		} else {
+			slog.Warn("No KAS allowlist provided and no KeyAccessServerRegistry available, ignoring allowlist")
+			config.ignoreAllowList = true
+		}
+	}
+
 	manifest, err := tdfReader.Manifest()
 	if err != nil {
 		return nil, fmt.Errorf("tdfReader.Manifest failed: %w", err)
@@ -954,6 +981,15 @@ func createRewrapRequest(_ context.Context, r *Reader) (map[string]*kas.Unsigned
 	kasReqs := make(map[string]*kas.UnsignedRewrapRequest_WithPolicyRequest)
 	for i, kao := range r.manifest.EncryptionInformation.KeyAccessObjs {
 		kaoID := fmt.Sprintf("kao-%d", i)
+
+		// if ignoreing allowlist then warn
+		// if kas url is not allowed then return error
+		if r.config.ignoreAllowList {
+			slog.Warn(fmt.Sprintf("KasAllowlist is ignored, kas url %s is allowed", kao.KasURL))
+		} else if !r.config.kasAllowlist.IsAllowed(kao.KasURL) {
+			return nil, fmt.Errorf("KasAllowlist: kas url %s is not allowed", kao.KasURL)
+		}
+
 		key, err := ocrypto.Base64Decode([]byte(kao.WrappedKey))
 		if err != nil {
 			return nil, fmt.Errorf("could not decode wrapper key: %w", err)

sdk/tdf_config.go

@@ -3,6 +3,8 @@ package sdk
 import (
 	"errors"
 	"fmt"
+	"net"
+	"net/url"
 
 	"github.com/opentdf/platform/lib/ocrypto"
 	"github.com/opentdf/platform/protocol/go/policy"
@@ -275,6 +277,68 @@ type TDFReaderConfig struct {
 
 	schemaValidationIntensity SchemaValidationIntensity
 	kasSessionKey             ocrypto.KeyPair
+	kasAllowlist              AllowList // KAS URLs that are allowed to be used for reading TDFs
+	ignoreAllowList           bool      // If true, the kasAllowlist will be ignored, and all KAS URLs will be allowed
+}
+
+type AllowList map[string]bool
+
+func getKasAddress(kasURL string) (string, error) {
+	parsedURL, err := url.Parse(kasURL)
+	if err != nil {
+		return "", fmt.Errorf("cannot parse kas url(%s): %w", kasURL, err)
+	}
+
+	// Needed to support buffconn for testing
+	if parsedURL.Host == "" && parsedURL.Port() == "" {
+		return "", nil
+	}
+
+	port := parsedURL.Port()
+	// if port is empty, default to 443.
+	if port == "" {
+		return parsedURL.Hostname(), nil
+	}
+
+	return net.JoinHostPort(parsedURL.Hostname(), port), nil
+}
+
+func newAllowList(kasList []string) (AllowList, error) {
+	allowList := make(AllowList)
+	for _, kasURL := range kasList {
+		err := allowList.Add(kasURL)
+		if err != nil {
+			return nil, fmt.Errorf("error adding kas url(%s) to allowlist: %w", kasURL, err)
+		}
+	}
+	return allowList, nil
+}
+
+func (a AllowList) IsAllowed(kasURL string) bool {
+	if a == nil {
+		return false // No allowlist, so no URLs are allowed
+	}
+	kasAddress, err := getKasAddress(kasURL)
+	if err != nil {
+		return false // If we can't parse the URL, we can't allow it
+	}
+	_, ok := a[kasAddress]
+	return ok
+}
+
+func (a AllowList) Add(kasURL string) error {
+	if a == nil {
+		a = make(AllowList)
+	}
+	kasAddress, err := getKasAddress(kasURL)
+	if err != nil {
+		// If we can't parse the URL, we can't add it to the allowlist
+		return fmt.Errorf("error parsing kas url(%s): %w", kasURL, err)
+	} else if kasAddress == "" {
+		return fmt.Errorf("kas url(%s) not parsed", kasURL)
+	}
+	a[kasAddress] = true
+	return nil
 }
 
 func newTDFReaderConfig(opt ...TDFReaderOption) (*TDFReaderConfig, error) {
@@ -332,6 +396,24 @@ func WithSessionKeyType(keyType ocrypto.KeyType) TDFReaderOption {
 	}
 }
 
+func WithKasAllowlist(kasList []string) TDFReaderOption {
+	return func(c *TDFReaderConfig) error {
+		allowlist, err := newAllowList(kasList)
+		if err != nil {
+			return fmt.Errorf("failed to create kas allowlist: %w", err)
+		}
+		c.kasAllowlist = allowlist
+		return nil
+	}
+}
+
+func WithIgnoreAllowlist(ignore bool) TDFReaderOption {
+	return func(c *TDFReaderConfig) error {
+		c.ignoreAllowList = ignore
+		return nil
+	}
+}
+
 func withSessionKey(k ocrypto.KeyPair) TDFReaderOption {
 	return func(c *TDFReaderConfig) error {
 		c.kasSessionKey = k

sdk/tdf_config_test.go

@@ -197,7 +197,7 @@ func TestWithTargetMode(t *testing.T) {
 			cfg, err = newTDFConfig(WithTargetMode(test.targetMode))
 
 			if test.expectError {
-				assert.Error(t, err)
+				require.Error(t, err)
 				return
 			}
 
@@ -207,3 +207,155 @@ func TestWithTargetMode(t *testing.T) {
 		})
 	}
 }
+
+func TestAllowList_Add(t *testing.T) {
+	tests := []struct {
+		name        string
+		kasURL      string
+		entry       string
+		expectError bool
+	}{
+		{
+			name:        "Valid URL with port",
+			kasURL:      "https://example.com:443",
+			entry:       "example.com:443",
+			expectError: false,
+		},
+		{
+			name:        "Valid URL without port",
+			kasURL:      "https://example.com",
+			entry:       "example.com",
+			expectError: false,
+		},
+		{
+			name:        "Invalid URL",
+			kasURL:      "invalid-url",
+			expectError: true,
+		},
+		{
+			name:        "Empty URL",
+			kasURL:      "",
+			expectError: true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			allowList := make(AllowList)
+			err := allowList.Add(tt.kasURL)
+			if tt.expectError {
+				require.Error(t, err, "Expected an error for test case: %s", tt.name)
+			} else {
+				require.NoError(t, err, "Did not expect an error for test case: %s", tt.name)
+				assert.Contains(t, allowList, tt.entry, "Expected URL to be added to the allowlist")
+			}
+		})
+	}
+}
+
+func TestAllowList_IsAllowed(t *testing.T) {
+	allowList := make(AllowList)
+	_ = allowList.Add("https://example.com:443")
+	_ = allowList.Add("https://another.com")
+
+	tests := []struct {
+		name     string
+		kasURL   string
+		expected bool
+	}{
+		{
+			name:     "Allowed URL with port",
+			kasURL:   "https://example.com:443",
+			expected: true,
+		},
+		{
+			name:     "Allowed URL without port",
+			kasURL:   "https://another.com",
+			expected: true,
+		},
+		{
+			name:     "Not allowed URL",
+			kasURL:   "https://notallowed.com",
+			expected: false,
+		},
+		{
+			name:     "Invalid URL",
+			kasURL:   "invalid-url",
+			expected: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := allowList.IsAllowed(tt.kasURL)
+			assert.Equal(t, tt.expected, result, "Unexpected result for test case: %s", tt.name)
+		})
+	}
+}
+
+func TestWithKasAllowlist(t *testing.T) {
+	tests := []struct {
+		name        string
+		kasList     []string
+		expectError bool
+	}{
+		{
+			name:        "Valid KAS URLs",
+			kasList:     []string{"https://example.com:443", "https://another.com"},
+			expectError: false,
+		},
+		{
+			name:        "Invalid KAS URL",
+			kasList:     []string{"invalid-url"},
+			expectError: true,
+		},
+		{
+			name:        "Empty KAS list",
+			kasList:     []string{},
+			expectError: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			config := &TDFReaderConfig{}
+			err := WithKasAllowlist(tt.kasList)(config)
+			if tt.expectError {
+				require.Error(t, err, "Expected an error for test case: %s", tt.name)
+			} else {
+				require.NoError(t, err, "Did not expect an error for test case: %s", tt.name)
+				for _, kasURL := range tt.kasList {
+					assert.True(t, config.kasAllowlist.IsAllowed(kasURL), "Expected KAS URL to be allowed: %s", kasURL)
+				}
+			}
+		})
+	}
+}
+
+func TestWithIgnoreAllowlist(t *testing.T) {
+	tests := []struct {
+		name          string
+		ignore        bool
+		expectedValue bool
+	}{
+		{
+			name:          "Ignore allowlist set to true",
+			ignore:        true,
+			expectedValue: true,
+		},
+		{
+			name:          "Ignore allowlist set to false",
+			ignore:        false,
+			expectedValue: false,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			config := &TDFReaderConfig{}
+			err := WithIgnoreAllowlist(tt.ignore)(config)
+			require.NoError(t, err, "Did not expect an error for test case: %s", tt.name)
+			assert.Equal(t, tt.expectedValue, config.ignoreAllowList, "Unexpected value for ignoreAllowList in test case: %s", tt.name)
+		})
+	}
+}