full/all/internal/server/server.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())
}