Correct code generation and apply timeout to login sequence.

Author: MikeNeilson

cwbi-auth-http-client/src/main/java/hec/army/usace/hec/cwbi/auth/http/client/AuthCodePkceTokenRequestBuilder.java

@@ -37,6 +37,9 @@
 import java.security.SecureRandom;
 import java.util.Base64;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicReference;
 import java.awt.Desktop;
 import java.awt.Desktop.Action;
@@ -61,31 +64,38 @@ OAuth2Token retrieveToken() throws IOException {
         try {
             byte[] verifierBytes = new byte[128];
             SecureRandom.getInstanceStrong().nextBytes(verifierBytes);
-            Base64.Encoder b64encoder = Base64.getUrlEncoder();
+            Base64.Encoder b64encoder = Base64.getUrlEncoder().withoutPadding();
             final String verifier = b64encoder.encodeToString(verifierBytes);
 
             MessageDigest md = MessageDigest.getInstance("SHA-256");
             final String challenge = b64encoder.encodeToString(md.digest(verifierBytes));
             HttpServer server = HttpServer.create(new InetSocketAddress("localhost", 0), 0);            
             int port = server.getAddress().getPort();
             String host = server.getAddress().getHostName();
-            final AtomicReference<String> code = new AtomicReference<>();
-            final AtomicReference<String> state = new AtomicReference<>();
-            final CompletableFuture<Void> future = new CompletableFuture<>();
+
+            final CompletableFuture<Result> future = new CompletableFuture<>();
 
             server.createContext("/", new HttpHandler() {
 
                 @Override
                 public void handle(HttpExchange exchange) throws IOException {
+                    Result ret = null;
+                    
                     final String query = exchange.getRequestURI().getQuery();
+                    System.out.println("Got auth server response." + query);
                     final QueryParameters parameters = QueryParameters.parse(query);
-                    
-                    code.set(parameters.get("code").get(0));
-                    state.set(parameters.get("state").get(0));
-                    exchange.sendResponseHeaders(201, 0);
-                    System.out.println("Got code");
+                    if (!parameters.get("error").isEmpty()) {
+                        String error = parameters.get("error").get(0);
+                        String errorDescription = parameters.get("error_description").get(0);
+                        ret = Result.failure(error, errorDescription);                    
+                    } else {
+                        String code = parameters.get("code").get(0);
+                        String state = parameters.get("state").get(0);
+                        ret = Result.success(code ,state);
+                    }
+                    System.out.println("Finishing");
                     server.stop(0);
-                    future.complete(null);
+                    future.complete(ret);
                 }
 
             });
@@ -94,26 +104,30 @@ public void handle(HttpExchange exchange) throws IOException {
                 .set("grant_type", "code")
                 .set("client_id", getClientId())
                 .set("scopes", "openid profile")
+                .set("response_type", "code")
                 .set("code_challenge_method", "S256")
                 .set("code_challenge", challenge)
-                .set("redirect_uri", redirectUri);
+                .set("redirect_uri", redirectUri)
+                .set("kc_idp_hint", "login.gov");
             String urlStr= String.format("%s?%s", getAuthUrl().getApiRoot(), authParameters.encode());
             // start server to listen
             server.start();
+            System.out.println(urlStr);
             this.authCallBack.accept(URI.create(urlStr));
 
-
-            future.join();
+            Result result = future.get(1, TimeUnit.MINUTES); // The user is now required to perform manual operations.
             System.out.println("Next steps.");
-
+            if (result.error != null) {
+                throw new IOException(String.format("Unable to login. %s : %s", result.error, result.errorDescription));
+            }
             final UrlEncodedFormData formData = new UrlEncodedFormData();
             formData.addClientId(getClientId())
                     .addGrantType("authorization_code")
                     .addParameter("code_verifier", verifier)
                     .addScopes("openid", "profile")
                     .addParameter("redirect_uri", redirectUri)
-                    .addParameter("session_state", state.get())
-                    .addParameter("code", code.get())
+                    .addParameter("session_state", result.state)
+                    .addParameter("code", result.code)
                     .addParameter("response_mode", "fragment")
                     .addParameter("response_type", "id_token token");
 
@@ -131,6 +145,31 @@ public void handle(HttpExchange exchange) throws IOException {
             return retVal;
         } catch (NoSuchAlgorithmException ex) {
             throw new IOException("Unable to retrieve SecureRandom or Message Digest instance to generate verifier", ex);
+        } catch (InterruptedException | ExecutionException | TimeoutException ex) {
+            throw new IOException("Unable to form login sequence.", ex);
         }
     }
+
+    private static class Result {
+        public final String code;
+        public final String state;
+
+        public final String error;
+        public final String errorDescription;
+
+        private Result(String code, String state, String error, String errorDescription) {
+            this.code = code;
+            this.state = state;
+            this.error = error;
+            this.errorDescription = errorDescription;
+        }
+
+        public static Result success(String code, String state) {
+            return new Result(code, state, null, null);
+        }
+
+        public static Result failure(String error, String errorDescription) {
+            return new Result(null, null, error, errorDescription);
+        }
+    };
 }