socks5/addr.go

package socks5

import (
	"encoding/binary"
	"errors"
	"net"
	"strconv"
)

var (
	// AddrZero is zero address
	AddrZero = &Addr{IP: net.IPv4zero, Port: 0}
)

// Addr implements net.Addr interface
type Addr struct {
	IP   net.IP
	Host string
	Port int
}

// AddrFromString returns new address from string representation
func AddrFromString(address string) (*Addr, error) {
	if address == "" {
		return AddrZero, nil
	}

	host, port, err := splitHostPort(address)
	if err != nil {
		return nil, err
	}

	addr := &Addr{Port: port}
	addr.IP = net.ParseIP(host)
	if addr.IP == nil {
		addr.Host = host
	}

	return addr, nil
}

func AddrFromSocks(atyp ATYP, host []byte, port []byte) *Addr {
	addr := &Addr{Port: int(binary.BigEndian.Uint16(port))}

	switch atyp {
	case ATYPFQDN:
		addr.Host = string(host[1:])
	default:
		addr.IP = net.IP(host)
	}

	return addr
}

// Socks returns socks protocol formatted address type, address and port
func (s *Addr) Socks() (ATYP, []byte, []byte) {
	portBytes := make([]byte, 2)
	binary.BigEndian.PutUint16(portBytes, uint16(s.Port))

	if s.IP != nil {
		if ip4 := s.IP.To4(); ip4 != nil {
			return ATYPIPv4, []byte(ip4), portBytes
		} else if ip6 := s.IP.To16(); ip6 != nil {
			return ATYPIPv6, []byte(ip6), portBytes
		}
	}

	addr := []byte{byte(len(s.Host))}
	addr = append(addr, []byte(s.Host)...)
	return ATYPFQDN, addr, portBytes
}

// UDP returns udp address
func (s *Addr) UDP() (*net.UDPAddr, error) {
	if s.IP != nil {
		return &net.UDPAddr{IP: s.IP, Port: s.Port}, nil
	}

	return net.ResolveUDPAddr("udp", s.Host)
}

// Network returns address Network
func (s *Addr) Network() string {
	return "socks"
}

func (s *Addr) String() string {
	if s == nil {
		return "<nil>"
	}

	port := strconv.Itoa(s.Port)
	if s.IP == nil {
		return net.JoinHostPort(s.Host, port)
	}

	return net.JoinHostPort(s.IP.String(), port)
}

func splitHostPort(address string) (string, int, error) {
	host, port, err := net.SplitHostPort(address)
	if err != nil {
		return "", 0, err
	}

	portInt, err := strconv.Atoi(port)
	if err != nil {
		return "", 0, err
	}

	if 1 > portInt || portInt > 0xffff {
		return "", 0, errors.New("port number out of range " + port)
	}

	return host, portInt, nil
}