initial commit

listener/grpc/listener.go

@@ -0,0 +1,100 @@
+package grpc
+
+import (
+	"net"
+
+	"github.com/go-gost/gost/v3/pkg/common/admission"
+	"github.com/go-gost/gost/v3/pkg/common/metrics"
+	"github.com/go-gost/gost/v3/pkg/listener"
+	"github.com/go-gost/gost/v3/pkg/logger"
+	md "github.com/go-gost/gost/v3/pkg/metadata"
+	"github.com/go-gost/gost/v3/pkg/registry"
+	pb "github.com/go-gost/x/internal/util/grpc/proto"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials"
+)
+
+func init() {
+	registry.ListenerRegistry().Register("grpc", NewListener)
+}
+
+type grpcListener struct {
+	addr    net.Addr
+	server  *grpc.Server
+	cqueue  chan net.Conn
+	errChan chan error
+	md      metadata
+	logger  logger.Logger
+	options listener.Options
+}
+
+func NewListener(opts ...listener.Option) listener.Listener {
+	options := listener.Options{}
+	for _, opt := range opts {
+		opt(&options)
+	}
+	return &grpcListener{
+		logger:  options.Logger,
+		options: options,
+	}
+}
+
+func (l *grpcListener) Init(md md.Metadata) (err error) {
+	if err = l.parseMetadata(md); err != nil {
+		return
+	}
+
+	ln, err := net.Listen("tcp", l.options.Addr)
+	if err != nil {
+		return
+	}
+	ln = metrics.WrapListener(l.options.Service, ln)
+	ln = admission.WrapListener(l.options.Admission, ln)
+
+	var opts []grpc.ServerOption
+	if !l.md.insecure {
+		opts = append(opts, grpc.Creds(credentials.NewTLS(l.options.TLSConfig)))
+	}
+
+	l.server = grpc.NewServer(opts...)
+	l.addr = ln.Addr()
+	l.cqueue = make(chan net.Conn, l.md.backlog)
+	l.errChan = make(chan error, 1)
+
+	pb.RegisterGostTunelServer(l.server, &server{
+		cqueue:    l.cqueue,
+		localAddr: l.addr,
+		logger:    l.options.Logger,
+	})
+
+	go func() {
+		err := l.server.Serve(ln)
+		if err != nil {
+			l.errChan <- err
+		}
+		close(l.errChan)
+	}()
+
+	return
+}
+
+func (l *grpcListener) Accept() (conn net.Conn, err error) {
+	var ok bool
+	select {
+	case conn = <-l.cqueue:
+	case err, ok = <-l.errChan:
+		if !ok {
+			err = listener.ErrClosed
+		}
+	}
+	return
+}
+
+func (l *grpcListener) Close() error {
+	l.server.Stop()
+	return nil
+}
+
+func (l *grpcListener) Addr() net.Addr {
+	return l.addr
+}

listener/grpc/metadata.go

@@ -0,0 +1,29 @@
+package grpc
+
+import (
+	mdata "github.com/go-gost/gost/v3/pkg/metadata"
+)
+
+const (
+	defaultBacklog = 128
+)
+
+type metadata struct {
+	backlog  int
+	insecure bool
+}
+
+func (l *grpcListener) parseMetadata(md mdata.Metadata) (err error) {
+	const (
+		backlog  = "backlog"
+		insecure = "grpcInsecure"
+	)
+
+	l.md.backlog = mdata.GetInt(md, backlog)
+	if l.md.backlog <= 0 {
+		l.md.backlog = defaultBacklog
+	}
+
+	l.md.insecure = mdata.GetBool(md, insecure)
+	return
+}

listener/grpc/server.go

@@ -0,0 +1,124 @@
+package grpc
+
+import (
+	"errors"
+	"io"
+	"net"
+	"time"
+
+	"github.com/go-gost/gost/v3/pkg/logger"
+	pb "github.com/go-gost/x/internal/util/grpc/proto"
+	"google.golang.org/grpc/peer"
+)
+
+type server struct {
+	cqueue    chan net.Conn
+	localAddr net.Addr
+	pb.UnimplementedGostTunelServer
+	logger logger.Logger
+}
+
+func (s *server) Tunnel(srv pb.GostTunel_TunnelServer) error {
+	c := &conn{
+		s:          srv,
+		localAddr:  s.localAddr,
+		remoteAddr: &net.TCPAddr{},
+		closed:     make(chan struct{}),
+	}
+	if p, ok := peer.FromContext(srv.Context()); ok {
+		c.remoteAddr = p.Addr
+	}
+
+	select {
+	case s.cqueue <- c:
+	default:
+		c.Close()
+		s.logger.Warnf("connection queue is full, client discarded")
+	}
+
+	<-c.closed
+
+	return nil
+}
+
+type conn struct {
+	s          pb.GostTunel_TunnelServer
+	rb         []byte
+	localAddr  net.Addr
+	remoteAddr net.Addr
+	closed     chan struct{}
+}
+
+func (c *conn) Read(b []byte) (n int, err error) {
+	select {
+	case <-c.s.Context().Done():
+		err = c.s.Context().Err()
+		return
+	case <-c.closed:
+		err = io.ErrClosedPipe
+		return
+	default:
+	}
+
+	if len(c.rb) == 0 {
+		chunk, err := c.s.Recv()
+		if err != nil {
+			return 0, err
+		}
+		c.rb = chunk.Data
+	}
+
+	n = copy(b, c.rb)
+	c.rb = c.rb[n:]
+	return
+}
+
+func (c *conn) Write(b []byte) (n int, err error) {
+	select {
+	case <-c.s.Context().Done():
+		err = c.s.Context().Err()
+		return
+	case <-c.closed:
+		err = io.ErrClosedPipe
+		return
+	default:
+	}
+
+	if err = c.s.Send(&pb.Chunk{
+		Data: b,
+	}); err != nil {
+		return
+	}
+	n = len(b)
+	return
+}
+
+func (c *conn) Close() error {
+	select {
+	case <-c.closed:
+	default:
+		close(c.closed)
+	}
+
+	return nil
+}
+
+func (c *conn) LocalAddr() net.Addr {
+	return c.localAddr
+}
+
+func (c *conn) RemoteAddr() net.Addr {
+	return c.remoteAddr
+}
+
+func (c *conn) SetDeadline(t time.Time) error {
+	return &net.OpError{Op: "set", Net: "grpc", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
+}
+
+func (c *conn) SetReadDeadline(t time.Time) error {
+	return &net.OpError{Op: "set", Net: "grpc", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
+}
+
+func (c *conn) SetWriteDeadline(t time.Time) error {
+	return &net.OpError{Op: "set", Net: "grpc", Source: nil, Addr: nil, Err: errors.New("deadline not supported")}
+}