128 lines
3.0 KiB
Go
128 lines
3.0 KiB
Go
package server
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"net"
|
|
"net/http"
|
|
"time"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
gwRuntime "github.com/grpc-ecosystem/grpc-gateway/v2/runtime"
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/reflection"
|
|
)
|
|
|
|
type Server struct {
|
|
GRPC *grpc.Server
|
|
Gateway *gwRuntime.ServeMux
|
|
HTTP *gin.Engine
|
|
http *http.Server
|
|
dynamic *dynamicGateway
|
|
}
|
|
|
|
func New() *Server {
|
|
grpcServer := grpc.NewServer()
|
|
reflection.Register(grpcServer)
|
|
engine := gin.New()
|
|
engine.Use(gin.Logger(), gin.Recovery())
|
|
return &Server{
|
|
GRPC: grpcServer,
|
|
Gateway: gwRuntime.NewServeMux(),
|
|
HTTP: engine,
|
|
}
|
|
}
|
|
|
|
func (s *Server) Start(grpcAddr, httpAddr string, allow []string) error {
|
|
grpcListener, err := net.Listen("tcp", grpcAddr)
|
|
if err != nil {
|
|
return fmt.Errorf("listen gRPC on %s: %w", grpcAddr, err)
|
|
}
|
|
httpListener, err := net.Listen("tcp", httpAddr)
|
|
if err != nil {
|
|
_ = grpcListener.Close()
|
|
return fmt.Errorf("listen HTTP on %s: %w", httpAddr, err)
|
|
}
|
|
|
|
s.dynamic, err = newDynamicGateway(grpcAddr, allow)
|
|
if err != nil {
|
|
_ = grpcListener.Close()
|
|
_ = httpListener.Close()
|
|
return err
|
|
}
|
|
s.HTTP.POST("/rpc/:method", s.dynamic.handle)
|
|
|
|
handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
recorder := newBufferedResponse()
|
|
s.Gateway.ServeHTTP(recorder, r)
|
|
if recorder.status != http.StatusNotFound {
|
|
recorder.flush(w)
|
|
return
|
|
}
|
|
s.HTTP.ServeHTTP(w, r)
|
|
})
|
|
s.http = &http.Server{
|
|
Addr: httpAddr,
|
|
Handler: handler,
|
|
ReadHeaderTimeout: 10 * time.Second,
|
|
IdleTimeout: 120 * time.Second,
|
|
MaxHeaderBytes: 1 << 20,
|
|
}
|
|
fmt.Printf("all gRPC services listening on %s\n", grpcAddr)
|
|
fmt.Printf("all HTTP services listening on %s\n", httpAddr)
|
|
|
|
errCh := make(chan error, 2)
|
|
go func() { errCh <- s.GRPC.Serve(grpcListener) }()
|
|
go func() { errCh <- s.http.Serve(httpListener) }()
|
|
serveErr := <-errCh
|
|
if errors.Is(serveErr, grpc.ErrServerStopped) || errors.Is(serveErr, http.ErrServerClosed) {
|
|
return http.ErrServerClosed
|
|
}
|
|
return serveErr
|
|
}
|
|
|
|
func (s *Server) Stop(ctx context.Context) error {
|
|
stopped := make(chan struct{})
|
|
go func() {
|
|
s.GRPC.GracefulStop()
|
|
close(stopped)
|
|
}()
|
|
select {
|
|
case <-stopped:
|
|
case <-ctx.Done():
|
|
s.GRPC.Stop()
|
|
}
|
|
if s.dynamic != nil {
|
|
_ = s.dynamic.Close()
|
|
}
|
|
if s.http == nil {
|
|
return nil
|
|
}
|
|
return s.http.Shutdown(ctx)
|
|
}
|
|
|
|
type bufferedResponse struct {
|
|
header http.Header
|
|
body bytes.Buffer
|
|
status int
|
|
}
|
|
|
|
func newBufferedResponse() *bufferedResponse {
|
|
return &bufferedResponse{header: make(http.Header), status: http.StatusOK}
|
|
}
|
|
|
|
func (r *bufferedResponse) Header() http.Header { return r.header }
|
|
func (r *bufferedResponse) WriteHeader(status int) { r.status = status }
|
|
func (r *bufferedResponse) Write(data []byte) (int, error) { return r.body.Write(data) }
|
|
func (r *bufferedResponse) flush(w http.ResponseWriter) {
|
|
for key, values := range r.header {
|
|
for _, value := range values {
|
|
w.Header().Add(key, value)
|
|
}
|
|
}
|
|
w.WriteHeader(r.status)
|
|
_, _ = w.Write(r.body.Bytes())
|
|
}
|