github.com/Rookout/GoSDK@v0.1.48/pkg/aug_manager/command_handler.go (about)

     1  package aug_manager
     2  
     3  import (
     4  	"encoding/json"
     5  
     6  	"github.com/Rookout/GoSDK/pkg/com_ws"
     7  	"github.com/Rookout/GoSDK/pkg/common"
     8  	"github.com/Rookout/GoSDK/pkg/config"
     9  	"github.com/Rookout/GoSDK/pkg/logger"
    10  	pb "github.com/Rookout/GoSDK/pkg/protobuf"
    11  	"github.com/Rookout/GoSDK/pkg/types"
    12  	"google.golang.org/protobuf/proto"
    13  	"google.golang.org/protobuf/types/known/anypb"
    14  )
    15  
    16  type CommandHandler struct {
    17  	agentCom   com_ws.AgentCom
    18  	augManager AugManager
    19  }
    20  
    21  func NewCommandHandler(agentCom com_ws.AgentCom, augManager AugManager) *CommandHandler {
    22  	handler := CommandHandler{agentCom, augManager}
    23  	handler.agentCom.RegisterCallback(common.MessageTypeInitAugs, func(msg *anypb.Any) {
    24  		initAugs := &pb.InitialAugsCommand{}
    25  		err := anypb.UnmarshalTo(msg, initAugs, proto.UnmarshalOptions{})
    26  		if err != nil {
    27  			logger.Logger().WithError(err).Errorf("failed to unmarshal init augs")
    28  			return
    29  		}
    30  
    31  		config.Update(initAugs.SdkConfiguration)
    32  
    33  		augs, err := handler.buildAugMap(initAugs.Augs)
    34  		if err != nil {
    35  			logger.Logger().WithError(err).Errorf("failed to build rules")
    36  			return
    37  		}
    38  
    39  		handler.augManager.InitializeAugs(augs)
    40  	})
    41  
    42  	handler.agentCom.RegisterCallback(common.MessageTypeRemoveAugCommand, func(msg *anypb.Any) {
    43  		removeAugCmd := &pb.RemoveAugCommand{}
    44  		err := anypb.UnmarshalTo(msg, removeAugCmd, proto.UnmarshalOptions{})
    45  		if err != nil {
    46  			logger.Logger().WithError(err).Error("Failed to unmarshal envelope to RemoveAugCommand")
    47  			return
    48  		}
    49  		err = handler.augManager.RemoveAug(removeAugCmd.AugID)
    50  		if err != nil {
    51  			logger.Logger().WithError(err).Error("failed to remove rule")
    52  		}
    53  	})
    54  
    55  	handler.agentCom.RegisterCallback(common.MessageTypeAddAugCommand, func(msg *anypb.Any) {
    56  		addAugCmd := &pb.AddAugCommand{}
    57  		err := anypb.UnmarshalTo(msg, addAugCmd, proto.UnmarshalOptions{})
    58  		if err != nil {
    59  			logger.Logger().WithError(err).Error("failed to unmarshal envelope to AddAugCommand")
    60  			return
    61  		}
    62  		augConfig := make(types.AugConfiguration)
    63  		err = json.Unmarshal([]byte(addAugCmd.AugJson), &augConfig)
    64  		if err != nil {
    65  			logger.Logger().WithError(err).Error("failed to parse Rule")
    66  			return
    67  		}
    68  		handler.augManager.AddAug(augConfig)
    69  	})
    70  
    71  	handler.agentCom.RegisterCallback(common.MessageTypeErrorMessage, func(msg *anypb.Any) {
    72  		errMsg := &pb.ErrorMessage{}
    73  		err := anypb.UnmarshalTo(msg, errMsg, proto.UnmarshalOptions{})
    74  		if err != nil {
    75  			logger.Logger().WithError(err).Error("failed to unmarshal envelope to ErrorMessage")
    76  			return
    77  		}
    78  		logger.Logger().Warningf("Received Error Message: %s", errMsg.Message)
    79  	})
    80  
    81  	return &handler
    82  }
    83  
    84  func (c *CommandHandler) buildAugMap(rules []string) (map[types.AugID]types.AugConfiguration, error) {
    85  	rulesMap := make(map[types.AugID]types.AugConfiguration)
    86  
    87  	for _, ruleStr := range rules {
    88  		augConfig := make(types.AugConfiguration)
    89  		err := json.Unmarshal([]byte(ruleStr), &augConfig)
    90  		if err != nil {
    91  			logger.Logger().WithError(err).Error("failed to parse aug")
    92  			return nil, err
    93  		}
    94  
    95  		if augID, ok := augConfig["id"].(types.AugID); ok {
    96  			rulesMap[augID] = augConfig
    97  		}
    98  	}
    99  	return rulesMap, nil
   100  }