go-playit/core/tcp_client.go

package core

import (
	"net"
	"net/netip"
	"time"

	"sirherobrine23.com.br/go-bds/go-playit/api"
	"sirherobrine23.com.br/go-bds/go-playit/proto"
)

type TcpClient struct {
	*net.TCPConn
	claimProto *proto.NewClient
	tunnel     api.AgentTunnel
}

// Return Local address and port to connect in localhost
func (client *TcpClient) LocalAddr() net.Addr {
	addr, port := client.tunnel.LocalIP, client.tunnel.LocalPort+client.claimProto.PortOffset
	return net.TCPAddrFromAddrPort(netip.AddrPortFrom(addr, port))
}

// Get original user address
func (client *TcpClient) RemoteAddr() net.Addr {
	return net.TCPAddrFromAddrPort(client.claimProto.PeerAddr)
}

func (a *PlayitAgent) AcceptTcp() (net.Conn, error) {
	if a.clientAcceptTcp == nil {
		return nil, net.ErrClosed
	}

	claim, ok := <-a.clientAcceptTcp
	if !ok {
		return nil, net.ErrClosed
	}

	// Get tunnel origim
	var tunnelInfo api.AgentTunnel
	if a.AgentInfo == nil {
		tunnelInfo = api.AgentTunnel{
			LocalIP:   claim.ConnectAddr.Addr(),
			LocalPort: claim.ConnectAddr.Port(),
		}
	} else {
		for _, tunnelInfo = range a.AgentInfo.Tunnels {
			if tunnelInfo.InternalID == claim.TunnelID {
				break
			}
			tunnelInfo = api.AgentTunnel{
				LocalIP:   claim.ConnectAddr.Addr(),
				LocalPort: claim.ConnectAddr.Port(),
			}
		}
	}

	conn, err := net.DialTCP("tcp", nil, net.TCPAddrFromAddrPort(claim.ClaimInstructions.Address))
	if err != nil {
		return nil, err
	}

	n := 0
	for n != len(claim.ClaimInstructions.Token) {
		n2, err := conn.Write(claim.ClaimInstructions.Token[n:])
		n += n2
		if err != nil {
			return nil, err
		} else if a.clientAcceptTcp == nil {
			conn.Close()
			return nil, net.ErrClosed
		}
	}

	defer conn.SetReadDeadline(time.Time{})
	conn.SetReadDeadline(time.Now().Add(time.Millisecond * 600))

	expect := make([]byte, 8)
	n, err = conn.Read(expect)
	if err != nil {
		conn.Close()
		return nil, err
	} else if n != 8 {
		conn.Close()
		return nil, net.ErrClosed
	}

	return &TcpClient{
		TCPConn:    conn,
		claimProto: claim,
		tunnel:     tunnelInfo,
	}, nil
}