From aa98f5fd5525e2ab8a66000e678312fcfd651f48 Mon Sep 17 00:00:00 2001
From: GLYASAI <chaoschunqfeijg@163.com>
Date: Wed, 21 Nov 2018 20:14:13 +0800
Subject: [PATCH] [ADD] add api for updating tcp rule

---
 api/controller/gateway.go      | 16 +++++++++
 api/handler/gateway.go         | 63 +++++++++++++++++++++++++++++-----
 api/handler/gateway_handler.go |  1 +
 db/mysql/dao/gateway.go        | 11 ++++--
 4 files changed, 81 insertions(+), 10 deletions(-)

diff --git a/api/controller/gateway.go b/api/controller/gateway.go
index 2e67d4fa4..8f59a3f3f 100644
--- a/api/controller/gateway.go
+++ b/api/controller/gateway.go
@@ -201,7 +201,23 @@ func (g *GatewayStruct) addTcpRule(w http.ResponseWriter, r *http.Request) {
 }
 
 func (g *GatewayStruct) updateTcpRule(w http.ResponseWriter, r *http.Request) {
+	logrus.Debugf("add tcp rule.")
+	var req api_model.TcpRuleStruct
+	ok := httputil.ValidatorRequestStructAndErrorResponse(r, w, &req, nil)
+	if !ok {
+		return
+	}
+	reqJson, _ := json.Marshal(req)
+	logrus.Debugf("Request is : %s", string(reqJson))
 
+	h := handler.GetGatewayHandler()
+	if err := h.UpdateTcpRule(&req); err != nil {
+		httputil.ReturnError(r, w, 500, fmt.Sprintf("Unexpected error occorred while " +
+			"updating tcp rule: %v", err))
+		return
+	}
+
+	httputil.ReturnSuccess(r, w, "success")
 }
 
 func (g *GatewayStruct) deleteTcpRule(w http.ResponseWriter, r *http.Request) {
diff --git a/api/handler/gateway.go b/api/handler/gateway.go
index cbde714a5..cf98419e5 100644
--- a/api/handler/gateway.go
+++ b/api/handler/gateway.go
@@ -51,7 +51,7 @@ func (g *GatewayAction) UpdateHttpRule(req *apimodel.HttpRuleStruct,
 		return nil, err
 	}
 	if rule == nil {
-		return nil, fmt.Errorf("HttpRule dosen't exist based on ServiceID(%s) " +
+		return nil, fmt.Errorf("HttpRule dosen't exist based on ServiceID(%s) "+
 			"and ContainerPort(%v)", req.ServiceID, req.ContainerPort)
 	}
 	// delete old Certificate
@@ -137,11 +137,11 @@ func (g *GatewayAction) UpdateCertificate(req apimodel.HttpRuleStruct, httpRule
 // AddTcpRule adds tcp rule.
 func (g *GatewayAction) AddTcpRule(req *apimodel.TcpRuleStruct) error {
 	tcpRule := &model.TcpRule{
-		UUID: util.NewUUID(),
-		ServiceID: req.ServiceID,
-		ContainerPort: req.ContainerPort,
-		IP: req.IP,
-		Port: req.Port,
+		UUID:             util.NewUUID(),
+		ServiceID:        req.ServiceID,
+		ContainerPort:    req.ContainerPort,
+		IP:               req.IP,
+		Port:             req.Port,
 		LoadBalancerType: req.LoadBalancerType,
 	}
 
@@ -149,17 +149,19 @@ func (g *GatewayAction) AddTcpRule(req *apimodel.TcpRuleStruct) error {
 	tx := db.GetManager().Begin()
 	// add tcp rule
 	if err := g.dbmanager.TcpRuleDaoTransactions(tx).AddModel(tcpRule); err != nil {
+		tx.Rollback()
 		return err
 	}
 
 	// add rule extensions
 	for _, ruleExtension := range req.RuleExtensions {
 		re := &model.RuleExtension{
-			UUID: util.NewUUID(),
+			UUID:   util.NewUUID(),
 			RuleID: tcpRule.UUID,
-			Value: ruleExtension.Value,
+			Value:  ruleExtension.Value,
 		}
 		if err := g.dbmanager.RuleExtensionDaoTransactions(tx).AddModel(re); err != nil {
+			tx.Rollback()
 			return err
 		}
 	}
@@ -172,6 +174,51 @@ func (g *GatewayAction) AddTcpRule(req *apimodel.TcpRuleStruct) error {
 	return nil
 }
 
+func (g *GatewayAction) UpdateTcpRule(req *apimodel.TcpRuleStruct) error {
+	// begin transaction
+	tx := db.GetManager().Begin()
+	// get old tcp rule
+	tcpRule, err := g.dbmanager.TcpRuleDaoTransactions(tx).GetTcpRuleByServiceIDAndContainerPort(req.ServiceID,
+		req.ContainerPort)
+	if err != nil {
+		tx.Rollback()
+		return err
+	}
+	// delete old rule extensions
+	if err := g.dbmanager.RuleExtensionDaoTransactions(tx).DeleteRuleExtensionByRuleID(tcpRule.UUID); err != nil {
+		tx.Rollback()
+		return err
+	}
+	// update tcp rule
+	tcpRule.IP = req.IP
+	tcpRule.Port = req.Port
+	tcpRule.LoadBalancerType = req.LoadBalancerType
+	if err := g.dbmanager.TcpRuleDaoTransactions(tx).UpdateModel(tcpRule); err != nil {
+		tx.Rollback()
+		return err
+	}
+	// add new rule extensions
+	for _, ruleExtension := range req.RuleExtensions {
+		re := &model.RuleExtension{
+			UUID:   util.NewUUID(),
+			RuleID: tcpRule.UUID,
+			Value:  ruleExtension.Value,
+		}
+		if err := g.dbmanager.RuleExtensionDaoTransactions(tx).AddModel(re); err != nil {
+			tx.Rollback()
+			return err
+		}
+	}
+
+	// end transaction
+	if err := tx.Commit().Error; err != nil {
+		tx.Rollback()
+		return err
+	}
+
+	return nil
+}
+
 // AddRuleExtensions adds rule extensions to db if any of they doesn't exists
 func (g *GatewayAction) AddRuleExtensions(ruleID string, ruleExtensions []*apimodel.RuleExtensionStruct,
 	tx *gorm.DB) error {
diff --git a/api/handler/gateway_handler.go b/api/handler/gateway_handler.go
index 8f3630870..efcbe4d77 100644
--- a/api/handler/gateway_handler.go
+++ b/api/handler/gateway_handler.go
@@ -33,6 +33,7 @@ type GatewayHandler interface {
 	UpdateCertificate(req apimodel.HttpRuleStruct, httpRule *dbmodel.HttpRule, tx *gorm.DB) error
 
 	AddTcpRule(req *apimodel.TcpRuleStruct) error
+	UpdateTcpRule(req *apimodel.TcpRuleStruct) error
 
 	AddRuleExtensions(ruleID string, ruleExtensions []*apimodel.RuleExtensionStruct, tx *gorm.DB) error
 }
diff --git a/db/mysql/dao/gateway.go b/db/mysql/dao/gateway.go
index 1b36550f8..5c3b121ea 100644
--- a/db/mysql/dao/gateway.go
+++ b/db/mysql/dao/gateway.go
@@ -196,8 +196,15 @@ func (t *TcpRuleDaoTmpl) AddModel(mo model.Interface) error {
 	return nil
 }
 
-func (s *TcpRuleDaoTmpl) UpdateModel(model.Interface) error {
-	return nil
+func (t *TcpRuleDaoTmpl) UpdateModel(mo model.Interface) error {
+	tr, ok := mo.(*model.TcpRule)
+	if !ok {
+		return fmt.Errorf("Failed to convert %s to *model.TcpRule", reflect.TypeOf(mo).String())
+	}
+
+	return t.DB.Table(tr.TableName()).
+		Where("uuid = ?", tr.UUID).
+		Update(tr).Error
 }
 
 // GetTcpRuleByServiceIDAndContainerPort gets a TcpRule based on serviceID and containerPort