atune: upgrade, define, undefine, update can not be remotely operated in order...

atune: upgrade, define, undefine, update can not be remotely operated in order to prevent remote attacks

common/service/service.go

@@ -14,7 +14,12 @@
 package module
 
 import (
+	"atune/common/config"
+	"context"
 	"fmt"
+	"google.golang.org/grpc/peer"
+	"net"
+	"strings"
 	"sync"
 
 	"github.com/urfave/cli"
@@ -115,3 +120,33 @@ func GetServices() ([]string, error) {
 
 	return svc, nil
 }
+
+//CheckRpcIsLocalAddr return whether rpc is a local addr
+func CheckRpcIsLocalAddr(ctx context.Context) (bool, error) {
+	if config.TransProtocol != "tcp" {
+		return true, nil
+	}
+
+	pr, ok := peer.FromContext(ctx)
+	if !ok || pr.Addr == net.Addr(nil) {
+		return false, fmt.Errorf("failed to get rpc client ip")
+	}
+	clientAddr := strings.Split(pr.Addr.String(), ":")[0]
+
+	serverAddr, err := net.InterfaceAddrs()
+	if err != nil {
+		return false, err
+	}
+
+	for i := range serverAddr {
+		addr, _, err := net.ParseCIDR(serverAddr[i].String())
+		if err != nil {
+			return false, err
+		}
+		if net.ParseIP(clientAddr).Equal(addr) {
+			return true, nil
+		}
+	}
+	return false, nil
+}
+

modules/client/profile/profile_define.go

@@ -104,7 +104,6 @@ func profileDefined(ctx *cli.Context) error {
 		ProfileName:  profileName,
 		Content:      data})
 	if err != nil {
-		fmt.Println(err)
 		return err
 	}
 	if reply.GetStatus() != "OK" {

modules/client/profile/profile_undefine.go

@@ -74,7 +74,6 @@ func profileDelete(ctx *cli.Context) error {
 	svc := PB.NewProfileMgrClient(c.Connection())
 	reply, err := svc.Delete(CTX.Background(), &PB.DefineMessage{WorkloadType: workloadType})
 	if err != nil {
-		fmt.Println(err)
 		return err
 	}
 	if reply.GetStatus() != "OK" {

modules/client/profile/profile_update.go

@@ -96,8 +96,7 @@ func profileUpdate(ctx *cli.Context) error {
 		ProfileName:  profileName,
 		Content:      data})
 	if err != nil {
-		fmt.Println(err)
-		return nil
+		return err
 	}
 	if reply.GetStatus() != "OK" {
 		fmt.Println(reply.GetStatus())

modules/server/profile/profile.go

@@ -511,6 +511,14 @@ func (s *ProfileServer) Tuning(stream PB.ProfileMgr_TuningServer) error {
 UpgradeProfile method update the db file
 */
 func (s *ProfileServer) UpgradeProfile(profileInfo *PB.ProfileInfo, stream PB.ProfileMgr_UpgradeProfileServer) error {
+	isLocalAddr, err := SVC.CheckRpcIsLocalAddr(stream.Context())
+	if err != nil {
+		return err
+	}
+	if !isLocalAddr {
+		return fmt.Errorf("the upgrade command can not be remotely operated")
+	}
+
 	log.Debug("Begin to upgrade profiles\n")
 	currenDbPath := path.Join(config.DatabasePath, config.DatabaseName)
 	newDbPath := profileInfo.GetName()
@@ -880,6 +888,14 @@ func (s *ProfileServer) Charaterization(profileInfo *PB.ProfileInfo, stream PB.P
 
 // Define method user define workload type and profile
 func (s *ProfileServer) Define(ctx context.Context, message *PB.DefineMessage) (*PB.Ack, error) {
+	isLocalAddr, err := SVC.CheckRpcIsLocalAddr(ctx)
+	if err != nil {
+		return &PB.Ack{}, err
+	}
+	if !isLocalAddr {
+		return &PB.Ack{}, fmt.Errorf("the define command can not be remotely operated")
+	}
+
 	workloadType := message.GetWorkloadType()
 	profileName := message.GetProfileName()
 	content := string(message.GetContent())
@@ -924,10 +940,18 @@ func (s *ProfileServer) Define(ctx context.Context, message *PB.DefineMessage) (
 
 // Delete method delete the self define workload type from database
 func (s *ProfileServer) Delete(ctx context.Context, message *PB.DefineMessage) (*PB.Ack, error) {
+	isLocalAddr, err := SVC.CheckRpcIsLocalAddr(ctx)
+	if err != nil {
+		return &PB.Ack{}, err
+	}
+	if !isLocalAddr {
+		return &PB.Ack{}, fmt.Errorf("the undefine command can not be remotely operated")
+	}
+
 	workloadType := message.GetWorkloadType()
 
 	classApps := &sqlstore.GetClassApp{Class: workloadType}
-	err := sqlstore.GetClassApps(classApps)
+	err = sqlstore.GetClassApps(classApps)
 	if err != nil {
 		return &PB.Ack{}, err
 	}
@@ -981,6 +1005,14 @@ func (s *ProfileServer) Delete(ctx context.Context, message *PB.DefineMessage) (
 
 // Update method update the content of the specified workload type from database
 func (s *ProfileServer) Update(ctx context.Context, message *PB.DefineMessage) (*PB.Ack, error) {
+	isLocalAddr, err := SVC.CheckRpcIsLocalAddr(ctx)
+	if err != nil {
+		return &PB.Ack{}, err
+	}
+	if !isLocalAddr {
+		return &PB.Ack{}, fmt.Errorf("the update command can not be remotely operated")
+	}
+
 	workloadType := message.GetWorkloadType()
 	profileName := message.GetProfileName()
 	content := string(message.GetContent())