Merge pull request #82 from Bouke/patch-1

Change how domain name is resolved to tunnel

localtunnel/server/backend.py

@@ -11,7 +11,7 @@
 from localtunnel import __version__
 from localtunnel.server import metrics
 
-HOST_TEMPLATE = "{0}.v2.localtunnel.com"
+HOST_TEMPLATE = "{0}.{1}"
 BANNER = """Thanks for trying localtunnel v2 beta!
   Source code: https://github.com/progrium/localtunnel
   Donate: http://j.mp/donate-localtunnel
@@ -42,7 +42,7 @@ def handle_control_request(socket, request):
         socket.close()
         return
     protocol.send_message(socket, protocol.control_reply(
-        host=HOST_TEMPLATE.format(tunnel.name),
+        host=HOST_TEMPLATE.format(tunnel.name, Tunnel.domain_suffix),
         banner=BANNER,
         concurrency=Tunnel.max_pool_size,
     ))

localtunnel/server/cli.py

@@ -25,17 +25,15 @@ def run():
                 help='port to run public frontend', default=8000)
     parser.add_argument('backend_port', metavar='backend_port', type=int,
                 help='port to run backend server', default=8001)
-    parser.add_argument('-d', '--domainpart', type=int,
-                help='domain part (from the right) to extract tunnel name')
+    parser.add_argument('domain_suffix', metavar='domain_suffix', type=str,
+                help='domain suffix (from the right) to extract tunnel name')
     args = parser.parse_args()
 
     logging.info("starting frontend on {0}...".format(args.frontend_port))
     logging.info("starting backend on {0}...".format(args.backend_port))
 
     Tunnel.backend_port = args.backend_port
-
-    if args.domainpart:
-        Tunnel.domain_part = args.domainpart
+    Tunnel.domain_suffix = args.domain_suffix
 
     stats_key = os.environ.get('STATHAT_EZKEY', None)
     if stats_key:

localtunnel/server/tunnel.py

@@ -1,4 +1,5 @@
 import json
+import re
 import time
 import logging
 
@@ -11,7 +12,7 @@
 
 class Tunnel(object):
     max_pool_size = 3
-    domain_part = 3
+    domain_suffix = None
     backend_port = None
     active_timeout = 5 * 60
 
@@ -75,17 +76,11 @@ def create(cls, obj):
 
     @classmethod
     def get_by_hostname(cls, hostname):
-        try:
-            name = hostname.split('.')[-1 * Tunnel.domain_part]
-        except IndexError:
+        if not hostname.endswith(Tunnel.domain_suffix):
             return
-        tunnel = cls._tunnels.get(name)
-        if not tunnel:
-            for n, tunnel in cls._tunnels.iteritems():
-                if tunnel.domain and hostname.endswith(tunnel.domain):
-                    return tunnel
-        else:
-            return tunnel
+        m = re.match('(.+?\.|)(\w+)\.%s$' % Tunnel.domain_suffix, hostname)
+        if m:
+            return cls._tunnels.get(m.group(2))
 
     @classmethod
     def get_by_control_request(cls, request):

tests/server/test_tunnel.py

@@ -0,0 +1,23 @@
+import unittest
+from localtunnel.server.tunnel import Tunnel
+
+
+class TestTunnel(unittest.TestCase):
+    def test_get_by_hostname(self):
+        Tunnel.domain_suffix = 'bar'
+        tunnel = Tunnel.create(dict(name='foo', client='Test-Client'))
+        self.assertTrue(Tunnel.get_by_hostname('foo.bar'))
+        self.assertTrue(Tunnel.get_by_hostname('xxx.foo.bar'))
+        self.assertFalse(Tunnel.get_by_hostname('foo.bar.bar'))
+        tunnel.destroy()
+
+        Tunnel.domain_suffix = 'foo.bar'
+        tunnel = Tunnel.create(dict(name='hello', client='Test-Client'))
+        self.assertTrue(Tunnel.get_by_hostname('hello.foo.bar'))
+        self.assertTrue(Tunnel.get_by_hostname('world.hello.foo.bar'))
+        self.assertFalse(Tunnel.get_by_hostname('foo.bar'))
+        self.assertFalse(Tunnel.get_by_hostname('bar'))
+        self.assertFalse(Tunnel.get_by_hostname('hello.world.foo.bar'))
+        tunnel.destroy()
+
+        Tunnel.domain_suffix = None