a broken AP change: extend PreWriteResponse

Author: smallnest

server/gateway.go

@@ -80,7 +80,7 @@ func (s *Server) handleGatewayRequest(w http.ResponseWriter, r *http.Request, pa
 	ctx := context.WithValue(context.Background(), StartRequestContextKey, time.Now().UnixNano())
 	err = s.auth(ctx, req)
 	if err != nil {
-		s.Plugins.DoPreWriteResponse(ctx, req)
+		s.Plugins.DoPreWriteResponse(ctx, req, nil)
 		wh.Set(XMessageStatusType, "Error")
 		wh.Set(XErrorMessage, err.Error())
 		w.WriteHeader(401)
@@ -103,7 +103,7 @@ func (s *Server) handleGatewayRequest(w http.ResponseWriter, r *http.Request, pa
 		return
 	}
 
-	s.Plugins.DoPreWriteResponse(newCtx, req)
+	s.Plugins.DoPreWriteResponse(newCtx, req, nil)
 	if len(resMetadata) > 0 { //copy meta in context to request
 		meta := res.Metadata
 		if meta == nil {

server/plugin.go

@@ -23,7 +23,7 @@ type PluginContainer interface {
 	DoPreReadRequest(ctx context.Context) error
 	DoPostReadRequest(ctx context.Context, r *protocol.Message, e error) error
 
-	DoPreWriteResponse(context.Context, *protocol.Message) error
+	DoPreWriteResponse(context.Context, *protocol.Message, *protocol.Message) error
 	DoPostWriteResponse(context.Context, *protocol.Message, *protocol.Message, error) error
 
 	DoPreWriteRequest(ctx context.Context) error
@@ -64,7 +64,7 @@ type (
 
 	//PreWriteResponsePlugin represents .
 	PreWriteResponsePlugin interface {
-		PreWriteResponse(context.Context, *protocol.Message) error
+		PreWriteResponse(context.Context, *protocol.Message, *protocol.Message) error
 	}
 
 	//PostWriteResponsePlugin represents .
@@ -193,10 +193,10 @@ func (p *pluginContainer) DoPostReadRequest(ctx context.Context, r *protocol.Mes
 }
 
 // DoPreWriteResponse invokes PreWriteResponse plugin.
-func (p *pluginContainer) DoPreWriteResponse(ctx context.Context, req *protocol.Message) error {
+func (p *pluginContainer) DoPreWriteResponse(ctx context.Context, req *protocol.Message, res *protocol.Message) error {
 	for i := range p.plugins {
 		if plugin, ok := p.plugins[i].(PreWriteResponsePlugin); ok {
-			err := plugin.PreWriteResponse(ctx, req)
+			err := plugin.PreWriteResponse(ctx, req, res)
 			if err != nil {
 				return err
 			}

server/server.go

@@ -314,15 +314,18 @@ func (s *Server) serveConn(conn net.Conn) {
 		ctx = context.WithValue(ctx, StartRequestContextKey, time.Now().UnixNano())
 		err = s.auth(ctx, req)
 		if err != nil {
-			s.Plugins.DoPreWriteResponse(ctx, req)
+
 			if !req.IsOneway() {
 				res := req.Clone()
 				res.SetMessageType(protocol.Response)
 				handleError(res, err)
 				data := res.Encode()
+				s.Plugins.DoPreWriteResponse(ctx, req, res)
 				conn.Write(data)
 				s.Plugins.DoPostWriteResponse(ctx, req, res, err)
 				protocol.FreeMsg(res)
+			} else {
+				s.Plugins.DoPreWriteResponse(ctx, req, nil)
 			}
 
 			protocol.FreeMsg(req)
@@ -346,7 +349,7 @@ func (s *Server) serveConn(conn net.Conn) {
 				log.Warnf("rpcx: failed to handle request: %v", err)
 			}
 
-			s.Plugins.DoPreWriteResponse(newCtx, req)
+			s.Plugins.DoPreWriteResponse(newCtx, req, res)
 			if !req.IsOneway() {
 				if len(resMetadata) > 0 { //copy meta in context to request
 					meta := res.Metadata

serverplugin/alias.go

@@ -59,7 +59,7 @@ func (p *AliasPlugin) PostReadRequest(ctx context.Context, r *protocol.Message,
 }
 
 // PreWriteResponse restore servicePath and serviceMethod.
-func (p *AliasPlugin) PreWriteResponse(ctx context.Context, r *protocol.Message) error {
+func (p *AliasPlugin) PreWriteResponse(ctx context.Context, r *protocol.Message, res *protocol.Message) error {
 	if r.Metadata[aliasAppliedKey] != "true" {
 		return nil
 	}
@@ -72,6 +72,10 @@ func (p *AliasPlugin) PreWriteResponse(ctx context.Context, r *protocol.Message)
 			r.ServicePath = pm.servicePath
 			r.ServiceMethod = pm.serviceMethod
 			delete(r.Metadata, aliasAppliedKey)
+			if res != nil {
+				res.ServicePath = pm.servicePath
+				res.ServiceMethod = pm.serviceMethod
+			}
 		}
 	}
 	return nil