go.ligato.io/vpp-agent/v3@v3.5.0/plugins/kvscheduler/internal/utils/record.go (about)

     1  // Copyright (c) 2018 Cisco and/or its affiliates.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at:
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package utils
    16  
    17  import (
    18  	"encoding/json"
    19  
    20  	"go.ligato.io/cn-infra/v2/logging"
    21  	"google.golang.org/protobuf/encoding/protojson"
    22  	"google.golang.org/protobuf/proto"
    23  	"google.golang.org/protobuf/reflect/protoreflect"
    24  	"google.golang.org/protobuf/reflect/protoregistry"
    25  
    26  	"go.ligato.io/vpp-agent/v3/pkg/models"
    27  )
    28  
    29  // RecordedProtoMessage is a proto.Message suitable for recording and access via
    30  // REST API.
    31  type RecordedProtoMessage struct {
    32  	proto.Message
    33  	ProtoMsgName string
    34  }
    35  
    36  // ProtoWithName is used to marshall proto message data alongside the proto
    37  // message name.
    38  type ProtoWithName struct {
    39  	ProtoMsgName string
    40  	ProtoMsgData json.RawMessage
    41  }
    42  
    43  // MarshalJSON marshalls proto message using the marshaller from protojson.
    44  // The protojson package produces a different output than the standard "encoding/json"
    45  // package, which does not operate correctly on protocol buffers.
    46  func (p *RecordedProtoMessage) MarshalJSON() ([]byte, error) {
    47  	var (
    48  		msgName string
    49  		msgData []byte
    50  		err     error
    51  	)
    52  	if p != nil {
    53  		msgName = p.ProtoMsgName
    54  		msgData, err = protojson.Marshal(p.Message)
    55  		if err != nil {
    56  			return nil, err
    57  		}
    58  	}
    59  	return json.Marshal(&ProtoWithName{
    60  		ProtoMsgName: msgName,
    61  		ProtoMsgData: json.RawMessage(msgData),
    62  	})
    63  }
    64  
    65  // UnmarshalJSON un-marshalls proto message using the marshaller from protojson.
    66  // The protojson package produces a different output than the standard "encoding/json"
    67  // package, which does not operate correctly on protocol buffers.
    68  func (p *RecordedProtoMessage) UnmarshalJSON(data []byte) error {
    69  	var pwn ProtoWithName
    70  	if err := json.Unmarshal(data, &pwn); err != nil {
    71  		return err
    72  	}
    73  	p.ProtoMsgName = pwn.ProtoMsgName
    74  	if p.ProtoMsgName == "" {
    75  		return nil
    76  	}
    77  
    78  	// try to find the message type in the default registry
    79  	typeRegistry := models.DefaultRegistry.MessageTypeRegistry()
    80  	fullMsgName := protoreflect.FullName(p.ProtoMsgName)
    81  	msgType, err := typeRegistry.FindMessageByName(fullMsgName)
    82  	if err != nil {
    83  		// if not found use the proto global types registry as a fallback
    84  		logging.Debugf("cannot get message type for message name %s from default registry: %v", fullMsgName, err)
    85  		msgType, err = protoregistry.GlobalTypes.FindMessageByName(fullMsgName)
    86  	}
    87  	if err != nil {
    88  		return err
    89  	}
    90  
    91  	msg := msgType.New().Interface()
    92  	if err = protojson.Unmarshal(pwn.ProtoMsgData, msg); err != nil {
    93  		return err
    94  	}
    95  	p.Message = msg
    96  	return nil
    97  }
    98  
    99  // RecordProtoMessage prepares proto message for recording and potential
   100  // access via REST API.
   101  // Note: no need to clone the message - once un-marshalled, the content is never
   102  // changed (otherwise it would break prev-new value comparisons).
   103  func RecordProtoMessage(msg proto.Message) *RecordedProtoMessage {
   104  	if msg == nil {
   105  		return nil
   106  	}
   107  	return &RecordedProtoMessage{
   108  		Message:      msg,
   109  		ProtoMsgName: string(proto.MessageName(msg)),
   110  	}
   111  }