From 647ac7fbfd1ee9b98adebeb7f336d5c4b3a8e955 Mon Sep 17 00:00:00 2001
From: Denis Mishin <dmishin@pomerium.com>
Date: Mon, 28 Nov 2022 11:31:57 -0500
Subject: [PATCH] storage: ignore removed fields when deserializing the data
 (#3768)

ignore removed fields when deserializing the data
---
 pkg/protoutil/any.go                  | 14 ++++++++++++++
 pkg/storage/postgres/postgres.go      | 16 +++++++---------
 pkg/storage/postgres/registry_test.go | 17 +++++++++++++++++
 3 files changed, 38 insertions(+), 9 deletions(-)

diff --git a/pkg/protoutil/any.go b/pkg/protoutil/any.go
index 72e6a7a4ca6..2e0cd8e5c50 100644
--- a/pkg/protoutil/any.go
+++ b/pkg/protoutil/any.go
@@ -1,6 +1,7 @@
 package protoutil
 
 import (
+	"google.golang.org/protobuf/encoding/protojson"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/anypb"
 	"google.golang.org/protobuf/types/known/wrapperspb"
@@ -58,6 +59,19 @@ func NewAny(msg proto.Message) *anypb.Any {
 	return a
 }
 
+// UnmarshalAnyJSON unmarshals JSON data into Any
+func UnmarshalAnyJSON(data []byte) (*anypb.Any, error) {
+	opts := protojson.UnmarshalOptions{
+		AllowPartial:   true,
+		DiscardUnknown: true,
+	}
+	var val anypb.Any
+	if err := opts.Unmarshal(data, &val); err != nil {
+		return nil, err
+	}
+	return &val, nil
+}
+
 // NewAnyBool creates a new any type from a bool.
 func NewAnyBool(v bool) *anypb.Any {
 	return NewAny(wrapperspb.Bool(v))
diff --git a/pkg/storage/postgres/postgres.go b/pkg/storage/postgres/postgres.go
index 797bbd0f750..7d120e4ab3e 100644
--- a/pkg/storage/postgres/postgres.go
+++ b/pkg/storage/postgres/postgres.go
@@ -19,6 +19,7 @@ import (
 
 	"github.com/pomerium/pomerium/pkg/grpc/databroker"
 	"github.com/pomerium/pomerium/pkg/grpc/registry"
+	"github.com/pomerium/pomerium/pkg/protoutil"
 	"github.com/pomerium/pomerium/pkg/storage"
 )
 
@@ -124,8 +125,7 @@ func getNextChangedRecord(ctx context.Context, q querier, recordType string, aft
 		}
 		afterRecordVersion = version
 
-		var any anypb.Any
-		err = protojson.Unmarshal(data.Bytes, &any)
+		any, err := protoutil.UnmarshalAnyJSON(data.Bytes)
 		if isUnknownType(err) {
 			// ignore
 			continue
@@ -137,7 +137,7 @@ func getNextChangedRecord(ctx context.Context, q querier, recordType string, aft
 			Version:    version,
 			Type:       recordType,
 			Id:         recordID,
-			Data:       &any,
+			Data:       any,
 			ModifiedAt: timestamppbFromTimestamptz(modifiedAt),
 			DeletedAt:  timestamppbFromTimestamptz(deletedAt),
 		}, nil
@@ -176,8 +176,7 @@ func getRecord(ctx context.Context, q querier, recordType, recordID string) (*da
 		return nil, fmt.Errorf("postgres: failed to execute query: %w", err)
 	}
 
-	var any anypb.Any
-	err = protojson.Unmarshal(data.Bytes, &any)
+	any, err := protoutil.UnmarshalAnyJSON(data.Bytes)
 	if isUnknownType(err) {
 		return nil, storage.ErrNotFound
 	} else if err != nil {
@@ -188,7 +187,7 @@ func getRecord(ctx context.Context, q querier, recordType, recordID string) (*da
 		Version:    version,
 		Type:       recordType,
 		Id:         recordID,
-		Data:       &any,
+		Data:       any,
 		ModifiedAt: timestamppbFromTimestamptz(modifiedAt),
 	}, nil
 }
@@ -228,8 +227,7 @@ func listRecords(ctx context.Context, q querier, expr storage.FilterExpression,
 			return nil, fmt.Errorf("postgres: failed to scan row: %w", err)
 		}
 
-		var any anypb.Any
-		err = protojson.Unmarshal(data.Bytes, &any)
+		any, err := protoutil.UnmarshalAnyJSON(data.Bytes)
 		if isUnknownType(err) {
 			// ignore records with an unknown type
 			continue
@@ -241,7 +239,7 @@ func listRecords(ctx context.Context, q querier, expr storage.FilterExpression,
 			Version:    version,
 			Type:       recordType,
 			Id:         id,
-			Data:       &any,
+			Data:       any,
 			ModifiedAt: timestamppbFromTimestamptz(modifiedAt),
 		})
 	}
diff --git a/pkg/storage/postgres/registry_test.go b/pkg/storage/postgres/registry_test.go
index dbf1963fc1c..0f015d5e79c 100644
--- a/pkg/storage/postgres/registry_test.go
+++ b/pkg/storage/postgres/registry_test.go
@@ -14,6 +14,7 @@ import (
 
 	"github.com/pomerium/pomerium/internal/testutil"
 	"github.com/pomerium/pomerium/pkg/grpc/registry"
+	"github.com/pomerium/pomerium/pkg/protoutil"
 )
 
 type mockRegistryWatchServer struct {
@@ -112,3 +113,19 @@ func TestRegistry(t *testing.T) {
 		return nil
 	}))
 }
+
+func TestUnmarshalJSONUnknownFields(t *testing.T) {
+	any, err := protoutil.UnmarshalAnyJSON([]byte(`
+	{
+		"@type": "type.googleapis.com/registry.Service",
+		"kind": "AUTHENTICATE",
+		"endpoint": "endpoint",
+		"unknown_field": true
+	  }
+	`))
+	require.NoError(t, err)
+	var val registry.Service
+	require.NoError(t, any.UnmarshalTo(&val))
+	assert.Equal(t, registry.ServiceKind_AUTHENTICATE, val.Kind)
+	assert.Equal(t, "endpoint", val.Endpoint)
+}