From 59e621d541746ff5f2576541abc1a201afcbc15f Mon Sep 17 00:00:00 2001
From: James Moger <james.moger@gitblit.com>
Date: Thu, 10 Apr 2014 18:58:08 -0400
Subject: [PATCH] Move dispatcher creation to SshCommandFactory and revise permission checks

---
 src/main/java/com/gitblit/transport/ssh/SshDaemon.java                |   47 --
 src/main/java/com/gitblit/transport/ssh/commands/BaseCommand.java     |   10 
 src/main/java/com/gitblit/transport/ssh/commands/DispatchCommand.java |  364 ++++++++++++-----------
 src/main/java/com/gitblit/transport/ssh/SshCommandFactory.java        |  433 ++++++++++++++++------------
 4 files changed, 432 insertions(+), 422 deletions(-)

diff --git a/src/main/java/com/gitblit/transport/ssh/SshCommandFactory.java b/src/main/java/com/gitblit/transport/ssh/SshCommandFactory.java
index 788bdfb..da57f76 100644
--- a/src/main/java/com/gitblit/transport/ssh/SshCommandFactory.java
+++ b/src/main/java/com/gitblit/transport/ssh/SshCommandFactory.java
@@ -34,7 +34,21 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.gitblit.git.GitblitReceivePackFactory;
+import com.gitblit.git.GitblitUploadPackFactory;
+import com.gitblit.git.RepositoryResolver;
+import com.gitblit.manager.IGitblit;
+import com.gitblit.models.UserModel;
+import com.gitblit.transport.ssh.commands.AddKeyCommand;
+import com.gitblit.transport.ssh.commands.CreateRepository;
 import com.gitblit.transport.ssh.commands.DispatchCommand;
+import com.gitblit.transport.ssh.commands.Receive;
+import com.gitblit.transport.ssh.commands.RemoveKeyCommand;
+import com.gitblit.transport.ssh.commands.ReviewCommand;
+import com.gitblit.transport.ssh.commands.SetAccountCommand;
+import com.gitblit.transport.ssh.commands.Upload;
+import com.gitblit.transport.ssh.commands.VersionCommand;
+import com.gitblit.utils.IdGenerator;
 import com.gitblit.utils.WorkQueue;
 import com.google.common.util.concurrent.Atomics;
 
@@ -44,224 +58,261 @@
  *
  */
 public class SshCommandFactory implements CommandFactory {
-  private static final Logger logger = LoggerFactory
-      .getLogger(SshCommandFactory.class);
-  private final ScheduledExecutorService startExecutor;
+	private static final Logger logger = LoggerFactory.getLogger(SshCommandFactory.class);
 
-  private DispatchCommand dispatcher;
+	private final IGitblit gitblit;
+	private final PublicKeyAuthenticator keyAuthenticator;
+	private final ScheduledExecutorService startExecutor;
 
-	public SshCommandFactory(
-	    WorkQueue workQueue,
-	    DispatchCommand d) {
-		this.dispatcher = d;
-		int threads = 2;//cfg.getInt("sshd","commandStartThreads", 2);
-	    startExecutor = workQueue.createQueue(threads, "SshCommandStart");
+	public SshCommandFactory(IGitblit gitblit, PublicKeyAuthenticator keyAuthenticator, IdGenerator idGenerator) {
+		this.gitblit = gitblit;
+		this.keyAuthenticator = keyAuthenticator;
+
+		int threads = 2;// cfg.getInt("sshd","commandStartThreads", 2);
+		WorkQueue workQueue = new WorkQueue(idGenerator);
+		startExecutor = workQueue.createQueue(threads, "SshCommandStart");
+	}
+
+	/**
+	 * Creates the root dispatcher command which builds up the available commands.
+	 *
+	 * @param the client
+	 * @param the command line
+	 * @return the root dispatcher command
+	 */
+	protected DispatchCommand createRootDispatcher(SshDaemonClient client, String cmdLine) {
+		final UserModel user = client.getUser();
+
+		DispatchCommand gitblitCmd = new DispatchCommand();
+		gitblitCmd.registerCommand(user, VersionCommand.class);
+		gitblitCmd.registerCommand(user, AddKeyCommand.class);
+		gitblitCmd.registerCommand(user, RemoveKeyCommand.class);
+		gitblitCmd.registerCommand(user, ReviewCommand.class);
+
+		gitblitCmd.registerCommand(user, CreateRepository.class);
+		gitblitCmd.registerCommand(user, SetAccountCommand.class);
+
+		DispatchCommand gitCmd = new DispatchCommand();
+		gitCmd.registerCommand(user, Upload.class);
+		gitCmd.registerCommand(user, Receive.class);
+
+		DispatchCommand root = new DispatchCommand();
+		root.registerDispatcher("gitblit", gitblitCmd);
+		root.registerDispatcher("git", gitCmd);
+
+		root.setRepositoryResolver(new RepositoryResolver<SshDaemonClient>(gitblit));
+		root.setUploadPackFactory(new GitblitUploadPackFactory<SshDaemonClient>(gitblit));
+		root.setReceivePackFactory(new GitblitReceivePackFactory<SshDaemonClient>(gitblit));
+		root.setAuthenticator(keyAuthenticator);
+
+		root.setContext(new SshCommandContext(client, cmdLine));
+
+		return root;
 	}
 
 	@Override
 	public Command createCommand(final String commandLine) {
-	  return new Trampoline(commandLine);
+		return new Trampoline(commandLine);
 	}
 
-	  private class Trampoline implements Command, SessionAware {
-	    private final String[] argv;
-	    private ServerSession session;
-	    private InputStream in;
-	    private OutputStream out;
-	    private OutputStream err;
-	    private ExitCallback exit;
-	    private Environment env;
-	    private String cmdLine;
-	    private DispatchCommand cmd;
-	    private final AtomicBoolean logged;
-	    private final AtomicReference<Future<?>> task;
+	private class Trampoline implements Command, SessionAware {
+		private final String[] argv;
+		private ServerSession session;
+		private InputStream in;
+		private OutputStream out;
+		private OutputStream err;
+		private ExitCallback exit;
+		private Environment env;
+		private String cmdLine;
+		private DispatchCommand cmd;
+		private final AtomicBoolean logged;
+		private final AtomicReference<Future<?>> task;
 
-	    Trampoline(String line) {
-	      if (line.startsWith("git-")) {
-            line = "git " + line;
-	      }
-	      cmdLine = line;
-	      argv = split(line);
-	      logged = new AtomicBoolean();
-	      task = Atomics.newReference();
-	    }
+		Trampoline(String line) {
+			if (line.startsWith("git-")) {
+				line = "git " + line;
+			}
+			cmdLine = line;
+			argv = split(line);
+			logged = new AtomicBoolean();
+			task = Atomics.newReference();
+		}
 
-	    @Override
-	    public void setSession(ServerSession session) {
-	      this.session = session;
-	    }
+		@Override
+		public void setSession(ServerSession session) {
+			this.session = session;
+		}
 
-	    @Override
+		@Override
 		public void setInputStream(final InputStream in) {
-	      this.in = in;
-	    }
+			this.in = in;
+		}
 
-	    @Override
+		@Override
 		public void setOutputStream(final OutputStream out) {
-	      this.out = out;
-	    }
+			this.out = out;
+		}
 
-	    @Override
+		@Override
 		public void setErrorStream(final OutputStream err) {
-	      this.err = err;
-	    }
+			this.err = err;
+		}
 
-	    @Override
+		@Override
 		public void setExitCallback(final ExitCallback callback) {
-	      this.exit = callback;
-	    }
+			this.exit = callback;
+		}
 
-	    @Override
+		@Override
 		public void start(final Environment env) throws IOException {
-	      this.env = env;
-	      task.set(startExecutor.submit(new Runnable() {
-	        @Override
-			public void run() {
-	          try {
-	            onStart();
-	          } catch (Exception e) {
-	            logger.warn("Cannot start command ", e);
-	          }
-	        }
+			this.env = env;
+			task.set(startExecutor.submit(new Runnable() {
+				@Override
+				public void run() {
+					try {
+						onStart();
+					} catch (Exception e) {
+						logger.warn("Cannot start command ", e);
+					}
+				}
 
-	        @Override
-	        public String toString() {
-	          return "start (user " + session.getUsername() + ")";
-	        }
-	      }));
-	    }
+				@Override
+				public String toString() {
+					return "start (user " + session.getUsername() + ")";
+				}
+			}));
+		}
 
-	    private void onStart() throws IOException {
-	      synchronized (this) {
-		    SshCommandContext ctx = new SshCommandContext(session.getAttribute(SshDaemonClient.KEY), cmdLine);
-	        try {
-	          cmd = dispatcher;
-	          cmd.setArguments(argv);
-	          cmd.setContext(ctx);
-	          cmd.setInputStream(in);
-	          cmd.setOutputStream(out);
-	          cmd.setErrorStream(err);
-	          cmd.setExitCallback(new ExitCallback() {
-	            @Override
-	            public void onExit(int rc, String exitMessage) {
-	              exit.onExit(translateExit(rc), exitMessage);
-	              log(rc);
-	            }
+		private void onStart() throws IOException {
+			synchronized (this) {
+				SshDaemonClient client = session.getAttribute(SshDaemonClient.KEY);
+				try {
+					cmd = createRootDispatcher(client, cmdLine);
+					cmd.setArguments(argv);
+					cmd.setInputStream(in);
+					cmd.setOutputStream(out);
+					cmd.setErrorStream(err);
+					cmd.setExitCallback(new ExitCallback() {
+						@Override
+						public void onExit(int rc, String exitMessage) {
+							exit.onExit(translateExit(rc), exitMessage);
+							log(rc);
+						}
 
-	            @Override
-	            public void onExit(int rc) {
-	              exit.onExit(translateExit(rc));
-	              log(rc);
-	            }
-	          });
-	          cmd.start(env);
-	        } finally {
-		      ctx = null;
-	        }
-	      }
-	    }
+						@Override
+						public void onExit(int rc) {
+							exit.onExit(translateExit(rc));
+							log(rc);
+						}
+					});
+					cmd.start(env);
+				} finally {
+					client = null;
+				}
+			}
+		}
 
-	    private int translateExit(final int rc) {
-	      return rc;
-//
-//	      switch (rc) {
-//	        case BaseCommand.STATUS_NOT_ADMIN:
-//	          return 1;
-//
-//	        case BaseCommand.STATUS_CANCEL:
-//	          return 15 /* SIGKILL */;
-//
-//	        case BaseCommand.STATUS_NOT_FOUND:
-//	          return 127 /* POSIX not found */;
-//
-//	        default:
-//	          return rc;
-//	      }
+		private int translateExit(final int rc) {
+			return rc;
+			//
+			// switch (rc) {
+			// case BaseCommand.STATUS_NOT_ADMIN:
+			// return 1;
+			//
+			// case BaseCommand.STATUS_CANCEL:
+			// return 15 /* SIGKILL */;
+			//
+			// case BaseCommand.STATUS_NOT_FOUND:
+			// return 127 /* POSIX not found */;
+			//
+			// default:
+			// return rc;
+			// }
 
-	    }
+		}
 
-	    private void log(final int rc) {
-	      if (logged.compareAndSet(false, true)) {
-	        //log.onExecute(cmd, rc);
-	        logger.info("onExecute: {} exits with: {}", cmd.getClass().getSimpleName(), rc);
-	      }
-	    }
+		private void log(final int rc) {
+			if (logged.compareAndSet(false, true)) {
+				// log.onExecute(cmd, rc);
+				logger.info("onExecute: {} exits with: {}", cmd.getClass().getSimpleName(), rc);
+			}
+		}
 
-	    @Override
-	    public void destroy() {
-	      Future<?> future = task.getAndSet(null);
-	      if (future != null) {
-	        future.cancel(true);
-//	        destroyExecutor.execute(new Runnable() {
-//	          @Override
-//	          public void run() {
-//	            onDestroy();
-//	          }
-//	        });
-	      }
-	    }
+		@Override
+		public void destroy() {
+			Future<?> future = task.getAndSet(null);
+			if (future != null) {
+				future.cancel(true);
+				// destroyExecutor.execute(new Runnable() {
+				// @Override
+				// public void run() {
+				// onDestroy();
+				// }
+				// });
+			}
+		}
 
-	    private void onDestroy() {
-	      synchronized (this) {
-	        if (cmd != null) {
-	          //final Context old = sshScope.set(ctx);
-	          try {
-	            cmd.destroy();
-	            //log(BaseCommand.STATUS_CANCEL);
-	          } finally {
-	            //ctx = null;
-	            cmd = null;
-	            //sshScope.set(old);
-	          }
-	        }
-	      }
-	    }
-	  }
+		private void onDestroy() {
+			synchronized (this) {
+				if (cmd != null) {
+					// final Context old = sshScope.set(ctx);
+					try {
+						cmd.destroy();
+						// log(BaseCommand.STATUS_CANCEL);
+					} finally {
+						// ctx = null;
+						cmd = null;
+						// sshScope.set(old);
+					}
+				}
+			}
+		}
+	}
 
-	  /** Split a command line into a string array. */
-	  static public String[] split(String commandLine) {
-	    final List<String> list = new ArrayList<String>();
-	    boolean inquote = false;
-	    boolean inDblQuote = false;
-	    StringBuilder r = new StringBuilder();
-	    for (int ip = 0; ip < commandLine.length();) {
-	      final char b = commandLine.charAt(ip++);
-	      switch (b) {
-	        case '\t':
-	        case ' ':
-	          if (inquote || inDblQuote)
-	            r.append(b);
-	          else if (r.length() > 0) {
-	            list.add(r.toString());
-	            r = new StringBuilder();
-	          }
-	          continue;
-	        case '\"':
-	          if (inquote)
-	            r.append(b);
-	          else
-	            inDblQuote = !inDblQuote;
-	          continue;
-	        case '\'':
-	          if (inDblQuote)
-	            r.append(b);
-	          else
-	            inquote = !inquote;
-	          continue;
-	        case '\\':
-	          if (inquote || ip == commandLine.length())
-	            r.append(b); // literal within a quote
-	          else
-	            r.append(commandLine.charAt(ip++));
-	          continue;
-	        default:
-	          r.append(b);
-	          continue;
-	      }
-	    }
-	    if (r.length() > 0) {
-	      list.add(r.toString());
-	    }
-	    return list.toArray(new String[list.size()]);
-	  }
+	/** Split a command line into a string array. */
+	static public String[] split(String commandLine) {
+		final List<String> list = new ArrayList<String>();
+		boolean inquote = false;
+		boolean inDblQuote = false;
+		StringBuilder r = new StringBuilder();
+		for (int ip = 0; ip < commandLine.length();) {
+			final char b = commandLine.charAt(ip++);
+			switch (b) {
+			case '\t':
+			case ' ':
+				if (inquote || inDblQuote)
+					r.append(b);
+				else if (r.length() > 0) {
+					list.add(r.toString());
+					r = new StringBuilder();
+				}
+				continue;
+			case '\"':
+				if (inquote)
+					r.append(b);
+				else
+					inDblQuote = !inDblQuote;
+				continue;
+			case '\'':
+				if (inDblQuote)
+					r.append(b);
+				else
+					inquote = !inquote;
+				continue;
+			case '\\':
+				if (inquote || ip == commandLine.length())
+					r.append(b); // literal within a quote
+				else
+					r.append(commandLine.charAt(ip++));
+				continue;
+			default:
+				r.append(b);
+				continue;
+			}
+		}
+		if (r.length() > 0) {
+			list.add(r.toString());
+		}
+		return list.toArray(new String[list.size()]);
+	}
 }
diff --git a/src/main/java/com/gitblit/transport/ssh/SshDaemon.java b/src/main/java/com/gitblit/transport/ssh/SshDaemon.java
index c11cb1f..c3d4860 100644
--- a/src/main/java/com/gitblit/transport/ssh/SshDaemon.java
+++ b/src/main/java/com/gitblit/transport/ssh/SshDaemon.java
@@ -34,22 +34,9 @@
 
 import com.gitblit.IStoredSettings;
 import com.gitblit.Keys;
-import com.gitblit.git.GitblitReceivePackFactory;
-import com.gitblit.git.GitblitUploadPackFactory;
-import com.gitblit.git.RepositoryResolver;
 import com.gitblit.manager.IGitblit;
-import com.gitblit.transport.ssh.commands.AddKeyCommand;
-import com.gitblit.transport.ssh.commands.CreateRepository;
-import com.gitblit.transport.ssh.commands.DispatchCommand;
-import com.gitblit.transport.ssh.commands.Receive;
-import com.gitblit.transport.ssh.commands.RemoveKeyCommand;
-import com.gitblit.transport.ssh.commands.ReviewCommand;
-import com.gitblit.transport.ssh.commands.SetAccountCommand;
-import com.gitblit.transport.ssh.commands.Upload;
-import com.gitblit.transport.ssh.commands.VersionCommand;
 import com.gitblit.utils.IdGenerator;
 import com.gitblit.utils.StringUtils;
-import com.gitblit.utils.WorkQueue;
 
 import dagger.Module;
 import dagger.ObjectGraph;
@@ -117,45 +104,19 @@
 			addr = new InetSocketAddress(bindInterface, port);
 		}
 
-		PublicKeyAuthenticator publickeyAuthenticator = new PublicKeyAuthenticator(
-				keyManager, gitblit);
+		PublicKeyAuthenticator keyAuthenticator = new PublicKeyAuthenticator(keyManager, gitblit);
+
 		sshd = SshServer.setUpDefaultServer();
 		sshd.setPort(addr.getPort());
 		sshd.setHost(addr.getHostName());
 		sshd.setKeyPairProvider(new PEMGeneratorHostKeyProvider(new File(
 				gitblit.getBaseFolder(), HOST_KEY_STORE).getPath()));
-		sshd.setPublickeyAuthenticator(publickeyAuthenticator);
+		sshd.setPublickeyAuthenticator(keyAuthenticator);
 		sshd.setPasswordAuthenticator(new UsernamePasswordAuthenticator(gitblit));
 		sshd.setSessionFactory(new SshServerSessionFactory());
 		sshd.setFileSystemFactory(new DisabledFilesystemFactory());
 		sshd.setTcpipForwardingFilter(new NonForwardingFilter());
-
-		DispatchCommand gitblitCmd = new DispatchCommand();
-		gitblitCmd.registerCommand(CreateRepository.class);
-		gitblitCmd.registerCommand(VersionCommand.class);
-		gitblitCmd.registerCommand(AddKeyCommand.class);
-		gitblitCmd.registerCommand(RemoveKeyCommand.class);
-		gitblitCmd.registerCommand(SetAccountCommand.class);
-		gitblitCmd.registerCommand(ReviewCommand.class);
-
-		DispatchCommand gitCmd = new DispatchCommand();
-		gitCmd.registerCommand(Upload.class);
-		gitCmd.registerCommand(Receive.class);
-
-		DispatchCommand root = new DispatchCommand();
-		root.registerDispatcher("gitblit", gitblitCmd);
-		root.registerDispatcher("git", gitCmd);
-
-		root.setRepositoryResolver(new RepositoryResolver<SshDaemonClient>(gitblit));
-		root.setUploadPackFactory(new GitblitUploadPackFactory<SshDaemonClient>(gitblit));
-		root.setReceivePackFactory(new GitblitReceivePackFactory<SshDaemonClient>(gitblit));
-		root.setAuthenticator(publickeyAuthenticator);
-
-		SshCommandFactory commandFactory = new SshCommandFactory(
-				new WorkQueue(idGenerator),
-				root);
-
-		sshd.setCommandFactory(commandFactory);
+		sshd.setCommandFactory(new SshCommandFactory(gitblit, keyAuthenticator, idGenerator));
 
 		run = new AtomicBoolean(false);
 	}
diff --git a/src/main/java/com/gitblit/transport/ssh/commands/BaseCommand.java b/src/main/java/com/gitblit/transport/ssh/commands/BaseCommand.java
index baa892c..28dfbdd 100644
--- a/src/main/java/com/gitblit/transport/ssh/commands/BaseCommand.java
+++ b/src/main/java/com/gitblit/transport/ssh/commands/BaseCommand.java
@@ -120,16 +120,6 @@
 		this.exit = callback;
 	}
 
-	protected void provideBaseStateTo(final Command cmd) {
-		if (cmd instanceof BaseCommand) {
-			((BaseCommand) cmd).setContext(ctx);
-		}
-		cmd.setInputStream(in);
-		cmd.setOutputStream(out);
-		cmd.setErrorStream(err);
-		cmd.setExitCallback(exit);
-	}
-
 	protected String getName() {
 		return commandName;
 	}
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 9ffb123..673b576 100644
--- a/src/main/java/com/gitblit/transport/ssh/commands/DispatchCommand.java
+++ b/src/main/java/com/gitblit/transport/ssh/commands/DispatchCommand.java
@@ -26,10 +26,13 @@
 import org.apache.sshd.server.Command;
 import org.apache.sshd.server.Environment;
 import org.kohsuke.args4j.Argument;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import com.gitblit.git.GitblitReceivePackFactory;
 import com.gitblit.git.GitblitUploadPackFactory;
 import com.gitblit.git.RepositoryResolver;
+import com.gitblit.models.UserModel;
 import com.gitblit.transport.ssh.CommandMetaData;
 import com.gitblit.transport.ssh.PublicKeyAuthenticator;
 import com.gitblit.transport.ssh.SshDaemonClient;
@@ -41,197 +44,202 @@
 
 public class DispatchCommand extends BaseCommand {
 
-  @Argument(index = 0, required = false, metaVar = "COMMAND", handler = SubcommandHandler.class)
-  private String commandName;
+	private Logger log = LoggerFactory.getLogger(getClass());
 
-  @Argument(index = 1, multiValued = true, metaVar = "ARG")
-  private List<String> args = new ArrayList<String>();
+	@Argument(index = 0, required = false, metaVar = "COMMAND", handler = SubcommandHandler.class)
+	private String commandName;
 
-  private Set<Class<? extends Command>> commands;
-  private Map<String, Class<? extends Command>> map;
-  private Map<String, Command> root;
+	@Argument(index = 1, multiValued = true, metaVar = "ARG")
+	private List<String> args = new ArrayList<String>();
 
-  public DispatchCommand() {
-	  commands = new HashSet<Class<? extends Command>>();
-  }
+	private Set<Class<? extends Command>> commands;
+	private Map<String, Class<? extends Command>> map;
+	private Map<String, Command> root;
 
-  public void registerDispatcher(String name, Command cmd) {
-	  if (root == null) {
-		  root = Maps.newHashMap();
-	  }
-	  root.put(name, cmd);
-  }
-
-  public void registerCommand(Class<? extends Command> cmd) {
-	  if (!cmd.isAnnotationPresent(CommandMetaData.class)) {
-		  throw new RuntimeException(MessageFormat.format("{0} must be annotated with {1}!",
-				  cmd.getName(), CommandMetaData.class.getName()));
-	  }
-	  commands.add(cmd);
-  }
-
-  private Map<String, Class<? extends Command>> getMap() {
-    if (map == null) {
-      map = Maps.newHashMapWithExpectedSize(commands.size());
-      for (Class<? extends Command> cmd : commands) {
-        CommandMetaData meta = cmd.getAnnotation(CommandMetaData.class);
-        map.put(meta.name(), cmd);
-      }
-    }
-    return map;
-  }
-
-  @Override
-  public void start(Environment env) throws IOException {
-    try {
-      parseCommandLine();
-      if (Strings.isNullOrEmpty(commandName)) {
-        StringWriter msg = new StringWriter();
-        msg.write(usage());
-        throw new UnloggedFailure(1, msg.toString());
-      }
-
-      Command cmd = getCommand();
-      if (cmd.getClass().isAnnotationPresent(CommandMetaData.class)) {
-    	  CommandMetaData meta = cmd.getClass().getAnnotation(CommandMetaData.class);
-    	  if (meta.admin() && !ctx.getClient().getUser().canAdmin()) {
-    		  throw new UnloggedFailure(1, MessageFormat.format("{0} requires admin permissions", commandName));
-    	  }
-      }
-      if (cmd instanceof BaseCommand) {
-        BaseCommand bc = (BaseCommand) cmd;
-        if (getName().isEmpty()) {
-          bc.setName(commandName);
-        } else {
-          bc.setName(getName() + " " + commandName);
-        }
-        bc.setArguments(args.toArray(new String[args.size()]));
-      }
-
-      provideBaseStateTo(cmd);
-      provideGitState(cmd);
-      reset();
-      //atomicCmd.set(cmd);
-      cmd.start(env);
-
-    } catch (UnloggedFailure e) {
-      String msg = e.getMessage();
-      if (!msg.endsWith("\n")) {
-        msg += "\n";
-      }
-      err.write(msg.getBytes(Charsets.UTF_8));
-      err.flush();
-      exit.onExit(e.exitCode);
-    }
-  }
-
-  private Command getCommand() throws UnloggedFailure {
-	if (root != null && root.containsKey(commandName)) {
-		return root.get(commandName);
+	public DispatchCommand() {
+		commands = new HashSet<Class<? extends Command>>();
 	}
-	final Class<? extends Command> c = getMap().get(commandName);
-      if (c == null) {
-        String msg =
-            (getName().isEmpty() ? "Gitblit" : getName()) + ": "
-                + commandName + ": not found";
-        throw new UnloggedFailure(1, msg);
-      }
 
-      Command cmd = null;
-      try {
-    	  cmd = c.newInstance();
-      } catch (Exception e) {
-    	  throw new UnloggedFailure(1, MessageFormat.format("Failed to instantiate {0} command", commandName));
-      }
-	return cmd;
-  }
+	public void registerDispatcher(String name, Command cmd) {
+		if (root == null) {
+			root = Maps.newHashMap();
+		}
+		root.put(name, cmd);
+	}
 
-  @Override
-  protected String usage() {
-    final StringBuilder usage = new StringBuilder();
-    usage.append("Available commands");
-    if (!getName().isEmpty()) {
-      usage.append(" of ");
-      usage.append(getName());
-    }
-    usage.append(" are:\n");
-    usage.append("\n");
+	/**
+	 * Registers a command as long as the user is permitted to execute it.
+	 *
+	 * @param user
+	 * @param cmd
+	 */
+	public void registerCommand(UserModel user, Class<? extends Command> cmd) {
+		if (!cmd.isAnnotationPresent(CommandMetaData.class)) {
+			throw new RuntimeException(MessageFormat.format("{0} must be annotated with {1}!", cmd.getName(),
+					CommandMetaData.class.getName()));
+		}
+		CommandMetaData meta = cmd.getAnnotation(CommandMetaData.class);
+		if (meta.admin() && user.canAdmin()) {
+			log.debug(MessageFormat.format("excluding admin command {} for {}", meta.name(), user.username));
+			return;
+		}
+		commands.add(cmd);
+	}
 
-    int maxLength = -1;
-    Map<String, Class<? extends Command>> m = getMap();
-    for (String name : m.keySet()) {
-      maxLength = Math.max(maxLength, name.length());
-    }
-    String format = "%-" + maxLength + "s   %s";
-    for (String name : Sets.newTreeSet(m.keySet())) {
-      final Class<? extends Command> c = m.get(name);
-      CommandMetaData meta = c.getAnnotation(CommandMetaData.class);
-      if (meta != null) {
-        if (meta.admin() && !ctx.getClient().getUser().canAdmin()) {
-      	  continue;
-        }
-        if (meta.hidden()) {
-        	continue;
-        }
-        usage.append("   ");
-        usage.append(String.format(format, name,
-            Strings.nullToEmpty(meta.description())));
-      }
-      usage.append("\n");
-    }
-    usage.append("\n");
+	private Map<String, Class<? extends Command>> getMap() {
+		if (map == null) {
+			map = Maps.newHashMapWithExpectedSize(commands.size());
+			for (Class<? extends Command> cmd : commands) {
+				CommandMetaData meta = cmd.getAnnotation(CommandMetaData.class);
+				map.put(meta.name(), cmd);
+			}
+		}
+		return map;
+	}
 
-    usage.append("See '");
-    if (getName().indexOf(' ') < 0) {
-      usage.append(getName());
-      usage.append(' ');
-    }
-    usage.append("COMMAND --help' for more information.\n");
-    usage.append("\n");
-    return usage.toString();
-  }
+	@Override
+	public void start(Environment env) throws IOException {
+		try {
+			parseCommandLine();
+			if (Strings.isNullOrEmpty(commandName)) {
+				StringWriter msg = new StringWriter();
+				msg.write(usage());
+				throw new UnloggedFailure(1, msg.toString());
+			}
 
-  // This is needed because we are not using provider or
-  // clazz.newInstance() for DispatchCommand
-  private void reset() {
-	  args = new ArrayList<String>();
-  }
+			Command cmd = getCommand();
+			if (cmd instanceof BaseCommand) {
+				BaseCommand bc = (BaseCommand) cmd;
+				if (getName().isEmpty()) {
+					bc.setName(commandName);
+				} else {
+					bc.setName(getName() + " " + commandName);
+				}
+				bc.setArguments(args.toArray(new String[args.size()]));
+			}
 
-  private void provideGitState(Command cmd) {
-	  if (cmd instanceof BaseGitCommand) {
-		BaseGitCommand a = (BaseGitCommand) cmd;
-		a.setRepositoryResolver(repositoryResolver);
-		a.setUploadPackFactory(gitblitUploadPackFactory);
-		a.setReceivePackFactory(gitblitReceivePackFactory);
-	  } else if (cmd instanceof DispatchCommand) {
-		DispatchCommand d = (DispatchCommand)cmd;
-		d.setRepositoryResolver(repositoryResolver);
-		d.setUploadPackFactory(gitblitUploadPackFactory);
-		d.setReceivePackFactory(gitblitReceivePackFactory);
-		d.setAuthenticator(authenticator);
-	  } else if (cmd instanceof BaseKeyCommand) {
-		  BaseKeyCommand k = (BaseKeyCommand)cmd;
-		  k.setAuthenticator(authenticator);
-	  }
-  }
+			provideStateTo(cmd);
+			// atomicCmd.set(cmd);
+			cmd.start(env);
 
-  private RepositoryResolver<SshDaemonClient> repositoryResolver;
-  public void setRepositoryResolver(RepositoryResolver<SshDaemonClient> repositoryResolver) {
-	  this.repositoryResolver = repositoryResolver;
-  }
+		} catch (UnloggedFailure e) {
+			String msg = e.getMessage();
+			if (!msg.endsWith("\n")) {
+				msg += "\n";
+			}
+			err.write(msg.getBytes(Charsets.UTF_8));
+			err.flush();
+			exit.onExit(e.exitCode);
+		}
+	}
 
-  private GitblitUploadPackFactory<SshDaemonClient> gitblitUploadPackFactory;
-  public void setUploadPackFactory(GitblitUploadPackFactory<SshDaemonClient> gitblitUploadPackFactory) {
-	  this.gitblitUploadPackFactory = gitblitUploadPackFactory;
-  }
+	private Command getCommand() throws UnloggedFailure {
+		if (root != null && root.containsKey(commandName)) {
+			return root.get(commandName);
+		}
+		final Class<? extends Command> c = getMap().get(commandName);
+		if (c == null) {
+			String msg = (getName().isEmpty() ? "Gitblit" : getName()) + ": " + commandName + ": not found";
+			throw new UnloggedFailure(1, msg);
+		}
 
-  private GitblitReceivePackFactory<SshDaemonClient> gitblitReceivePackFactory;
-  public void setReceivePackFactory(GitblitReceivePackFactory<SshDaemonClient> gitblitReceivePackFactory) {
-	  this.gitblitReceivePackFactory = gitblitReceivePackFactory;
-  }
+		Command cmd = null;
+		try {
+			cmd = c.newInstance();
+		} catch (Exception e) {
+			throw new UnloggedFailure(1, MessageFormat.format("Failed to instantiate {0} command", commandName));
+		}
+		return cmd;
+	}
 
-  private PublicKeyAuthenticator authenticator;
-  public void setAuthenticator(PublicKeyAuthenticator authenticator) {
-	this.authenticator = authenticator;
-  }
+	@Override
+	protected String usage() {
+		final StringBuilder usage = new StringBuilder();
+		usage.append("Available commands");
+		if (!getName().isEmpty()) {
+			usage.append(" of ");
+			usage.append(getName());
+		}
+		usage.append(" are:\n");
+		usage.append("\n");
+
+		int maxLength = -1;
+		Map<String, Class<? extends Command>> m = getMap();
+		for (String name : m.keySet()) {
+			maxLength = Math.max(maxLength, name.length());
+		}
+		String format = "%-" + maxLength + "s   %s";
+		for (String name : Sets.newTreeSet(m.keySet())) {
+			final Class<? extends Command> c = m.get(name);
+			CommandMetaData meta = c.getAnnotation(CommandMetaData.class);
+			if (meta != null) {
+				if (meta.hidden()) {
+					continue;
+				}
+				usage.append("   ");
+				usage.append(String.format(format, name, Strings.nullToEmpty(meta.description())));
+			}
+			usage.append("\n");
+		}
+		usage.append("\n");
+
+		usage.append("See '");
+		if (getName().indexOf(' ') < 0) {
+			usage.append(getName());
+			usage.append(' ');
+		}
+		usage.append("COMMAND --help' for more information.\n");
+		usage.append("\n");
+		return usage.toString();
+	}
+
+	protected void provideStateTo(final Command cmd) {
+		if (cmd instanceof BaseCommand) {
+			((BaseCommand) cmd).setContext(ctx);
+		}
+		cmd.setInputStream(in);
+		cmd.setOutputStream(out);
+		cmd.setErrorStream(err);
+		cmd.setExitCallback(exit);
+
+		if (cmd instanceof BaseGitCommand) {
+			BaseGitCommand a = (BaseGitCommand) cmd;
+			a.setRepositoryResolver(repositoryResolver);
+			a.setUploadPackFactory(gitblitUploadPackFactory);
+			a.setReceivePackFactory(gitblitReceivePackFactory);
+		} else if (cmd instanceof DispatchCommand) {
+			DispatchCommand d = (DispatchCommand) cmd;
+			d.setRepositoryResolver(repositoryResolver);
+			d.setUploadPackFactory(gitblitUploadPackFactory);
+			d.setReceivePackFactory(gitblitReceivePackFactory);
+			d.setAuthenticator(authenticator);
+		} else if (cmd instanceof BaseKeyCommand) {
+			BaseKeyCommand k = (BaseKeyCommand) cmd;
+			k.setAuthenticator(authenticator);
+		}
+	}
+
+	private RepositoryResolver<SshDaemonClient> repositoryResolver;
+
+	public void setRepositoryResolver(RepositoryResolver<SshDaemonClient> repositoryResolver) {
+		this.repositoryResolver = repositoryResolver;
+	}
+
+	private GitblitUploadPackFactory<SshDaemonClient> gitblitUploadPackFactory;
+
+	public void setUploadPackFactory(GitblitUploadPackFactory<SshDaemonClient> gitblitUploadPackFactory) {
+		this.gitblitUploadPackFactory = gitblitUploadPackFactory;
+	}
+
+	private GitblitReceivePackFactory<SshDaemonClient> gitblitReceivePackFactory;
+
+	public void setReceivePackFactory(GitblitReceivePackFactory<SshDaemonClient> gitblitReceivePackFactory) {
+		this.gitblitReceivePackFactory = gitblitReceivePackFactory;
+	}
+
+	private PublicKeyAuthenticator authenticator;
+
+	public void setAuthenticator(PublicKeyAuthenticator authenticator) {
+		this.authenticator = authenticator;
+	}
 }

--
Gitblit v1.9.1