kuma

1
package proto
2

3
import (
4
	"strings"
5

6
	protov1 "github.com/golang/protobuf/proto"
7
	"github.com/pkg/errors"
8
	"google.golang.org/protobuf/proto"
9
	"google.golang.org/protobuf/reflect/protoreflect"
10
	"google.golang.org/protobuf/reflect/protoregistry"
11
	"google.golang.org/protobuf/types/known/anypb"
12
)
13

14
const googleApis = "type.googleapis.com/"
15

16
// When saving Snapshot in SnapshotCache we generate version based on proto.Equal()
17
// Therefore we need deterministic way of marshaling Any which is part of the Protobuf on which we execute Equal()
18
//
19
// Based on proto.MarshalAny
20
func MarshalAnyDeterministic(pb proto.Message) (*anypb.Any, error) {
21
	bytes, err := proto.MarshalOptions{Deterministic: true}.Marshal(pb)
22
	if err != nil {
23
		return nil, err
24
	}
25
	name := string(protov1.MessageV2(pb).ProtoReflect().Descriptor().FullName())
26
	return &anypb.Any{TypeUrl: googleApis + name, Value: bytes}, nil
27
}
28

29
func MustMarshalAny(pb proto.Message) *anypb.Any {
30
	msg, err := MarshalAnyDeterministic(pb)
31
	if err != nil {
32
		panic(err.Error())
33
	}
34
	return msg
35
}
36

37
func UnmarshalAnyTo(src *anypb.Any, dst proto.Message) error {
38
	return anypb.UnmarshalTo(src, dst, proto.UnmarshalOptions{})
39
}
40

41
// MergeAnys merges two Any messages of the same type. We cannot just use proto#Merge on Any directly because values are encoded in byte slices.
42
// Instead we have to unmarshal types, merge them and marshal again.
43
func MergeAnys(dst *anypb.Any, src *anypb.Any) (*anypb.Any, error) {
44
	if src == nil {
45
		return dst, nil
46
	}
47
	if dst == nil {
48
		return src, nil
49
	}
50
	if src.TypeUrl != dst.TypeUrl {
51
		return nil, errors.Errorf("type URL of dst %q is different than src %q", dst.TypeUrl, src.TypeUrl)
52
	}
53

54
	msgType, err := FindMessageType(dst.TypeUrl)
55
	if err != nil {
56
		return nil, err
57
	}
58

59
	dstMsg := msgType.New().Interface()
60
	if err := proto.Unmarshal(dst.Value, dstMsg); err != nil {
61
		return nil, err
62
	}
63

64
	srcMsg := msgType.New().Interface()
65
	if err := proto.Unmarshal(src.Value, srcMsg); err != nil {
66
		return nil, err
67
	}
68

69
	Merge(dstMsg, srcMsg)
70
	return MarshalAnyDeterministic(dstMsg)
71
}
72

73
func FindMessageType(typeUrl string) (protoreflect.MessageType, error) {
74
	// TypeURL in Any contains type.googleapis.com/ prefix, but in Proto
75
	// registry it does not have this prefix.
76
	msgTypeName := strings.ReplaceAll(typeUrl, googleApis, "")
77
	fullName := protoreflect.FullName(msgTypeName)
78

79
	return protoregistry.GlobalTypes.FindMessageByName(fullName)
80
}
81