110 lines
2.3 KiB
Go
110 lines
2.3 KiB
Go
package cmux
|
|
|
|
import (
|
|
"fmt"
|
|
"net"
|
|
"time"
|
|
|
|
"quyun/v2/providers/grpc"
|
|
"quyun/v2/providers/http"
|
|
|
|
log "github.com/sirupsen/logrus"
|
|
"github.com/soheilhy/cmux"
|
|
"go.ipao.vip/atom/container"
|
|
"go.ipao.vip/atom/opt"
|
|
"golang.org/x/sync/errgroup"
|
|
)
|
|
|
|
const DefaultPrefix = "Cmux"
|
|
|
|
func DefaultProvider() container.ProviderContainer {
|
|
return container.ProviderContainer{
|
|
Provider: Provide,
|
|
Options: []opt.Option{
|
|
opt.Prefix(DefaultPrefix),
|
|
},
|
|
}
|
|
}
|
|
|
|
type Config struct {
|
|
Host *string
|
|
Port uint
|
|
}
|
|
|
|
func (h *Config) Address() string {
|
|
if h.Host == nil {
|
|
return fmt.Sprintf(":%d", h.Port)
|
|
}
|
|
return fmt.Sprintf("%s:%d", *h.Host, h.Port)
|
|
}
|
|
|
|
type CMux struct {
|
|
Http *http.Service
|
|
Grpc *grpc.Grpc
|
|
Mux cmux.CMux
|
|
Base net.Listener
|
|
}
|
|
|
|
func (c *CMux) Serve() error {
|
|
// Protect against slowloris connections when sniffing protocol
|
|
// Safe even if SetReadTimeout is a no-op in the cmux version in use
|
|
c.Mux.SetReadTimeout(1 * time.Second)
|
|
|
|
addr := ""
|
|
if c.Base != nil && c.Base.Addr() != nil {
|
|
addr = c.Base.Addr().String()
|
|
}
|
|
log.WithFields(log.Fields{
|
|
"addr": addr,
|
|
}).Info("cmux starting")
|
|
|
|
// Route classic HTTP/1.x traffic to the HTTP service
|
|
httpL := c.Mux.Match(cmux.HTTP1Fast())
|
|
|
|
// Route gRPC (HTTP/2 with content-type application/grpc) to the gRPC service.
|
|
// Additionally, send other HTTP/2 traffic to gRPC since Fiber (HTTP) does not serve HTTP/2.
|
|
grpcL := c.Mux.Match(
|
|
cmux.HTTP2HeaderField("content-type", "application/grpc"),
|
|
cmux.HTTP2(),
|
|
)
|
|
|
|
var eg errgroup.Group
|
|
eg.Go(func() error {
|
|
log.WithField("addr", addr).Info("grpc serving via cmux")
|
|
err := c.Grpc.ServeWithListener(grpcL)
|
|
if err != nil {
|
|
log.WithError(err).Error("grpc server exited with error")
|
|
} else {
|
|
log.Info("grpc server exited")
|
|
}
|
|
return err
|
|
})
|
|
|
|
eg.Go(func() error {
|
|
log.WithField("addr", addr).Info("http serving via cmux")
|
|
err := c.Http.Listener(httpL)
|
|
if err != nil {
|
|
log.WithError(err).Error("http server exited with error")
|
|
} else {
|
|
log.Info("http server exited")
|
|
}
|
|
return err
|
|
})
|
|
|
|
// Run cmux dispatcher; wait for the first error from any goroutine
|
|
eg.Go(func() error {
|
|
err := c.Mux.Serve()
|
|
if err != nil {
|
|
log.WithError(err).Error("cmux exited with error")
|
|
} else {
|
|
log.Info("cmux exited")
|
|
}
|
|
return err
|
|
})
|
|
err := eg.Wait()
|
|
if err == nil {
|
|
log.Info("cmux and sub-servers exited cleanly")
|
|
}
|
|
return err
|
|
}
|