1 files changed, 5 insertions, 3 deletions

diff --git a/bin/slave-build-connect b/bin/slave-build-connect
index 4fff765..555b694 100755
--- a/bin/slave-build-connect
+++ b/bin/slave-build-connect
@@ -2,8 +2,10 @@
 
 . "${0%/*}/../conf/default.conf"
 
-if [ "${SSH_ORIGINAL_COMMAND% *}" == "get-assignment" ]; then
-  /bin/bash -c "${base_dir}/bin/${SSH_ORIGINAL_COMMAND}" "$@"
+if [ "${SSH_ORIGINAL_COMMAND%% *}" == "get-assignment" ] || \
+  [ "${SSH_ORIGINAL_COMMAND%% *}" == "return-assignment" ]; then
+  /bin/bash -c "${base_dir}/bin/${SSH_ORIGINAL_COMMAND}"
 else
-  >&2 echo "Invalid command: ${SSH_ORIGINAL_COMMAND} $@"
+  >&2 echo "Invalid command: '${SSH_ORIGINAL_COMMAND%% *}'"
+  exit 1
 fi