unify all address reading and writing

2025-12-30 14:05:26 -05:00 · 2018-02-23 23:42:01 +01:00
parent a059ee2c00
commit af1abf687c
10 changed files with 333 additions and 267 deletions
--- a/common/protocol/address.go
+++ b/common/protocol/address.go
@@ -0,0 +1,183 @@
+package protocol
+
+import (
+	"io"
+
+	"v2ray.com/core/common/buf"
+	"v2ray.com/core/common/net"
+)
+
+type AddressOption func(*AddressParser)
+
+func PortThenAddress() AddressOption {
+	return func(p *AddressParser) {
+		p.portFirst = true
+	}
+}
+
+func AddressFamilyByte(b byte, f net.AddressFamily) AddressOption {
+	return func(p *AddressParser) {
+		p.addrTypeMap[b] = f
+		p.addrByteMap[f] = b
+	}
+}
+
+type AddressTypeParser func(byte) byte
+
+func WithAddressTypeParser(atp AddressTypeParser) AddressOption {
+	return func(p *AddressParser) {
+		p.typeParser = atp
+	}
+}
+
+type AddressParser struct {
+	addrTypeMap map[byte]net.AddressFamily
+	addrByteMap map[net.AddressFamily]byte
+	portFirst   bool
+	typeParser  AddressTypeParser
+}
+
+func NewAddressParser(options ...AddressOption) *AddressParser {
+	p := &AddressParser{
+		addrTypeMap: make(map[byte]net.AddressFamily, 8),
+		addrByteMap: make(map[net.AddressFamily]byte, 8),
+	}
+	for _, opt := range options {
+		opt(p)
+	}
+	return p
+}
+
+func (p *AddressParser) readPort(b *buf.Buffer, reader io.Reader) (net.Port, error) {
+	if err := b.AppendSupplier(buf.ReadFullFrom(reader, 2)); err != nil {
+		return 0, err
+	}
+	return net.PortFromBytes(b.BytesFrom(-2)), nil
+}
+
+func (p *AddressParser) readAddress(b *buf.Buffer, reader io.Reader) (net.Address, error) {
+	if err := b.AppendSupplier(buf.ReadFullFrom(reader, 1)); err != nil {
+		return nil, err
+	}
+
+	addrType := b.Byte(b.Len() - 1)
+	if p.typeParser != nil {
+		addrType = p.typeParser(addrType)
+	}
+
+	addrFamily, valid := p.addrTypeMap[addrType]
+	if !valid {
+		return nil, newError("unknown address type: ", addrType)
+	}
+
+	switch addrFamily {
+	case net.AddressFamilyIPv4:
+		if err := b.AppendSupplier(buf.ReadFullFrom(reader, 4)); err != nil {
+			return nil, err
+		}
+		return net.IPAddress(b.BytesFrom(-4)), nil
+	case net.AddressFamilyIPv6:
+		if err := b.AppendSupplier(buf.ReadFullFrom(reader, 16)); err != nil {
+			return nil, err
+		}
+		return net.IPAddress(b.BytesFrom(-16)), nil
+	case net.AddressFamilyDomain:
+		if err := b.AppendSupplier(buf.ReadFullFrom(reader, 1)); err != nil {
+			return nil, err
+		}
+		domainLength := int(b.Byte(b.Len() - 1))
+		if err := b.AppendSupplier(buf.ReadFullFrom(reader, domainLength)); err != nil {
+			return nil, err
+		}
+		return net.DomainAddress(string(b.BytesFrom(-domainLength))), nil
+	default:
+		panic("impossible case")
+	}
+}
+
+func (p *AddressParser) ReadAddressPort(buffer *buf.Buffer, input io.Reader) (net.Address, net.Port, error) {
+	if buffer == nil {
+		buffer = buf.New()
+		defer buffer.Release()
+	}
+
+	if p.portFirst {
+		port, err := p.readPort(buffer, input)
+		if err != nil {
+			return nil, 0, err
+		}
+		addr, err := p.readAddress(buffer, input)
+		if err != nil {
+			return nil, 0, err
+		}
+		return addr, port, nil
+	}
+
+	addr, err := p.readAddress(buffer, input)
+	if err != nil {
+		return nil, 0, err
+	}
+
+	port, err := p.readPort(buffer, input)
+	if err != nil {
+		return nil, 0, err
+	}
+
+	return addr, port, nil
+}
+
+func (p *AddressParser) writePort(writer io.Writer, port net.Port) error {
+	if _, err := writer.Write(port.Bytes(nil)); err != nil {
+		return err
+	}
+	return nil
+}
+
+func (p *AddressParser) writeAddress(writer io.Writer, address net.Address) error {
+	tb, valid := p.addrByteMap[address.Family()]
+	if !valid {
+		return newError("unknown address family", address.Family())
+	}
+
+	switch address.Family() {
+	case net.AddressFamilyIPv4, net.AddressFamilyIPv6:
+		if _, err := writer.Write([]byte{tb}); err != nil {
+			return err
+		}
+		if _, err := writer.Write(address.IP()); err != nil {
+			return err
+		}
+	case net.AddressFamilyDomain:
+		domain := address.Domain()
+		if IsDomainTooLong(domain) {
+			return newError("Super long domain is not supported: ", domain)
+		}
+		if _, err := writer.Write([]byte{tb, byte(len(domain))}); err != nil {
+			return err
+		}
+		if _, err := writer.Write([]byte(domain)); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (p *AddressParser) WriteAddressPort(writer io.Writer, addr net.Address, port net.Port) error {
+	if p.portFirst {
+		if err := p.writePort(writer, port); err != nil {
+			return err
+		}
+		if err := p.writeAddress(writer, addr); err != nil {
+			return err
+		}
+		return nil
+	}
+
+	if err := p.writeAddress(writer, addr); err != nil {
+		return err
+	}
+	if err := p.writePort(writer, port); err != nil {
+		return err
+	}
+	return nil
+}