Allows me to override the REMOTE_ADDR too.

ms-digital-labs · Sep 12, 2015 · ca4122b · ca4122b
1 parent f09994b
commit ca4122b
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 5 deletions.
diff --git a/lib/rack/cloudflare_ip.rb b/lib/rack/cloudflare_ip.rb
@@ -2,14 +2,20 @@
 
 module Rack
   class CloudflareIp
-    def initialize(app)
+    def initialize(app, override_remote_address: false)
       @app = app
+      @override_remote_address = override_remote_address
     end
 
     def call(env)
       if env.has_key?("HTTP_CF_CONNECTING_IP")
-        env["HTTP_ORIGINAL_X_FORWARDED_FOR"] = env["HTTP_X_FORWARDED_FOR"]
-        env["HTTP_X_FORWARDED_FOR"] = env["HTTP_CF_CONNECTING_IP"]
+        if @override_remote_address
+          env["ORIGINAL_REMOTE_ADDR"] = env["REMOTE_ADDR"]
+          env["REMOTE_ADDR"] = env["HTTP_CF_CONNECTING_IP"]
+        else
+          env["HTTP_ORIGINAL_X_FORWARDED_FOR"] = env["HTTP_X_FORWARDED_FOR"]
+          env["HTTP_X_FORWARDED_FOR"] = env["HTTP_CF_CONNECTING_IP"]
+        end
       end
       @app.call(env)
     end

diff --git a/spec/rack/cloudflare_ip_spec.rb b/spec/rack/cloudflare_ip_spec.rb
@@ -9,7 +9,7 @@ def app
 
   def base_app
     lambda { |env|
-      headers = env.select { |k, _| k =~ /^HTTP_.*/ }
+      headers = env.select { |k, _| k =~ /^(HTTP_|ORIGINAL_|REMOTE_)/ }
       [200, headers, []]
     }
   end
@@ -21,7 +21,8 @@ def base_app
   context "with HTTP_CF_CONNECTING_IP header" do
     let(:headers) { {
       "HTTP_CF_CONNECTING_IP" => "123.123.123.123",
-      "HTTP_X_FORWARDED_FOR" => "234.234.234.234"
+      "HTTP_X_FORWARDED_FOR" => "234.234.234.234",
+      "REMOTE_ADDR" => "222.222.222.222",
     } }
 
     it "overwrites the HTTP_X_FORWARDED_FOR header" do
@@ -32,6 +33,10 @@ def base_app
       expect(last_response.headers["HTTP_ORIGINAL_X_FORWARDED_FOR"])
               .to eq("234.234.234.234")
     end
+
+    it "leaves the remote address alone" do
+      expect(last_response.headers["REMOTE_ADDR"]).to eq("222.222.222.222")
+    end
   end
 
   context "without HTTP_CF_CONNECTING_IP header" do
@@ -43,4 +48,36 @@ def base_app
       expect(last_response.headers["HTTP_X_FORWARDED_FOR"]).to eq("234.234.234.234")
     end
   end
+
+  context "with override_remote_address set" do
+    def app
+      Rack::CloudflareIp.new(base_app, override_remote_address: true)
+    end
+
+    context "with HTTP_CF_CONNECTING_IP header" do
+      let(:headers) { {
+        "HTTP_CF_CONNECTING_IP" => "123.123.123.123",
+        "REMOTE_ADDR" => "234.234.234.234"
+      } }
+
+      it "overwrites the REMOTE_ADDR env var" do
+        expect(last_response.headers["REMOTE_ADDR"]).to eq("123.123.123.123")
+      end
+
+      it "saves the original header in ORIGINAL_REMOTE_ADDR" do
+        expect(last_response.headers["ORIGINAL_REMOTE_ADDR"]).to eq("234.234.234.234")
+      end
+    end
+
+    context "without HTTP_CF_CONNECTING_IP header" do
+      let(:headers) { {
+        "REMOTE_ADDR" => "234.234.234.234"
+      } }
+
+      it "doesn't modify the HTTP_X_FORWARDED_FOR header" do
+        expect(last_response.headers["REMOTE_ADDR"]).to eq("234.234.234.234")
+      end
+    end
+
+  end
 end