go-tunnel/server.go

package tunnel

import (
	"context"

	cmap "github.com/orcaman/concurrent-map"
	"github.com/pkg/errors"
	"github.com/xtaci/kcp-go/v5"
	"gitlab.com/wpetit/goweb/logger"
)

type Server struct {
	conf    *ServerConfig
	clients cmap.ConcurrentMap
}

func (s *Server) Listen(ctx context.Context) error {
	listener, err := kcp.ListenWithOptions(
		s.conf.Address, s.conf.BlockCrypt,
		s.conf.DataShards, s.conf.ParityShards,
	)
	if err != nil {
		return errors.WithStack(err)
	}

	for {
		conn, err := listener.AcceptKCP()
		if err != nil {
			return errors.WithStack(err)
		}

		go s.handleNewConn(ctx, conn)
	}
}

func (s *Server) handleNewConn(ctx context.Context, conn *kcp.UDPSession) {
	ctx = logger.With(ctx, logger.F("remoteAddr", conn.RemoteAddr().String()))

	remoteClient := NewRemoteClient()

	defer remoteClient.Close()
	defer conn.Close()

	remoteClient.ConfigureHooks(s.conf.Hooks)

	if err := remoteClient.Accept(ctx, conn); err != nil {
		logger.Error(ctx, "remote client error", logger.E(errors.WithStack(err)))

		return
	}

	if err := remoteClient.Listen(ctx); err != nil {
		logger.Error(ctx, "remote client error", logger.E(errors.WithStack(err)))
	}
}

func NewServer(funcs ...ServerConfigFunc) *Server {
	conf := DefaultServerConfig()

	for _, fn := range funcs {
		fn(conf)
	}

	return &Server{
		conf:    conf,
		clients: cmap.New(),
	}
}