-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.go
111 lines (89 loc) · 2.86 KB
/
main.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
package main
import (
"context"
"encoding/json"
"log"
"flag"
"net/http"
"strings"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
"golang.org/x/net/http2"
"golang.org/x/net/http2/h2c"
"github.com/grpc-ecosystem/grpc-gateway/runtime"
pb "github.com/spark8899/go-tag-service/proto"
"github.com/spark8899/go-tag-service/server"
"google.golang.org/grpc"
"google.golang.org/grpc/reflection"
)
var port string
func init() {
flag.StringVar(&port, "port", "8004", "启动端口号")
flag.Parse()
}
func main() {
err := RunServer(port)
if err != nil {
log.Fatalf("Run Serve err: %v", err)
}
}
func RunServer(port string) error {
httpMux := runHttpServer()
grpcS := runGrpcServer()
endpoint := "0.0.0.0:" + port
runtime.HTTPError = grpcGatewayError
gatewayMux := runGrpcGatewayServer()
httpMux.Handle("/", gatewayMux)
return http.ListenAndServe(":"+port, grpcHandlerFunc(grpcS, httpMux))
}
func runHttpServer() *http.ServeMux {
serveMux := http.NewServeMux()
serveMux.HandleFunc("/ping", func(w http.ResponseWriter, r *http.Request) {
_, _ = w.Write([]byte(`pong`))
})
return serveMux
}
func runGrpcServer() *grpc.Server {
s := grpc.NewServer()
pb.RegisterTagServiceServer(s, server.NewTagServer())
reflection.Register(s)
return s
}
func runGrpcGatewayServer() *runtime.ServeMux {
endpoint := "0.0.0.0:" + port
gwmux := runtime.NewServeMux()
dopts := []grpc.DialOption{grpc.WithInsecure()}
_ = pb.RegisterTagServiceHandlerFromEndpoint(context.Background(), gwmux, endpoint, dopts)
return gwmux
}
func grpcHandlerFunc(grpcServer *grpc.Server, otherHandler http.Handler) http.Handler {
return h2c.NewHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if r.ProtoMajor == 2 && strings.Contains(r.Header.Get("Content-Type"), "application/grpc") {
grpcServer.ServeHTTP(w, r)
} else {
otherHandler.ServeHTTP(w, r)
}
}), &http2.Server{})
}
type httpError struct {
Code int32 `json:"code,omitempty"`
Message string `json:"message,omitempty"`
}
func grpcGatewayError(ctx context.Context, _ *runtime.ServeMux, marshaler runtime.Marshaler, w http.ResponseWriter, _ *http.Request, err error) {
s, ok := status.FromError(err)
if !ok {
s = status.New(codes.Unknown, err.Error())
}
httpError := httpError{Code: int32(s.Code()), Message: s.Message()}
details := s.Details()
for _, detail := range details {
if v, ok := detail.(*pb.Error); ok {
httpError.Code = v.Code
httpError.Message = v.Message
}
}
resp, _ := json.Marshal(httpError)
w.Header().Set("Content-type", marshaler.ContentType())
w.WriteHeader(runtime.HTTPStatusFromCode(s.Code()))
_, _ = w.Write(resp)
}