mirror of
https://github.com/v2fly/v2ray-core.git
synced 2026-02-03 14:55:24 -05:00
unify all address reading and writing
This commit is contained in:
Darien Raymond
183
common/protocol/address.go
Normal file
183
common/protocol/address.go
Normal file
@@ -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
|
||||
}
|
||||
70
common/protocol/address_test.go
Normal file
70
common/protocol/address_test.go
Normal file
@@ -0,0 +1,70 @@
|
||||
package protocol_test
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"testing"
|
||||
|
||||
"v2ray.com/core/common/buf"
|
||||
"v2ray.com/core/common/net"
|
||||
. "v2ray.com/core/common/protocol"
|
||||
. "v2ray.com/ext/assert"
|
||||
)
|
||||
|
||||
func TestAddressParser(t *testing.T) {
|
||||
assert := With(t)
|
||||
|
||||
data := []struct {
|
||||
Options []AddressOption
|
||||
Input []byte
|
||||
Address net.Address
|
||||
Port net.Port
|
||||
Error bool
|
||||
}{
|
||||
{
|
||||
Options: []AddressOption{},
|
||||
Input: []byte{0, 0, 0, 0, 0},
|
||||
Error: true,
|
||||
},
|
||||
{
|
||||
Options: []AddressOption{AddressFamilyByte(0x01, net.AddressFamilyIPv4)},
|
||||
Input: []byte{1, 0, 0, 0, 0, 0, 53},
|
||||
Address: net.IPAddress([]byte{0, 0, 0, 0}),
|
||||
Port: net.Port(53),
|
||||
},
|
||||
{
|
||||
Options: []AddressOption{AddressFamilyByte(0x01, net.AddressFamilyIPv4)},
|
||||
Input: []byte{1, 0, 0, 0, 0},
|
||||
Error: true,
|
||||
},
|
||||
{
|
||||
Options: []AddressOption{AddressFamilyByte(0x04, net.AddressFamilyIPv6)},
|
||||
Input: []byte{4, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 0, 80},
|
||||
Address: net.IPAddress([]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6}),
|
||||
Port: net.Port(80),
|
||||
},
|
||||
{
|
||||
Options: []AddressOption{AddressFamilyByte(0x03, net.AddressFamilyDomain)},
|
||||
Input: []byte{3, 9, 118, 50, 114, 97, 121, 46, 99, 111, 109, 0, 80},
|
||||
Address: net.DomainAddress("v2ray.com"),
|
||||
Port: net.Port(80),
|
||||
},
|
||||
{
|
||||
Options: []AddressOption{AddressFamilyByte(0x03, net.AddressFamilyDomain)},
|
||||
Input: []byte{3, 9, 118, 50, 114, 97, 121, 46, 99, 111, 109, 0},
|
||||
Error: true,
|
||||
},
|
||||
}
|
||||
|
||||
for _, tc := range data {
|
||||
b := buf.New()
|
||||
parser := NewAddressParser(tc.Options...)
|
||||
addr, port, err := parser.ReadAddressPort(b, bytes.NewReader(tc.Input))
|
||||
b.Release()
|
||||
if tc.Error {
|
||||
assert(err, IsNotNil)
|
||||
} else {
|
||||
assert(addr, Equals, tc.Address)
|
||||
assert(port, Equals, tc.Port)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user