github.com/ipfans/trojan-go@v0.11.0/tunnel/freedom/conn.go (about)

     1  package freedom
     2  
     3  import (
     4  	"bytes"
     5  	"net"
     6  
     7  	"github.com/txthinking/socks5"
     8  
     9  	"github.com/ipfans/trojan-go/common"
    10  	"github.com/ipfans/trojan-go/log"
    11  	"github.com/ipfans/trojan-go/tunnel"
    12  )
    13  
    14  const MaxPacketSize = 1024 * 8
    15  
    16  type Conn struct {
    17  	net.Conn
    18  }
    19  
    20  func (c *Conn) Metadata() *tunnel.Metadata {
    21  	return nil
    22  }
    23  
    24  type PacketConn struct {
    25  	*net.UDPConn
    26  }
    27  
    28  func (c *PacketConn) WriteWithMetadata(p []byte, m *tunnel.Metadata) (int, error) {
    29  	return c.WriteTo(p, m.Address)
    30  }
    31  
    32  func (c *PacketConn) ReadWithMetadata(p []byte) (int, *tunnel.Metadata, error) {
    33  	n, addr, err := c.ReadFrom(p)
    34  	if err != nil {
    35  		return 0, nil, err
    36  	}
    37  	address, err := tunnel.NewAddressFromAddr("udp", addr.String())
    38  	common.Must(err)
    39  	metadata := &tunnel.Metadata{
    40  		Address: address,
    41  	}
    42  	return n, metadata, nil
    43  }
    44  
    45  func (c *PacketConn) WriteTo(p []byte, addr net.Addr) (int, error) {
    46  	if udpAddr, ok := addr.(*net.UDPAddr); ok {
    47  		return c.WriteToUDP(p, udpAddr)
    48  	}
    49  	ip, err := addr.(*tunnel.Address).ResolveIP()
    50  	if err != nil {
    51  		return 0, err
    52  	}
    53  	udpAddr := &net.UDPAddr{
    54  		IP:   ip,
    55  		Port: addr.(*tunnel.Address).Port,
    56  	}
    57  	return c.WriteToUDP(p, udpAddr)
    58  }
    59  
    60  type SocksPacketConn struct {
    61  	net.PacketConn
    62  	socksAddr   *net.UDPAddr
    63  	socksClient *socks5.Client
    64  }
    65  
    66  func (c *SocksPacketConn) WriteWithMetadata(payload []byte, metadata *tunnel.Metadata) (int, error) {
    67  	buf := bytes.NewBuffer(make([]byte, 0, MaxPacketSize))
    68  	buf.Write([]byte{0, 0, 0}) // RSV, FRAG
    69  	common.Must(metadata.Address.WriteTo(buf))
    70  	buf.Write(payload)
    71  	_, err := c.PacketConn.WriteTo(buf.Bytes(), c.socksAddr)
    72  	if err != nil {
    73  		return 0, err
    74  	}
    75  	log.Debug("sent udp packet to " + c.socksAddr.String() + " with metadata " + metadata.String())
    76  	return len(payload), nil
    77  }
    78  
    79  func (c *SocksPacketConn) ReadWithMetadata(payload []byte) (int, *tunnel.Metadata, error) {
    80  	buf := make([]byte, MaxPacketSize)
    81  	n, from, err := c.PacketConn.ReadFrom(buf)
    82  	if err != nil {
    83  		return 0, nil, err
    84  	}
    85  	log.Debug("recv udp packet from " + from.String())
    86  	addr := new(tunnel.Address)
    87  	r := bytes.NewBuffer(buf[3:n])
    88  	if err := addr.ReadFrom(r); err != nil {
    89  		return 0, nil, common.NewError("socks5 failed to parse addr in the packet").Base(err)
    90  	}
    91  	length, err := r.Read(payload)
    92  	if err != nil {
    93  		return 0, nil, err
    94  	}
    95  	return length, &tunnel.Metadata{
    96  		Address: addr,
    97  	}, nil
    98  }
    99  
   100  func (c *SocksPacketConn) Close() error {
   101  	c.socksClient.Close()
   102  	return c.PacketConn.Close()
   103  }