github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/net/proxy/vless/conn.go (about)

     1  package vless
     2  
     3  import (
     4  	"context"
     5  	"encoding/binary"
     6  	"errors"
     7  	"fmt"
     8  	"io"
     9  	"net"
    10  
    11  	"github.com/Asutorufa/yuhaiin/pkg/net/netapi"
    12  	"github.com/Asutorufa/yuhaiin/pkg/utils/pool"
    13  	"github.com/Asutorufa/yuhaiin/pkg/utils/uuid"
    14  )
    15  
    16  type Conn struct {
    17  	net.Conn
    18  	dst      netapi.Address
    19  	id       uuid.UUID
    20  	received bool
    21  	udp      bool
    22  
    23  	remain int
    24  }
    25  
    26  func (vc *Conn) ReadFrom(b []byte) (int, net.Addr, error) {
    27  	if vc.remain > 0 {
    28  		length := min(len(b), vc.remain)
    29  		n, err := vc.Read(b[:length])
    30  		vc.remain -= n
    31  		return n, vc.dst, err
    32  	}
    33  
    34  	var total uint16
    35  	if err := binary.Read(vc.Conn, binary.BigEndian, &total); err != nil {
    36  		return 0, nil, fmt.Errorf("read length failed: %w", err)
    37  	}
    38  
    39  	length := min(len(b), int(total))
    40  
    41  	if n, err := io.ReadFull(vc.Conn, b[:total]); err != nil {
    42  		return n, vc.dst, fmt.Errorf("read packet error: %w", err)
    43  	}
    44  
    45  	vc.remain = int(total) - length
    46  	return length, vc.dst, nil
    47  }
    48  
    49  func (vc *Conn) WriteTo(b []byte, target net.Addr) (int, error) {
    50  	buf := pool.GetBytesWriter(2 + len(b))
    51  	defer buf.Free()
    52  
    53  	_ = binary.Write(buf, binary.BigEndian, uint16(len(b)))
    54  	_, _ = buf.Write(b)
    55  
    56  	_, err := vc.Write(buf.Bytes())
    57  	return len(b), err
    58  }
    59  
    60  func (vc *Conn) Read(b []byte) (int, error) {
    61  	if vc.received {
    62  		return vc.Conn.Read(b)
    63  	}
    64  
    65  	if err := vc.recvResponse(); err != nil {
    66  		return 0, err
    67  	}
    68  	vc.received = true
    69  	return vc.Conn.Read(b)
    70  }
    71  
    72  func (vc *Conn) sendRequest() error {
    73  	buf := pool.GetBytesWriter(pool.DefaultSize)
    74  	defer buf.Free()
    75  
    76  	buf.WriteByte(Version)          // protocol version
    77  	_, _ = buf.Write(vc.id.Bytes()) // 16 bytes of uuid
    78  	buf.WriteByte(0)                // addon data length. 0 means no addon data
    79  
    80  	// Command
    81  	if vc.udp {
    82  		buf.WriteByte(CommandUDP)
    83  	} else {
    84  		buf.WriteByte(CommandTCP)
    85  	}
    86  
    87  	// Port AddrType Addr
    88  	_ = binary.Write(buf, binary.BigEndian, vc.dst.Port().Port())
    89  
    90  	if vc.dst.IsFqdn() {
    91  		buf.WriteByte(AtypDomainName)
    92  		buf.WriteByte(byte(len(vc.dst.Hostname())))
    93  		buf.WriteString(vc.dst.Hostname())
    94  	} else {
    95  		addrPort := vc.dst.AddrPort(context.TODO()).V.Addr()
    96  
    97  		if addrPort.Is6() {
    98  			buf.WriteByte(AtypIPv6)
    99  		} else {
   100  			buf.WriteByte(AtypIPv4)
   101  		}
   102  
   103  		_, _ = buf.Write(addrPort.AsSlice())
   104  	}
   105  
   106  	_, err := vc.Conn.Write(buf.Bytes())
   107  	return err
   108  }
   109  
   110  func (vc *Conn) recvResponse() error {
   111  	var buf [2]byte
   112  	if _, err := io.ReadFull(vc.Conn, buf[:]); err != nil {
   113  		return err
   114  	}
   115  
   116  	if buf[0] != Version {
   117  		return errors.New("unexpected response version")
   118  	}
   119  
   120  	length := int64(buf[1])
   121  	if length > 0 { // addon data length > 0
   122  		_, _ = io.CopyN(io.Discard, vc.Conn, length) // just discard
   123  	}
   124  
   125  	return nil
   126  }
   127  
   128  // newConn return a Conn instance
   129  func newConn(conn net.Conn, client *Client, udp bool, dst netapi.Address) (*Conn, error) {
   130  	c := &Conn{
   131  		Conn: conn,
   132  		id:   client.uuid,
   133  		dst:  dst,
   134  		udp:  udp,
   135  	}
   136  
   137  	if err := c.sendRequest(); err != nil {
   138  		return nil, err
   139  	}
   140  	return c, nil
   141  }