From a323f878e1efcee0974cebe5489ded46f960740d Mon Sep 17 00:00:00 2001
From: clowwindy <clowwindy42@gmail.com>
Date: Sat, 12 Jul 2014 10:31:35 +0800
Subject: [PATCH] prefer addrtype 1 and 4 over 3 in UDP relay

---
 shadowsocks/common.py   | 15 +++++++++++++++
 shadowsocks/udprelay.py |  5 ++---
 2 files changed, 17 insertions(+), 3 deletions(-)

diff --git a/shadowsocks/common.py b/shadowsocks/common.py
index 6104478..7548ff3 100644
--- a/shadowsocks/common.py
+++ b/shadowsocks/common.py
@@ -79,6 +79,21 @@ ADDRTYPE_IPV6 = 4
 ADDRTYPE_HOST = 3
 
 
+def pack_addr(address):
+    for family in (socket.AF_INET, socket.AF_INET6):
+        try:
+            r = socket.inet_pton(family, address)
+            if family == socket.AF_INET6:
+                return '\x04' + r
+            else:
+                return '\x01' + r
+        except (TypeError, ValueError, OSError, IOError):
+            pass
+    if len(address) > 255:
+        address = address[:255]  # TODO
+    return '\x03' + chr(len(address)) + address
+
+
 def parse_header(data):
     addrtype = ord(data[0])
     dest_addr = None
diff --git a/shadowsocks/udprelay.py b/shadowsocks/udprelay.py
index 7a72c11..d7e3d4c 100644
--- a/shadowsocks/udprelay.py
+++ b/shadowsocks/udprelay.py
@@ -75,7 +75,7 @@ import random
 import encrypt
 import eventloop
 import lru_cache
-from common import parse_header
+from common import parse_header, pack_addr
 
 
 BUF_SIZE = 65536
@@ -214,8 +214,7 @@ class UDPRelay(object):
             if addrlen > 255:
                 # drop
                 return
-            data = '\x03' + chr(addrlen) + r_addr[0] + \
-                   struct.pack('>H', r_addr[1]) + data
+            data = pack_addr(r_addr[0]) + struct.pack('>H', r_addr[1]) + data
             response = encrypt.encrypt_all(self._password, self._method, 1,
                                            data)
             if not response: