diff --git a/global/overlay/etc/puppet/cosmos-rules.yaml b/global/overlay/etc/puppet/cosmos-rules.yaml
index 2fcde36c..4b8418c4 100644
--- a/global/overlay/etc/puppet/cosmos-rules.yaml
+++ b/global/overlay/etc/puppet/cosmos-rules.yaml
@@ -1305,6 +1305,7 @@ proxy-qa-sto1-1.komreg.net:
    eid::proxy:
       version: 3.0.8_hsm2
       service_name: qa.proxy.eidas.swedenconnect.se
+      environment: qa
       use_hsm: false
 
 ### TEST environment ####
@@ -1374,6 +1375,7 @@ proxy-test-sto3-1.komreg.net:
    eid::proxy:
       version: 3.0.8_hsm2
       service_name: test.proxy.eidas.swedenconnect.se
+      environment: test
    sunet::frontend::register_sites:
      sites:
        'test.proxy.eidas.swedenconnect.se':
@@ -1388,6 +1390,7 @@ proxy-test-sto3-2.komreg.net:
    eid::proxy:
       version: 3.0.8_hsm2
       service_name: test.proxy.eidas.swedenconnect.se
+      environment: test
    sunet::frontend::register_sites:
      sites:
        'test.proxy.eidas.swedenconnect.se':
diff --git a/global/overlay/etc/puppet/modules/eid/manifests/proxy.pp b/global/overlay/etc/puppet/modules/eid/manifests/proxy.pp
index 619ac70f..e53cf7cd 100644
--- a/global/overlay/etc/puppet/modules/eid/manifests/proxy.pp
+++ b/global/overlay/etc/puppet/modules/eid/manifests/proxy.pp
@@ -1,20 +1,22 @@
 # This puppet manifest is used to configure Sweden Connect proxy servers
 
+# @param environment          The environment that the server belongs to. (referenced in compose file)
 # @param version              Version of the docker image to use. (referenced in compose file)
 # @param service_name         Name of the service, for example qa.proxy.eidas.swedenconnect.se
 # @param server_fqdn          The FQDN of the server. (referenced in compose file)
 # @param proxy_directory      The directory where all proxy related config and files are stored. (referenced in compose file)
 # @param spring_config_param  Used as parameter name in compose file
-# @param country               Used while creating directories and referenced in compsose file
+# @param country              Used while creating directories and referenced in compsose file
 # @param use_hsm              Configure if HSM is used or not
 class eid::proxy (
-  String  $version = '',
-  String  $service_name = '',
-  String  $server_fqdn = $facts['networking']['fqdn'],
-  String  $proxy_directory = '/opt/eidas-proxy',
-  String  $spring_config_param = 'SPRING_CONFIG_ADDITIONAL_LOCATION',
-  String  $country = 'se',
-  Boolean $use_hsm=true,
+  Enum['test', 'qa', 'prod'] $environment,
+  String                     $version = '',
+  String                     $service_name = '',
+  String                     $server_fqdn = $facts['networking']['fqdn'],
+  String                     $proxy_directory = '/opt/eidas-proxy',
+  String                     $spring_config_param = 'SPRING_CONFIG_ADDITIONAL_LOCATION',
+  String                     $country = 'se',
+  Boolean                    $use_hsm=true,
 ) {
 
   if $use_hsm {
@@ -25,12 +27,14 @@ class eid::proxy (
 
   if $version and $service_name and $eidas_proxy_oidc_rp_jks  != 'NOT_SET_IN_HIERA' and $proxy_service_cookie_encrypt_pw != 'NOT_SET_IN_HIERA'{
 
-    sunet::nftables::allow { 'allow-http-from-any':
-      from => any,
+    # Allow HTTP/HTTPS from load balancer servers
+    $lb_ips = hiera_array("lb_${environment}_servers",[])
+    sunet::nftables::allow { 'allow-http-from-lbs':
+      from => $lb_ips,
       port => 80,
     }
-    sunet::nftables::allow { 'allow-https-from-any':
-      from => any,
+    sunet::nftables::allow { 'allow-https-from-lbs':
+      from => $lb_ips,
       port => 443,
     }