diff --git a/aws_cloudfront_distribution.tf b/aws_cloudfront_distribution.tf
index 6184e0c..6bbfaa4 100644
--- a/aws_cloudfront_distribution.tf
+++ b/aws_cloudfront_distribution.tf
@@ -8,7 +8,7 @@ resource "aws_cloudfront_distribution" "s3_distribution" {
       origin_access_identity = aws_cloudfront_origin_access_identity.current.cloudfront_access_identity_path
     }
   }
-  comment = "${var.distribution_name} distribution"
+  comment         = "${var.distribution_name} distribution"
   enabled         = true
   is_ipv6_enabled = true
 
@@ -40,7 +40,7 @@ resource "aws_cloudfront_distribution" "s3_distribution" {
 
   #caching
   default_cache_behavior {
-   response_headers_policy_id = aws_cloudfront_response_headers_policy.security_headers_policy.id
+    response_headers_policy_id = var.response_header_policy_enable ? one(aws_cloudfront_response_headers_policy.security_headers_policy).id : ""
 
     min_ttl     = var.cloudfront_cache_min_ttl
     default_ttl = var.cloudfront_cache_default_ttl
@@ -68,12 +68,11 @@ resource "aws_cloudfront_distribution" "s3_distribution" {
     target_origin_id       = "${data.aws_s3_bucket.origin_bucket.id}-origin"
     viewer_protocol_policy = "redirect-to-https"
 
-    dynamic "lambda_function_association" {
-      for_each = var.lambda_function_association
+    dynamic "function_association" {
+      for_each = var.function_associations
       content {
-        event_type   = lambda_function_association.value.event_type
-        include_body = lookup(lambda_function_association.value, "include_body", null)
-        lambda_arn   = lambda_function_association.value.lambda_arn
+        event_type   = function_association.value.event_type
+        function_arn = function_association.value.function_arn
       }
     }
   }
@@ -106,7 +105,7 @@ resource "aws_cloudfront_origin_access_identity" "current" {}
     # https://infosec.mozilla.org/guidelines/web_security#x-frame-options
     frame_options {
       frame_option = "DENY"
-      override = true
+      override     = true
     }
     # https://infosec.mozilla.org/guidelines/web_security#referrer-policy
     # referrer_policy {
@@ -122,9 +121,9 @@ resource "aws_cloudfront_origin_access_identity" "current" {}
     # https://infosec.mozilla.org/guidelines/web_security#http-strict-transport-security
     strict_transport_security {
       access_control_max_age_sec = "63072000"
-      include_subdomains = true
-      preload = true
-      override = true
+      include_subdomains         = true
+      preload                    = true
+      override                   = true
     }
     # https://infosec.mozilla.org/guidelines/web_security#content-security-policy
     # content_security_policy {
diff --git a/example/module_cloudfront_example.tf b/example/module_cloudfront_example.tf
index 2419348..fd4c93d 100644
--- a/example/module_cloudfront_example.tf
+++ b/example/module_cloudfront_example.tf
@@ -6,4 +6,4 @@ module "cloudfront_example" {
   hosted_zone_name                  = "domain-name.com"
   common_tags                       = var.common_tags
   cloudfront_cache_compress_content = var.cloudfront_cache_compress_content
-}
\ No newline at end of file
+}
diff --git a/variables.tf b/variables.tf
index 86444eb..1e7a256 100644
--- a/variables.tf
+++ b/variables.tf
@@ -76,14 +76,12 @@ variable "ttl" {
   default = "300"
 }
 
-variable "lambda_function_association" {
+variable "function_associations" {
+  description = "A config block that triggers a function with specific actions"
   type = list(object({
     event_type   = string
-    include_body = bool
-    lambda_arn   = string
+    function_arn   = string
   }))
-
-  description = "A config block that triggers a lambda function with specific actions"
   default     = []
 }