go-tunnel/client.go

package tunnel

import (
	"context"
	"io"
	"net"
	"net/http"
	"os"
	"strconv"

	"gitlab.com/wpetit/goweb/logger"

	"forge.cadoles.com/wpetit/go-tunnel/control"
	cmap "github.com/orcaman/concurrent-map"
	"github.com/pkg/errors"
	"github.com/xtaci/kcp-go/v5"
	"github.com/xtaci/smux"
)

type Client struct {
	conf    *ClientConfig
	conn    *kcp.UDPSession
	sess    *smux.Session
	control *control.Control
	http    *http.Client
	proxies cmap.ConcurrentMap
}

func (c *Client) Connect(ctx context.Context) error {
	conn, err := kcp.DialWithOptions(
		c.conf.ServerAddress, c.conf.BlockCrypt,
		c.conf.DataShards, c.conf.ParityShards,
	)
	if err != nil {
		return errors.WithStack(err)
	}

	conn.SetWriteDelay(false)

	config := smux.DefaultConfig()
	config.Version = 2

	sess, err := smux.Client(conn, config)
	if err != nil {
		return errors.WithStack(err)
	}

	controlStream, err := sess.OpenStream()
	if err != nil {
		return errors.WithStack(err)
	}

	c.conn = conn
	c.sess = sess
	c.control = control.New(sess, controlStream)

	logger.Debug(ctx, "sending auth request")

	success, err := c.control.AuthRequest(c.conf.Credentials)
	if err != nil {
		return errors.WithStack(err)
	}

	if !success {
		defer c.Close()
		return errors.WithStack(ErrAuthFailed)
	}

	return nil
}

func (c *Client) Listen(ctx context.Context) error {
	logger.Debug(ctx, "listening for messages")

	err := c.control.Listen(ctx, control.Handlers{
		control.TypeProxyRequest: c.handleProxyRequest,
		control.TypeCloseProxy:   c.handleCloseProxy,
	})

	if errors.Is(err, io.ErrClosedPipe) {
		logger.Debug(ctx, "client connection closed")

		return errors.WithStack(ErrConnectionClosed)
	}

	return err
}

func (c *Client) Close() error {
	if c.conn == nil {
		return errors.WithStack(ErrNotConnected)
	}

	if err := c.conn.Close(); err != nil {
		return errors.WithStack(err)
	}

	if c.sess != nil && !c.sess.IsClosed() {
		if err := c.sess.Close(); err != nil {
			return errors.WithStack(err)
		}
	}

	return nil
}

func (c *Client) handleCloseProxy(ctx context.Context, m *control.Message) (*control.Message, error) {
	closeProxyPayload, ok := m.Payload.(*control.CloseProxyPayload)
	if !ok {
		return nil, errors.WithStack(ErrUnexpectedMessage)
	}

	requestID := strconv.FormatInt(closeProxyPayload.RequestID, 10)

	rawCloseChan, exists := c.proxies.Get(requestID)
	if !exists {
		return nil, nil
	}

	closeChan, ok := rawCloseChan.(chan struct{})
	if !ok {
		return nil, nil
	}

	closeChan <- struct{}{}

	return nil, nil
}

func (c *Client) handleProxyRequest(ctx context.Context, m *control.Message) (*control.Message, error) {
	proxyReqPayload, ok := m.Payload.(*control.ProxyRequestPayload)
	if !ok {
		return nil, errors.WithStack(ErrUnexpectedMessage)
	}

	requestID := strconv.FormatInt(proxyReqPayload.RequestID, 10)

	ctx = logger.With(ctx, logger.F("requestID", requestID))

	logger.Debug(
		ctx, "handling proxy request",
		logger.F("network", proxyReqPayload.Network),
		logger.F("address", proxyReqPayload.Address),
	)

	stream, err := c.sess.OpenStream()
	if err != nil {
		return nil, errors.WithStack(err)
	}

	closeChan := make(chan struct{})

	go func() {
		defer func() {
			stream.Close()
			logger.Debug(ctx, "proxy stream closed")
		}()

		proxy := func() error {
			net, err := net.Dial(proxyReqPayload.Network, proxyReqPayload.Address)
			if err != nil {
				return errors.WithStack(err)
			}
			defer net.Close()

			err = pipe(ctx, stream, net)
			if errors.Is(err, os.ErrClosed) {
				return nil
			}

			if err != nil {
				return errors.WithStack(err)
			}

			return nil
		}

		for {
			select {
			case <-closeChan:
				return
			default:
				if err := proxy(); err != nil {
					logger.Error(ctx, "error while proxying", logger.E(err))

					continue
				}

				return
			}
		}
	}()

	c.proxies.Set(requestID, closeChan)

	return nil, nil
}

func NewClient(funcs ...ClientConfigFunc) *Client {
	conf := DefaultClientConfig()

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

	return &Client{
		conf:    conf,
		http:    &http.Client{},
		proxies: cmap.New(),
	}
}