Allow registration of a command instance

1 files changed, 27 insertions, 0 deletions

diff --git a/src/main/java/com/gitblit/transport/ssh/commands/DispatchCommand.java b/src/main/java/com/gitblit/transport/ssh/commands/DispatchCommand.java
index f8503f89..76012a21 100644
--- a/src/main/java/com/gitblit/transport/ssh/commands/DispatchCommand.java
+++ b/src/main/java/com/gitblit/transport/ssh/commands/DispatchCommand.java
@@ -144,6 +144,26 @@ public abstract class DispatchCommand extends BaseCommand implements ExtensionPo
 		commands.add(cmd);
 	}
 
+	/**
+	 * Registers a command as long as the user is permitted to execute it.
+	 *
+	 * @param user
+	 * @param cmd
+	 */
+	protected void registerCommand(UserModel user, BaseCommand cmd) {
+		if (!cmd.getClass().isAnnotationPresent(CommandMetaData.class)) {
+			throw new RuntimeException(MessageFormat.format("{0} must be annotated with {1}!", cmd.getName(),
+					CommandMetaData.class.getName()));
+		}
+		CommandMetaData meta = cmd.getClass().getAnnotation(CommandMetaData.class);
+		if (meta.admin() && !user.canAdmin()) {
+			log.debug(MessageFormat.format("excluding admin command {0} for {1}", meta.name(), user.username));
+			return;
+		}
+		commands.add(cmd.getClass());
+		instantiated.add(cmd);
+	}
+
 	private Map<String, Class<? extends BaseCommand>> getMap() {
 		if (map == null) {
 			map = Maps.newHashMapWithExpectedSize(commands.size());
@@ -222,6 +242,13 @@ public abstract class DispatchCommand extends BaseCommand implements ExtensionPo
 			throw new UnloggedFailure(1, msg);
 		}
 
+		for (BaseCommand cmd : instantiated) {
+			// use an already instantiated command
+			if (cmd.getClass().equals(c)) {
+				return cmd;
+			}
+		}
+
 		BaseCommand cmd = null;
 		try {
 			cmd = c.newInstance();