mtg/proxy/middle.go

package proxy

import (
	"sync"

	"github.com/9seconds/mtg/conntypes"
	"github.com/9seconds/mtg/protocol"
	"github.com/9seconds/mtg/wrappers/packetack"
	"go.uber.org/zap"
)

func middleConnection(request *protocol.TelegramRequest) {
	telegramConn, err := packetack.NewProxy(request)
	if err != nil {
		request.Logger.Debugw("Cannot dial to Telegram", "error", err)

		return
	}
	defer telegramConn.Close()

	var clientConn conntypes.PacketAckFullReadWriteCloser

	switch request.ClientProtocol.ConnectionType() {
	case conntypes.ConnectionTypeAbridged:
		clientConn = packetack.NewClientAbridged(request.ClientConn)
	case conntypes.ConnectionTypeIntermediate:
		clientConn = packetack.NewClientIntermediate(request.ClientConn)
	case conntypes.ConnectionTypeSecure:
		clientConn = packetack.NewClientIntermediateSecure(request.ClientConn)
	case conntypes.ConnectionTypeUnknown:
		panic("unknown connection type")
	}

	wg := &sync.WaitGroup{}
	wg.Add(2)

	go middlePipe(telegramConn, clientConn, wg, request.Logger)

	go middlePipe(clientConn, telegramConn, wg, request.Logger)

	wg.Wait()
}

func middlePipe(dst conntypes.PacketAckWriteCloser,
	src conntypes.PacketAckReadCloser,
	wg *sync.WaitGroup,
	logger *zap.SugaredLogger) {
	defer func() {
		dst.Close()
		src.Close()
		wg.Done()
	}()

	for {
		acks := conntypes.ConnectionAcks{}

		packet, err := src.Read(&acks)
		if err != nil {
			logger.Debugw("Cannot read packet", "error", err)

			return
		}

		if err = dst.Write(packet, &acks); err != nil {
			logger.Debugw("Cannot send packet", "error", err)

			return
		}
	}
}