[peer/host addr] Req_Probe__sock_addr__proxy, Req_Probe__sock_addr__direct

2 files changed, 77 insertions, 0 deletions

diff --git a/src/server/rz_mesh.py b/src/server/rz_mesh.py
index 78b6a22b..09e263f2 100644
--- a/src/server/rz_mesh.py
+++ b/src/server/rz_mesh.py
@@ -29,6 +29,30 @@ class RZ_WebSocket_Server(SocketIOServer):
         - allow response header injection on websocket connections
     """
 
+    class Req_Probe__sock_addr__proxy(object):
+
+        def probe_client_socket_addr__ws_conn(self, ws_environ):
+
+            ret = sock_addr_from_REMOTE_X_keys(ws_environ)
+
+            #
+            # relying on the presence of the 'X-Forwarded-For' is preferable, but
+            # a bit flaky as it is not always present - see #496
+            #
+            # TODO: evaluate proxy server's behavior on this
+            #
+            try:
+                _, __ = sock_addr_from_env_HTTP_headers(ws_environ, key_name__addr='X-Forwarded-For')
+            except Exception as e:
+                log.warning('ws: client socket addr probe: %s, peer-addr ~: %s:%s' % (e.message, ret[0], ret[1]))
+
+            return ret
+
+    class Req_Probe__sock_addr__direct(object):
+
+        def probe_client_socket_addr__ws_conn(self, ws_environ):
+            return sock_addr_from_REMOTE_X_keys(ws_environ)
+
     class RZ_SocketIOHandler(SocketIOHandler):
 
         def __init__(self, config, *args, **kwargs):
diff --git a/src/server/rz_server.py b/src/server/rz_server.py
index 6432b626..ddc6669e 100644
--- a/src/server/rz_server.py
+++ b/src/server/rz_server.py
@@ -141,6 +141,59 @@ class FlaskExt(Flask):
     Flask server customization
     """
 
+    class Req_Probe__sock_addr__proxy(object):
+
+        def __init__(self, proxy_host, proxy_port):
+            self.proxy_host = proxy_host
+            self.proxy_port = proxy_port
+
+        def probe_client_socket_addr__http_req(self, req):
+
+            ret = sock_addr_from_REMOTE_X_keys(req.environ)
+
+            #
+            # relying on the presence of the 'X-Forwarded-For' is preferable, but
+            # a bit flaky as it is not always present - see #496
+            #
+            # TODO: evaluate proxy server's behavior on this
+            #
+            try:
+                _, __ = sock_addr_from_env_HTTP_headers(req.environ, key_name__addr='X-Forwarded-For')
+            except Exception as e:
+                log.warning('flask: client socket addr probe: %s, peer-addr ~: %s:%s' % (e.message, ret[0], ret[1]))
+
+            return ret
+
+        def probe_requested_host__http_req(self, req, probe_for_proxy=True):
+
+            ret = self.proxy_host, self.proxy_port
+
+            #
+            # relying on the presence of the 'X-Forwarded-Host' is preferable, but
+            # a bit flaky as it is not always present - see #496
+            #
+            # TODO: evaluate proxy server's behavior on this
+            #
+            try:
+                _, __ = sock_addr_from_env_HTTP_headers(req.environ, key_name__addr='X-Forwarded-Host')
+            except Exception as e:
+                log.warning('flask: client socket addr probe: %s, replacing with: \'%s:%s\'' % (e.message, ret[0], ret[1]))
+
+            return ret
+
+    class Req_Probe__sock_addr__direct(object):
+
+        def __init__(self, dafault_port):
+            self.dafault_port = dafault_port
+
+        def probe_client_socket_addr__http_req(self, req):
+            return sock_addr_from_REMOTE_X_keys(req.environ)
+
+        def probe_requested_host__http_req(self, req, probe_for_proxy=True):
+            req_host__addr, _ = sock_addr_from_env_HTTP_headers(req.environ, key_name__addr='Host')
+            req_host__port = self.dafault_port
+            return req_host__addr, req_host__port
+
     def __init__(self, import_name, *args, **kwargs):
         """
         reserved for future use