add middleware

cmd/protoc-gen-go-axe/grpc.go

@@ -73,16 +73,28 @@ func generateFileContent(gen *protogen.Plugin, file *protogen.File, g *protogen.
 func genHttpService(gen *protogen.Plugin, file *protogen.File, g *protogen.GeneratedFile, service *protogen.Service) {
 	serverType := service.GoName + "Server"
 
-	g.P("func Register", service.GoName, "HttpServer(s *", httpPackage.Ident("Server"), ",srv ", serverType, ") {")
+	g.P("type MiddlewareFunc func(", httpPackage.Ident("Handler"), ") ", httpPackage.Ident("Handler"))
+	g.P()
 
+	g.P("func Register", service.GoName, "HttpServer(s *", httpPackage.Ident("Server"), ",srv ", serverType, ", middlewares ...MiddlewareFunc) {")
 	g.P("mux := ", httpPackage.Ident("NewServeMux"), "()")
 	g.P()
 
 	for _, method := range service.Methods {
 		hname := genHttpServerMethod(gen, file, g, method)
-		g.P("mux.HandleFunc(", strconv.Quote(fmt.Sprintf("/%s/%s", service.Desc.FullName(), method.Desc.Name())), ", ", hname, ")")
+
+		hf := fmt.Sprintf("_%s_%s_HandlerFunc", service.GoName, method.GoName)
+		g.P("var ", hf, " ", httpPackage.Ident("Handler"))
+
+		g.P(hf, " = ", httpPackage.Ident("Handler"), "(", httpPackage.Ident("HandlerFunc"), "(", hname, "))")
+		g.P("for _, m := range middlewares {")
+		g.P(hf, " = m(", hf, ")")
+		g.P("}")
+
+		g.P("mux.Handle(", strconv.Quote(fmt.Sprintf("/%s/%s", service.Desc.FullName(), method.Desc.Name())), ", ", hf, ")")
 		g.P()
 	}
+
 	g.P("s.Handler = mux")
 	g.P("}")
 	g.P()

demo/main.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"encoding/json"
 	"log"
+	"net/http"
 	"path"
 	"runtime"
 	"time"
@@ -18,7 +19,27 @@ type echoServer struct {
 }
 
 func (s *echoServer) Echo(ctx context.Context, req *pb.EchoRequest) (*pb.EchoResponse, error) {
-	return &pb.EchoResponse{Value: req.GetValue()}, nil
+	return &pb.EchoResponse{Value: "echo1_" + req.GetValue()}, nil
+}
+
+func (s *echoServer) Echo2(ctx context.Context, req *pb.EchoRequest) (*pb.EchoResponse, error) {
+	return &pb.EchoResponse{Value: "echo2_" + req.GetValue()}, nil
+}
+
+func middleware1(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		log.Println("middleware1 - Before Handler")
+		next.ServeHTTP(w, r)
+		log.Println("middleware1 - After Handler")
+	})
+}
+
+func middleware2(next http.Handler) http.Handler {
+	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		log.Println("middleware2 - Before Handler")
+		next.ServeHTTP(w, r)
+		log.Println("middleware2 - After Handler")
+	})
 }
 
 /////////////////////////////
@@ -56,7 +77,7 @@ func main() {
 	// register rpc
 	pb.RegisterEchoServiceServer(s.GrpcServer(), &echoServer{})
 	// register http, pattern和handler会自动生成
-	pb.RegisterEchoServiceHttpServer(s.HttpServer(), &echoServer{})
+	pb.RegisterEchoServiceHttpServer(s.HttpServer(), &echoServer{}, []pb.MiddlewareFunc{middleware1, middleware2}...)
 
 	// 调用client的例子
 	//go clientExample()

demo/pb/echo.proto

@@ -12,4 +12,5 @@ message EchoResponse{
 
 service EchoService{
     rpc Echo(EchoRequest) returns (EchoResponse){}
+    rpc Echo2(EchoRequest) returns (EchoResponse){}
 }