Merge pull request #3124 from nbrownus/elb_proxy_protocol

ELB Proxy Protocol support for backends

1 files changed, 110 insertions, 2 deletions

diff --git a/cloud/amazon/ec2_elb_lb.py b/cloud/amazon/ec2_elb_lb.py
index faf99c6a..d05a0126 100644
--- a/cloud/amazon/ec2_elb_lb.py
+++ b/cloud/amazon/ec2_elb_lb.py
@@ -44,6 +44,16 @@ options:
       - Purge existing listeners on ELB that are not found in listeners
     required: false
     default: true
+  instance_ids:
+    description:
+      - List of instance ids to attach to this ELB
+    required: false
+    default: false
+  purge_instance_ids:
+    description:
+      - Purge existing instance ids on ELB that are not found in instance_ids
+    required: false
+    default: false
   zones:
     description:
       - List of availability zones to enable on this ELB
@@ -165,6 +175,7 @@ EXAMPLES = """
       - protocol: http # options are http, https, ssl, tcp
         load_balancer_port: 80
         instance_port: 80
+        proxy_protocol: True
       - protocol: https
         load_balancer_port: 443
         instance_protocol: http # optional, defaults to value of protocol setting
@@ -179,6 +190,9 @@ EXAMPLES = """
     name: "test-vpc"
     scheme: internal
     state: present
+    instance_ids:
+      - i-abcd1234
+    purge_instance_ids: true
     subnets:
       - subnet-abcd1234
       - subnet-1a2b3c4d
@@ -354,12 +368,15 @@ class ElbManager(object):
                  scheme="internet-facing", connection_draining_timeout=None,
                  idle_timeout=None,
                  cross_az_load_balancing=None, access_logs=None,
-                 stickiness=None, wait=None, wait_timeout=None, region=None, **aws_connect_params):
+                 stickiness=None, wait=None, wait_timeout=None, region=None,
+                 instance_ids=None, purge_instance_ids=None, **aws_connect_params):
 
         self.module = module
         self.name = name
         self.listeners = listeners
         self.purge_listeners = purge_listeners
+        self.instance_ids = instance_ids
+        self.purge_instance_ids = purge_instance_ids
         self.zones = zones
         self.purge_zones = purge_zones
         self.security_group_ids = security_group_ids
@@ -408,6 +425,10 @@ class ElbManager(object):
             self._set_access_log()
         # add sitcky options
         self.select_stickiness_policy()
+        # ensure backend server policies are correct
+        self._set_backend_policies()
+        # set/remove instance ids
+        self._set_instance_ids()
 
     def ensure_gone(self):
         """Destroy the ELB"""
@@ -456,6 +477,8 @@ class ElbManager(object):
                 'hosted_zone_id': check_elb.canonical_hosted_zone_name_id,
                 'lb_cookie_policy': lb_cookie_policy,
                 'app_cookie_policy': app_cookie_policy,
+                'proxy_policy': self._get_proxy_protocol_policy(),
+                'backends': self._get_backend_policies(),
                 'instances': [instance.id for instance in check_elb.instances],
                 'out_of_service_count': 0,
                 'in_service_count': 0,
@@ -968,6 +991,84 @@ class ElbManager(object):
             else:
                 self._set_listener_policy(listeners_dict)
 
+    def _get_backend_policies(self):
+        """Get a list of backend policies"""
+        return [
+            str(backend.instance_port) + ':' + policy.policy_name for backend in self.elb.backends
+            for policy in backend.policies
+        ]
+
+    def _set_backend_policies(self):
+        """Sets policies for all backends"""
+        ensure_proxy_protocol = False
+        replace = []
+        backend_policies = self._get_backend_policies()
+
+        # Find out what needs to be changed
+        for listener in self.listeners:
+            want = False
+
+            if 'proxy_protocol' in listener and listener['proxy_protocol']:
+                ensure_proxy_protocol = True
+                want = True
+
+            if str(listener['instance_port']) + ':ProxyProtocol-policy' in backend_policies:
+                if not want:
+                    replace.append({'port': listener['instance_port'], 'policies': []})
+            elif want:
+                replace.append({'port': listener['instance_port'], 'policies': ['ProxyProtocol-policy']})
+
+        # enable or disable proxy protocol
+        if ensure_proxy_protocol:
+            self._set_proxy_protocol_policy()
+
+        # Make the backend policies so
+        for item in replace:
+            self.elb_conn.set_lb_policies_of_backend_server(self.elb.name, item['port'], item['policies'])
+            self.changed = True
+
+    def _get_proxy_protocol_policy(self):
+        """Find out if the elb has a proxy protocol enabled"""
+        for policy in self.elb.policies.other_policies:
+            if policy.policy_name == 'ProxyProtocol-policy':
+                return policy.policy_name
+
+        return None
+
+    def _set_proxy_protocol_policy(self):
+        """Install a proxy protocol policy if needed"""
+        proxy_policy = self._get_proxy_protocol_policy()
+
+        if proxy_policy is None:
+            self.elb_conn.create_lb_policy(
+                self.elb.name, 'ProxyProtocol-policy', 'ProxyProtocolPolicyType', {'ProxyProtocol': True}
+            )
+            self.changed = True
+
+        # TODO: remove proxy protocol policy if not needed anymore? There is no side effect to leaving it there
+
+    def _diff_list(self, a, b):
+        """Find the entries in list a that are not in list b"""
+        b = set(b)
+        return [aa for aa in a if aa not in b]
+
+    def _set_instance_ids(self):
+        """Register or deregister instances from an lb instance"""
+        assert_instances = self.instance_ids or []
+
+        has_instances = [has_instance.id for has_instance in self.elb.instances]
+
+        add_instances = self._diff_list(assert_instances, has_instances)
+        if add_instances:
+            self.elb_conn.register_instances(self.elb.name, add_instances)
+            self.changed = True
+
+        if self.purge_instance_ids:
+            remove_instances = self._diff_list(has_instances, assert_instances)
+            if remove_instances:
+                self.elb_conn.deregister_instances(self.elb.name, remove_instances)
+                self.changed = True
+
     def _get_health_check_target(self):
         """Compose target string from healthcheck parameters"""
         protocol = self.health_check['ping_protocol'].upper()
@@ -986,6 +1087,8 @@ def main():
             name={'required': True},
             listeners={'default': None, 'required': False, 'type': 'list'},
             purge_listeners={'default': True, 'required': False, 'type': 'bool'},
+            instance_ids={'default': None, 'required': False, 'type': 'list'},
+            purge_instance_ids={'default': False, 'required': False, 'type': 'bool'},
             zones={'default': None, 'required': False, 'type': 'list'},
             purge_zones={'default': False, 'required': False, 'type': 'bool'},
             security_group_ids={'default': None, 'required': False, 'type': 'list'},
@@ -1020,6 +1123,8 @@ def main():
     state = module.params['state']
     listeners = module.params['listeners']
     purge_listeners = module.params['purge_listeners']
+    instance_ids = module.params['instance_ids']
+    purge_instance_ids = module.params['purge_instance_ids']
     zones = module.params['zones']
     purge_zones = module.params['purge_zones']
     security_group_ids = module.params['security_group_ids']
@@ -1060,13 +1165,16 @@ def main():
         except boto.exception.NoAuthHandlerFound, e:
             module.fail_json(msg = str(e))
 
+
     elb_man = ElbManager(module, name, listeners, purge_listeners, zones,
                          purge_zones, security_group_ids, health_check,
                          subnets, purge_subnets, scheme,
                          connection_draining_timeout, idle_timeout,
                          cross_az_load_balancing,
                          access_logs, stickiness, wait, wait_timeout,
-                         region=region, **aws_connect_params)
+                         region=region, instance_ids=instance_ids, purge_instance_ids=purge_instance_ids,
+                         **aws_connect_params)
+
 
     # check for unsupported attributes for this version of boto
     if cross_az_load_balancing and not elb_man._check_attribute_support('cross_zone_load_balancing'):