github.com/Rookout/GoSDK@v0.1.48/pkg/common/utils.go (about)

     1  package common
     2  
     3  import (
     4  	"fmt"
     5  	pb "github.com/Rookout/GoSDK/pkg/protobuf"
     6  	"github.com/golang/protobuf/proto"
     7  	"github.com/golang/protobuf/ptypes/any"
     8  	"github.com/sirupsen/logrus"
     9  	"strings"
    10  	"time"
    11  
    12  	"github.com/golang/protobuf/ptypes"
    13  )
    14  
    15  
    16  const (
    17  	GoogleProtoTypeURLPrefix          = "type.googleapis.com/"
    18  	MessageTypeInitAugs               = "com.rookout.InitialAugsCommand"
    19  	MessageTypeDataOnPremTokenRequest = "com.rookout.DataOnPremTokenRequest"
    20  	MessageTypeInitC2C                = "com.rookout.C2CInitMessage"
    21  	MessageTypeNewAgentMessage        = "com.rookout.NewAgentMessage"
    22  	MessageTypeAddAugCommand          = "com.rookout.AddAugCommand"
    23  	MessageTypeRemoveAugCommand       = "com.rookout.RemoveAugCommand"
    24  	MessageTypeAgentInformation       = "com.rookout.AgentInformation"
    25  	MessageTypeAgentInformationArray  = "com.rookout.AgentInformationArray"
    26  	MessageTypeLogMessage             = "com.rookout.LogMessage" 
    27  	MessageTypeControllerLogMessage   = "com.rookout.ControllerLogMessage"
    28  	MessageTypeRuleStatusMessage      = "com.rookout.RuleStatusMessage"
    29  	MessageTypeAugReport              = "com.rookout.AugReportMessage"
    30  	MessageTypeUserMsg                = "com.rookout.UserMsg"
    31  	MessageTypePingMessage            = "com.rookout.PingMessage"
    32  	MessageTypeHitCountUpdateMessage  = "com.rookout.HitCountUpdateMessage"
    33  	MessageTypeAgentsListMessage      = "com.rookout.AgentsList"
    34  	MessageTypeRpcRequest             = "com.rookout.C2cRpcRequest"
    35  	MessageTypeRpcResponse            = "com.rookout.C2cRpcResponse"
    36  	MessageTypeAssertAgents           = "com.rookout.AssertAgents"
    37  	MessageTypeErrorMessage           = "com.rookout.ErrorMessage"
    38  )
    39  
    40  func WrapMsgInEnvelopeWithTime(message proto.Message, t time.Time) ([]byte, error) {
    41  	envelope := &pb.Envelope{}
    42  	msgAsAny, err := AnyProtoMarshal(message)
    43  	envelope.Msg = msgAsAny
    44  	envelope.Timestamp, _ = ptypes.TimestampProto(t)
    45  
    46  	if err != nil {
    47  		return nil, err
    48  	}
    49  
    50  	payload, err := proto.Marshal(envelope)
    51  	if err != nil {
    52  		return nil, fmt.Errorf("failed to marshal envelope (%s) message: %s", msgAsAny.TypeUrl, err.Error())
    53  	}
    54  	return payload, nil
    55  }
    56  
    57  
    58  func WrapMsgInEnvelope(message proto.Message) ([]byte, error) {
    59  	return WrapMsgInEnvelopeWithTime(message, time.Now())
    60  }
    61  
    62  
    63  func ParseEnvelope(buf []byte) (*pb.Envelope, string, error) {
    64  	envelope := &pb.Envelope{}
    65  
    66  	err := proto.Unmarshal(buf, envelope)
    67  	if err != nil {
    68  		return nil, "", err
    69  	}
    70  
    71  	typeName := NormalizeType(envelope.Msg.TypeUrl)
    72  
    73  	return envelope, typeName, nil
    74  }
    75  
    76  func NormalizeType(typeName string) string {
    77  	return strings.Replace(typeName, GoogleProtoTypeURLPrefix, "", 1)
    78  }
    79  
    80  func AnyProtoMarshal(message proto.Message) (*any.Any, error) {
    81  	msgAsAny, err := ptypes.MarshalAny(message)
    82  
    83  	if err != nil {
    84  		return nil, fmt.Errorf("failed to marshal message err - %s", err.Error())
    85  	}
    86  
    87  	return msgAsAny, nil
    88  }
    89  
    90  func AnyProtoUnmarshal(anyMsg *any.Any, msg proto.Message, msgType string) error {
    91  	err := ptypes.UnmarshalAny(anyMsg, msg)
    92  
    93  	if err != nil {
    94  		logrus.WithError(err).Errorf("Failed to convert any.Any type message to: %s msg: %s", msgType, TruncateString(msg.String(), 50))
    95  		return err
    96  	}
    97  
    98  	return nil
    99  }
   100  
   101  func TruncateString(str string, num int) string {
   102  	if num > len(str) {
   103  		return str
   104  	}
   105  	return str[:num]
   106  }
   107  
   108  func ExtractProtoFromEnvelope(envelope *pb.Envelope, msg proto.Message, msgType string) error {
   109  	err := ptypes.UnmarshalAny(envelope.Msg, msg)
   110  	if err != nil {
   111  		logrus.WithError(err).Errorf("Failed to convert envelope message to: %s, msg: %s", msgType, TruncateString(msg.String(), 50))
   112  		return err
   113  	}
   114  	return nil
   115  }