Custom URL scheme support

server.rb

@@ -12,6 +12,9 @@
 
 Rack::Attack.cache.store = ActiveSupport::Cache.lookup_store :redis_store
 
+http_redirect_regexp = /^https?:\/.+/
+scheme_redirect_regexp = /^\S+:\/.+/
+
 if ENV['UDL_THROTTLE_LIMIT'].present? && ENV['UDL_THROTTLE_PERIOD'].present?
   limit = ENV['UDL_THROTTLE_LIMIT'].to_i
   period = ENV['UDL_THROTTLE_PERIOD'].to_i
@@ -38,7 +41,14 @@
 
 get '/' do
   begin
-    redirect URI(params[:r])
+    requested_url = params[:r]
+    if http_redirect_regexp.match(requested_url)
+      redirect URI(requested_url)
+    elsif scheme_redirect_regexp.match(requested_url)
+      redirect requested_url
+    else
+      raise 'Requested URI is invalid'
+    end
   rescue => error
     @error = error
     logger.info @error.inspect
@@ -72,9 +82,16 @@
 
 get '/*' do
   begin
-    target_url = URI(params['splat'].first.gsub('https:/', 'https://'))
-    raise 'Invalid redirect URL' unless target_url.host.present?
-    redirect target_url
+    requested_url = params['splat'].first
+    if http_redirect_regexp.match(requested_url)
+      target_url = URI(requested_url.gsub(':/', '://'))
+      raise 'Invalid redirect URL' if target_url&.host.nil?
+      redirect target_url
+    elsif scheme_redirect_regexp.match(requested_url)
+      redirect requested_url.gsub(':/', ':///')
+    else
+      raise 'Target redirect location must have a scheme'
+    end
   rescue => error
     @error = error
     logger.info @error.inspect

spec/app_spec.rb

@@ -12,26 +12,41 @@ def app
   end
 
   context "success" do
-    it "redirects to the r parameter if valid" do
-      target_url = "https://dev.to/fdoxyz"
-      get "/?r=#{target_url}"
-      expect(last_response).to be_redirect
-      expect(last_response.location).to eq(target_url)
+    let(:http_target_url) { "https://dev.to/fdocr" }
+    let(:custom_scheme_target_url) { "ms-mobile-apps:///providers/Microsoft.PowerApps/apps/123?tenantId=456" }
+
+    context "using r query parameter" do
+      it "redirects if URL is valid" do
+        get "/?r=#{http_target_url}"
+        expect(last_response).to be_redirect
+        expect(last_response.location).to eq(http_target_url)
+      end
+
+      it "redirects to MS Power Apps schemes" do 
+        get "/?r=#{custom_scheme_target_url}"
+        expect(last_response).to be_redirect
+      end
     end
 
-    it "redirects when passing target in REST first level param" do
-      target_url = "https://dev.to/fdoxyz"
-      get "/#{target_url}"
-      expect(last_response).to be_redirect
-      expect(last_response.location).to eq(target_url)
+    context "using REST first level param" do
+      it "redirects if URL is valid" do
+        get "/#{http_target_url}"
+        expect(last_response).to be_redirect
+        expect(last_response.location).to eq(http_target_url)
+      end
+
+      it "redirects to MS Power Apps schemes" do
+        get "/#{custom_scheme_target_url}"
+        expect(last_response).to be_redirect
+      end
     end
   end
 
   context "failure" do
     after(:each) do
       expect(last_response).to be_ok
       expect(last_response.body).to include('Something went wrong')
-      expect(last_response.body).to include('Check out the <a href="https://github.com/fdoxyz/udl-server#Troubleshooting">README</a> for more details')
+      expect(last_response.body).to include('Check out the <a href="https://github.com/fdocr/udl-server#Troubleshooting">README</a> for more details')
     end
 
     it "renders fallback page if r parameter isn't available" do

views/fallback.erb

@@ -8,7 +8,7 @@
     <h1>Something went wrong</h1>
 
     <p><code><%= @error %></code></p>
-    <p>Check out the <a href="https://github.com/fdoxyz/udl-server#Troubleshooting">README</a> for more details</p>
+    <p>Check out the <a href="https://github.com/fdocr/udl-server#Troubleshooting">README</a> for more details</p>
 
   </body>
 </html>