merge address writing

proxy/shadowsocks/protocol.go

@@ -10,7 +10,7 @@ import (
 	"v2ray.com/core/common/buf"
 	"v2ray.com/core/common/net"
 	"v2ray.com/core/common/protocol"
-	"v2ray.com/core/common/serial"
+	"v2ray.com/core/proxy/socks"
 )
 
 const (
@@ -165,26 +165,10 @@ func WriteTCPRequest(request *protocol.RequestHeader, writer io.Writer) (buf.Wri
 
 	header := buf.NewLocal(512)
 
-	switch request.Address.Family() {
-	case net.AddressFamilyIPv4:
-		header.AppendBytes(AddrTypeIPv4)
-		header.Append([]byte(request.Address.IP()))
-	case net.AddressFamilyIPv6:
-		header.AppendBytes(AddrTypeIPv6)
-		header.Append([]byte(request.Address.IP()))
-	case net.AddressFamilyDomain:
-		domain := request.Address.Domain()
-		if protocol.IsDomainTooLong(domain) {
-			return nil, newError("domain name too long: ", domain)
-		}
-		header.AppendBytes(AddrTypeDomain, byte(len(domain)))
-		common.Must(header.AppendSupplier(serial.WriteString(domain)))
-	default:
-		return nil, newError("unsupported address type: ", request.Address.Family())
+	if err := socks.AppendAddress(header, request.Address, request.Port); err != nil {
+		return nil, newError("failed to write address").Base(err)
 	}
 
-	common.Must(header.AppendSupplier(serial.WriteUint16(uint16(request.Port))))
-
 	if request.Option.Has(RequestOptionOneTimeAuth) {
 		header.SetByte(0, header.Byte(0)|0x10)
 
@@ -261,21 +245,10 @@ func EncodeUDPPacket(request *protocol.RequestHeader, payload []byte) (*buf.Buff
 	}
 	iv := buffer.Bytes()
 
-	switch request.Address.Family() {
-	case net.AddressFamilyIPv4:
-		buffer.AppendBytes(AddrTypeIPv4)
-		buffer.Append([]byte(request.Address.IP()))
-	case net.AddressFamilyIPv6:
-		buffer.AppendBytes(AddrTypeIPv6)
-		buffer.Append([]byte(request.Address.IP()))
-	case net.AddressFamilyDomain:
-		buffer.AppendBytes(AddrTypeDomain, byte(len(request.Address.Domain())))
-		buffer.Append([]byte(request.Address.Domain()))
-	default:
-		return nil, newError("unsupported address type: ", request.Address.Family()).AtError()
+	if err := socks.AppendAddress(buffer, request.Address, request.Port); err != nil {
+		return nil, newError("failed to write address").Base(err)
 	}
 
-	common.Must(buffer.AppendSupplier(serial.WriteUint16(uint16(request.Port))))
 	buffer.Append(payload)
 
 	if !account.Cipher.IsAEAD() && request.Option.Has(RequestOptionOneTimeAuth) {

proxy/socks/protocol.go

@@ -245,19 +245,20 @@ func writeSocks5AuthenticationResponse(writer io.Writer, version byte, auth byte
 	return err
 }
 
-func appendAddress(buffer *buf.Buffer, address net.Address, port net.Port) error {
+// AppendAddress appends Socks address into the given buffer.
+func AppendAddress(buffer *buf.Buffer, address net.Address, port net.Port) error {
 	switch address.Family() {
 	case net.AddressFamilyIPv4:
-		buffer.AppendBytes(0x01)
+		buffer.AppendBytes(addrTypeIPv4)
 		buffer.Append(address.IP())
 	case net.AddressFamilyIPv6:
-		buffer.AppendBytes(0x04)
+		buffer.AppendBytes(addrTypeIPv6)
 		buffer.Append(address.IP())
 	case net.AddressFamilyDomain:
 		if protocol.IsDomainTooLong(address.Domain()) {
 			return newError("Super long domain is not supported in Socks protocol: ", address.Domain())
 		}
-		buffer.AppendBytes(0x03, byte(len(address.Domain())))
+		buffer.AppendBytes(addrTypeDomain, byte(len(address.Domain())))
 		common.Must(buffer.AppendSupplier(serial.WriteString(address.Domain())))
 	}
 	common.Must(buffer.AppendSupplier(serial.WriteUint16(port.Value())))
@@ -267,7 +268,7 @@ func appendAddress(buffer *buf.Buffer, address net.Address, port net.Port) error
 func writeSocks5Response(writer io.Writer, errCode byte, address net.Address, port net.Port) error {
 	buffer := buf.NewLocal(64)
 	buffer.AppendBytes(socks5Version, errCode, 0x00 /* reserved */)
-	if err := appendAddress(buffer, address, port); err != nil {
+	if err := AppendAddress(buffer, address, port); err != nil {
 		return err
 	}
 
@@ -337,7 +338,7 @@ func DecodeUDPPacket(packet []byte) (*protocol.RequestHeader, []byte, error) {
 func EncodeUDPPacket(request *protocol.RequestHeader, data []byte) (*buf.Buffer, error) {
 	b := buf.New()
 	b.AppendBytes(0, 0, 0 /* Fragment */)
-	if err := appendAddress(b, request.Address, request.Port); err != nil {
+	if err := AppendAddress(b, request.Address, request.Port); err != nil {
 		return nil, err
 	}
 	b.Append(data)
@@ -444,7 +445,7 @@ func ClientHandshake(request *protocol.RequestHeader, reader io.Reader, writer i
 		command = byte(cmdUDPPort)
 	}
 	b.AppendBytes(socks5Version, command, 0x00 /* reserved */)
-	appendAddress(b, request.Address, request.Port)
+	AppendAddress(b, request.Address, request.Port)
 	if _, err := writer.Write(b.Bytes()); err != nil {
 		return nil, err
 	}