From 349dbb2b818eb83e92204cb07067b9aeeaec6a4a Mon Sep 17 00:00:00 2001
From: Kirill Sysoev <kirill@binary.com>
Date: Tue, 27 Sep 2022 06:14:00 +0000
Subject: [PATCH 1/2] Adds possibility to provide custom dispather

---
 lib/Mojo/WebSocketProxy/Dispatcher.pm | 40 +++++++++++++++++++--------
 1 file changed, 28 insertions(+), 12 deletions(-)

diff --git a/lib/Mojo/WebSocketProxy/Dispatcher.pm b/lib/Mojo/WebSocketProxy/Dispatcher.pm
index 707f1ff..17b9ebd 100644
--- a/lib/Mojo/WebSocketProxy/Dispatcher.pm
+++ b/lib/Mojo/WebSocketProxy/Dispatcher.pm
@@ -128,18 +128,24 @@ sub on_message {
         $c->send({json => $err}, $req_storage);
         return $c->_run_hooks($config->{after_dispatch} || [])->retain;
     }
-
-    my $action = $c->dispatch($args) or do {
-        my $err = $c->wsp_error('error', UnrecognisedRequest => 'Unrecognised request');
-        $c->send({json => $err}, $req_storage);
-        return $c->_run_hooks($config->{after_dispatch} || [])->retain;
-    };
-
-    @{$req_storage}{keys %$action} = (values %$action);
-    $req_storage->{method} = $req_storage->{name};
-
+    
     # main processing pipeline
-    my $f = $c->before_forward($req_storage)->transform(
+    my $f = $c->dispatch($args)->then(sub {
+        my ($action) = @_; 
+        unless ($action) {
+            my $err = $c->wsp_error('error', UnrecognisedRequest => 'Unrecognised request');
+            $c->send({json => $err}, $req_storage);
+            return $c->_run_hooks($config->{after_dispatch} || [])->then(sub {
+                return Future::Mojo->fail('UNRECOGNISED_REQUEST');
+            });
+        }
+
+        @{$req_storage}{keys %$action} = (values %$action);
+        $req_storage->{method} = $req_storage->{name};
+        return Future::Mojo->done;
+    })->then(sub {
+        return $c->before_forward($req_storage);
+    })->transform(
         done => sub {
             # Note that we completely ignore the return value of ->before_forward here.
             return $req_storage->{instead_of_forward}->($c, $req_storage) if $req_storage->{instead_of_forward};
@@ -170,6 +176,11 @@ sub on_message {
         }
     )->on_fail(
         sub {
+            my ($err) = @_;
+
+            # we don't want to flood logs if clients sent incorrect requests.
+            return if $err and $err eq 'UNRECOGNISED_REQUEST';
+
             $c->app->log->error("An error occurred handling on_message. Error @_");
         })->retain;
 }
@@ -221,12 +232,17 @@ sub dispatch {
     my $log = $c->app->log;
     $log->debug("websocket got json " . $c->dumper($args));
 
+    my $actions = $c->wsp_config->{actions};
+    if (blessed $actions && $actions->can('dispatch')) {
+        return $actions->dipatch($actions);
+    }
+
     my ($action) =
         sort { $a->{order} <=> $b->{order} }
         grep { defined }
         map  { $c->wsp_config->{actions}->{$_} } keys %$args;
 
-    return $action;
+    return Future::Mojo->done($action);
 }
 
 sub forward {

From fbb5e0c954199343a16288ebc43b0ac61e96b7e1 Mon Sep 17 00:00:00 2001
From: Kirill Sysoev <kirill@binary.com>
Date: Wed, 28 Sep 2022 01:42:09 +0000
Subject: [PATCH 2/2] Adds possibility custom dispatcher

---
 lib/Mojo/WebSocketProxy/Config.pm        | 11 +++++++++++
 lib/Mojo/WebSocketProxy/Dispatcher.pm    |  5 ++---
 lib/Mojolicious/Plugin/WebSocketProxy.pm |  9 +++++++--
 3 files changed, 20 insertions(+), 5 deletions(-)

diff --git a/lib/Mojo/WebSocketProxy/Config.pm b/lib/Mojo/WebSocketProxy/Config.pm
index 39d7030..e31e386 100644
--- a/lib/Mojo/WebSocketProxy/Config.pm
+++ b/lib/Mojo/WebSocketProxy/Config.pm
@@ -69,6 +69,17 @@ sub add_backend {
     return;
 }
 
+sub add_dispatcher {
+    my ($self, $dispatcher) = @_;
+
+    die 'Custom dipatcher should have dispatch method' 
+            unless blessed($dispatcher) and $dispatcher->can('dispatch');
+    
+    $self->{dispatcher} = $dispatcher;
+
+    return;
+}
+
 1;
 
 __END__
diff --git a/lib/Mojo/WebSocketProxy/Dispatcher.pm b/lib/Mojo/WebSocketProxy/Dispatcher.pm
index 17b9ebd..f7768ff 100644
--- a/lib/Mojo/WebSocketProxy/Dispatcher.pm
+++ b/lib/Mojo/WebSocketProxy/Dispatcher.pm
@@ -232,9 +232,8 @@ sub dispatch {
     my $log = $c->app->log;
     $log->debug("websocket got json " . $c->dumper($args));
 
-    my $actions = $c->wsp_config->{actions};
-    if (blessed $actions && $actions->can('dispatch')) {
-        return $actions->dipatch($actions);
+    if ($c->wsp_config->{dispatcher}) {
+        return Future::Mojo->wrap($c->wsp_config->{dispatcher}->dipatch($actions));
     }
 
     my ($action) =
diff --git a/lib/Mojolicious/Plugin/WebSocketProxy.pm b/lib/Mojolicious/Plugin/WebSocketProxy.pm
index 9a7790e..b30a08b 100644
--- a/lib/Mojolicious/Plugin/WebSocketProxy.pm
+++ b/lib/Mojolicious/Plugin/WebSocketProxy.pm
@@ -3,6 +3,8 @@ package Mojolicious::Plugin::WebSocketProxy;
 use strict;
 use warnings;
 
+use Scalar::Util qw(blessed);
+
 use Mojo::Base 'Mojolicious::Plugin';
 use Mojo::WebSocketProxy::Backend;
 use Mojo::WebSocketProxy::Config;
@@ -55,11 +57,14 @@ sub register {
         $_->websocket('/')->to('Dispatcher#open_connection', namespace => 'Mojo::WebSocketProxy');
     }
 
-    my $actions           = delete $config->{actions};
     my $dispatcher_config = Mojo::WebSocketProxy::Config->new;
     $dispatcher_config->init($config);
+    # if custom dispatcher is provided, we will use it
+    if (my $dispatcher = delete $config->{dispatcher}) {
+        $dispatcher_config->add_dispatcher($dispatcher);
+    } elsif (my $actions = delete $config->{actions}) {
+        die 'Invalid format of actions parameter' unless ref $actions eq 'ARRAY';
 
-    if (ref $actions eq 'ARRAY') {
         for (my $i = 0; $i < @$actions; $i++) {
             $dispatcher_config->add_action($actions->[$i], $i);
         }