org.apache.sshd.common.config.keys.KeyUtils Java Examples

private SshServer setUpSshServer(String algorithm) throws IOException {
	SshServer sshServer = SshServer.setUpDefaultServer();
	sshServer.setPort(0);
	AbstractGeneratorHostKeyProvider hostKeyProvider = SecurityUtils.createGeneratorHostKeyProvider(getServerKeyFile(algorithm));
	hostKeyProvider.setAlgorithm(algorithm);
	if (algorithm.equals(KeyUtils.EC_ALGORITHM)) {
		hostKeyProvider.setKeySize(256);
	}
	sshServer.setKeyPairProvider(hostKeyProvider);

	sshServer.setPublickeyAuthenticator(AcceptAllPublickeyAuthenticator.INSTANCE);
	sshServer.setForwardingFilter(AcceptAllForwardingFilter.INSTANCE);

	writeFingerprintToKnownHosts(algorithm);

	sshServer.start();

	int sshServerPort = sshServer.getPort();
	assertTrue(sshServerPort > 0);

	return sshServer;
}
 

@Test(timeout = TEST_TIMEOUT_MILLIS)
public void testSingleHop_EcDsaServer() throws Exception {
	SshServer sshServer = setUpSshServer(KeyUtils.EC_ALGORITHM);
	int sshServerPort = sshServer.getPort();

	String hostConfigName = "localhost-" + sshServerPort;
	appendToSshFile(CONFIG_FILENAME, "Host " + hostConfigName + "\n\tHostName localhost\n\tPort " + sshServerPort + "\n\n");

	try (DummyServerSocketThread dummyServerSocketThread = new DummyServerSocketThread(TRANSFER_CHARSET, TEST_TEXT);
		SshProxy sshProxy = new SshProxy()) {
		int port = sshProxy.connect(hostConfigName, "localhost", dummyServerSocketThread.getPort());

		final String receivedText;
		try (Socket s = new Socket(SshProxy.LOCALHOST, port);
			 InputStream is = s.getInputStream()) {
			log.info("connected to port: {}", port);
			receivedText = readLine(is);
		}
		assertEquals(TEST_TEXT, receivedText);
	} finally {
		tryStop(sshServer);
	}
}
 

public static PrivateKey decodePEMPrivateKey(String privateKey) throws IOException, GeneralSecurityException {
    try (PemReader pemReaderPrivate = new PemReader(new StringReader(privateKey))) {
        KeyFactory kf = SecurityUtils.getKeyFactory(KeyUtils.RSA_ALGORITHM);
        
        PemObject pemObjectPrivate = pemReaderPrivate.readPemObject();
        PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(pemObjectPrivate.getContent());
        return kf.generatePrivate(spec);
    }
}
 

public static KeyPair generateKeyPair(String algorithm, int keySize) throws GeneralSecurityException {
    KeyPairGenerator gen = SecurityUtils.getKeyPairGenerator(algorithm);
    if (KeyUtils.EC_ALGORITHM.equalsIgnoreCase(algorithm)) {
        ECCurves curve = ECCurves.fromCurveSize(keySize);
        if (curve == null) {
            throw new InvalidKeySpecException("Unknown curve for key size=" + keySize);
        }
        gen.initialize(curve.getParameters());
    } else {
        gen.initialize(keySize);
    }

    return gen.generateKeyPair();
}
 

private static Path getServerKeyFile(String algorithm) {
	switch (algorithm) {
		case KeyUtils.RSA_ALGORITHM:
			return SERVER_RSA_KEY;
		case KeyUtils.EC_ALGORITHM:
			return SERVER_ECDSA_KEY;
		default:
			throw new IllegalArgumentException("Unknown algorithm: " + algorithm);
	}
}
 

private void writeFingerprintToKnownHosts(String algorithm) throws IOException {
	switch (algorithm) {
		case KeyUtils.RSA_ALGORITHM:
			appendToSshFile(KNOWN_HOSTS_FILENAME, "localhost ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQDL8360Wxcgo33sggS0bSid0u7Ad4XFig8/e0UfD5l02x/w2DRJuqJow4SiDfi9jvD8p3lu7To7b/oGH/c/vsK9j35ICG0eJ/bbnQDuHROBAnbAC6PXN+/XX2F9s48KlOC5dQXrGYyYhoozW67yoHTooisZSzF/iyPdNat64rM0+ZO3dV6eEQ0FItYO632YcSiBRE7YZe9rP7ne50xaltKgrAmHRDRo+tjIcykrlcZFG1Bp/ct9Ejs2DQDsFOZRCmFbag0pQxxbkA1U6z7O3qwhhDWcJz2ZHDHK8DUkgHdX+Hbp7LxBWEaCiU8cL+S6rmCpNsui9NT/XeoLuXQ4J8jX\n");
			break;
		case KeyUtils.EC_ALGORITHM:
			appendToSshFile(KNOWN_HOSTS_FILENAME, "localhost ecdsa-sha2-nistp256 AAAAE2VjZHNhLXNoYTItbmlzdHAyNTYAAAAIbmlzdHAyNTYAAABBBCH+0xjLYNGoqVGlD4VtKHF1Tig2/Y76BxVld88bYAaRV4ojJni62vIYMKqk+FMZhL1lcQ/VQTvIeLMnYk+grKo=\n");
			break;
		default:
			throw new IllegalArgumentException("Unknown algorithm: " + algorithm);
	}
}
 

@Test   // see SSHD-620
public void testHostBasedAuthentication() throws Exception {
    final String hostClienUser = getClass().getSimpleName();
    final String hostClientName = SshdSocketAddress.toAddressString(SshdSocketAddress.getFirstExternalNetwork4Address());
    final KeyPair hostClientKey = Utils.generateKeyPair(KeyUtils.RSA_ALGORITHM, 1024);
    final AtomicInteger invocationCount = new AtomicInteger(0);
    sshd.setHostBasedAuthenticator(new HostBasedAuthenticator() {
        @Override
        public boolean authenticate(ServerSession session, String username,
                PublicKey clientHostKey, String clientHostName, String clientUsername, List<X509Certificate> certificates) {
            invocationCount.incrementAndGet();
            return hostClienUser.equals(clientUsername)
                && hostClientName.equals(clientHostName)
                && KeyUtils.compareKeys(hostClientKey.getPublic(), clientHostKey);
        }
    });
    sshd.setPasswordAuthenticator(RejectAllPasswordAuthenticator.INSTANCE);
    sshd.setKeyboardInteractiveAuthenticator(KeyboardInteractiveAuthenticator.NONE);
    sshd.setPublickeyAuthenticator(RejectAllPublickeyAuthenticator.INSTANCE);
    sshd.setUserAuthFactories(
            Collections.<NamedFactory<org.apache.sshd.server.auth.UserAuth>>singletonList(
                    org.apache.sshd.server.auth.hostbased.UserAuthHostBasedFactory.INSTANCE));

    try (SshClient client = setupTestClient()) {
        org.apache.sshd.client.auth.hostbased.UserAuthHostBasedFactory factory =
                new org.apache.sshd.client.auth.hostbased.UserAuthHostBasedFactory();
        // TODO factory.setClientHostname(CLIENT_HOSTNAME);
        factory.setClientUsername(hostClienUser);
        factory.setClientHostKeys(HostKeyIdentityProvider.Utils.wrap(hostClientKey));

        client.setUserAuthFactories(Collections.<NamedFactory<org.apache.sshd.client.auth.UserAuth>>singletonList(factory));
        client.start();
        try (ClientSession s = client.connect(getCurrentTestName(), TEST_LOCALHOST, port).verify(7L, TimeUnit.SECONDS).getSession()) {
            s.auth().verify(11L, TimeUnit.SECONDS);
            assertEquals("Mismatched authenticator invocation count", 1, invocationCount.get());
        } finally {
            client.stop();
        }
    }
}
 

@Sessional
@Override
public boolean authenticate(String username, PublicKey key, ServerSession session) throws AsyncAuthException {
       String digest = KeyUtils.getFingerPrint(SshKey.DIGEST_FORMAT, key);  
       SshKey sshKey = sshKeyManager.findByDigest(digest);
       if (sshKey != null) {
           session.setAttribute(ATTR_PUBLIC_KEY_OWNER_ID, sshKey.getOwner().getId());
           return true;
       } else {
       	return false;
       }
}
 

@Override
public Iterable<KeyPair> loadKeys(SessionContext session) {
    SshSetting sshSetting = settingManager.getSshSetting();
    
    try {
        PrivateKey privateKey = SshKeyUtils.decodePEMPrivateKey(sshSetting.getPemPrivateKey());
        PublicKey publicKey = KeyUtils.recoverPublicKey(privateKey);
        return Lists.newArrayList(new KeyPair(publicKey, privateKey));
    } catch (Exception e) {
        throw new RuntimeException(e);
    }
}
 

public static String generatePEMPrivateKey() {
    try (StringWriter privateWriter = new StringWriter();
            PemWriter privatePemWriter = new PemWriter(privateWriter)) {
       KeyPair keyPair = KeyUtils.generateKeyPair("ssh-rsa", 4096);
       
       privatePemWriter.writeObject(new PemObject("RSA PRIVATE KEY", keyPair.getPrivate().getEncoded()));
       privatePemWriter.flush();
       
       return privateWriter.toString();
   } catch (GeneralSecurityException | IOException e) {
       throw new RuntimeException(e);
   }
}
 

public static PublicKey decodePEMPublicKey(String publicKey) throws IOException, GeneralSecurityException {
    try (PemReader pemReaderPublic = new PemReader(new StringReader(publicKey))) {
        KeyFactory kf = SecurityUtils.getKeyFactory(KeyUtils.RSA_ALGORITHM);
        
        PemObject pemObjectPublic = pemReaderPublic.readPemObject();
        X509EncodedKeySpec x509EncodedKeySpec = new X509EncodedKeySpec(pemObjectPublic.getContent());
        return kf.generatePublic(x509EncodedKeySpec);
    }
}
 

public PublicKey getPublicKey() {
try {
	return KeyUtils.recoverPublicKey(getPrivateKey());
} catch (GeneralSecurityException e) {
	throw new RuntimeException(e);
}
  }
 

public String getFingerPrint() {
      try {
	PrivateKey privateKey = SshKeyUtils.decodePEMPrivateKey(pemPrivateKey);
	PublicKey publicKey = KeyUtils.recoverPublicKey(privateKey);
	return KeyUtils.getFingerPrint(BuiltinDigests.sha256, publicKey);
} catch (IOException | GeneralSecurityException e) {
	throw new RuntimeException(e);
}
  }
 

@Transactional
  @Override
  public void syncSshKeys(User user, Collection<String> sshKeys) {
  	Map<String, SshKey> syncMap = new HashMap<>();
  	for (String content: sshKeys) {
  		try {
  			PublicKey pubEntry = SshKeyUtils.decodeSshPublicKey(content);
  	        String digest = KeyUtils.getFingerPrint(SshKey.DIGEST_FORMAT, pubEntry);
  			
  	        SshKey sshKey = new SshKey();
  	        sshKey.setDigest(digest);
  	        sshKey.setContent(content);
  	        sshKey.setOwner(user);
  	        sshKey.setDate(new Date());
  	        syncMap.put(content, sshKey);
  		} catch (IOException | GeneralSecurityException e) {
  			logger.error("Error parsing SSH key", e);
  		}
  	}

  	Map<String, SshKey> currentMap = new HashMap<>();
user.getSshKeys().forEach(sshKey -> currentMap.put(sshKey.getContent(), sshKey));

MapDifference<String, SshKey> diff = Maps.difference(currentMap, syncMap);

diff.entriesOnlyOnLeft().values().forEach(sshKey -> delete(sshKey));

diff.entriesOnlyOnRight().values().forEach(sshKey -> {
	if (findByDigest(sshKey.getDigest()) == null) 
		save(sshKey);	
	else 
		logger.warn("SSH key is already in use (digest: {})", sshKey.getDigest());
});

  }
 

@Test
public void testPublicKeyAuthWithCache() throws Exception {
    final ConcurrentHashMap<String, AtomicInteger> count = new ConcurrentHashMap<String, AtomicInteger>();
    TestCachingPublicKeyAuthenticator auth = new TestCachingPublicKeyAuthenticator(new PublickeyAuthenticator() {
        @SuppressWarnings("synthetic-access")
        @Override
        public boolean authenticate(String username, PublicKey key, ServerSession session) {
            String fp = KeyUtils.getFingerPrint(key);
            count.putIfAbsent(fp, new AtomicInteger());
            count.get(fp).incrementAndGet();
            return key.equals(pairRsa.getPublic());
        }
    });
    delegate = auth;

    try (SshClient client = setupTestClient()) {
        client.start();

        try (ClientSession session = client.connect(getCurrentTestName(), TEST_LOCALHOST, port).verify(7L, TimeUnit.SECONDS).getSession()) {
            session.addPublicKeyIdentity(pairRsaBad);
            session.addPublicKeyIdentity(pairRsa);
            session.auth().verify(5L, TimeUnit.SECONDS);

            assertEquals("Mismatched authentication invocations count", 2, count.size());

            String fpBad = KeyUtils.getFingerPrint(pairRsaBad.getPublic());
            String fpGood = KeyUtils.getFingerPrint(pairRsa.getPublic());
            assertTrue("Missing bad public key", count.containsKey(fpBad));
            assertTrue("Missing good public key", count.containsKey(fpGood));
            assertEquals("Mismatched bad key authentication attempts", 1, count.get(fpBad).get());
            assertEquals("Mismatched good key authentication attempts", 1, count.get(fpGood).get());
        } finally {
            client.stop();
        }
    }

    Thread.sleep(100L);
    assertTrue("Cache not empty", auth.getCache().isEmpty());
}
 

@Test
public void testPublicKeyAuthWithCache() throws Exception {
    final ConcurrentHashMap<String, AtomicInteger> count = new ConcurrentHashMap<String, AtomicInteger>();
    TestCachingPublicKeyAuthenticator auth = new TestCachingPublicKeyAuthenticator(new PublickeyAuthenticator() {
        @SuppressWarnings("synthetic-access")
        @Override
        public boolean authenticate(String username, PublicKey key, ServerSession session) {
            String fp = KeyUtils.getFingerPrint(key);
            count.putIfAbsent(fp, new AtomicInteger());
            count.get(fp).incrementAndGet();
            return key.equals(pairRsa.getPublic());
        }
    });
    delegate = auth;

    try (SshClient client = setupTestClient()) {
        client.start();

        try (ClientSession session = client.connect(getCurrentTestName(), TEST_LOCALHOST, port).verify(7L, TimeUnit.SECONDS).getSession()) {
            session.addPublicKeyIdentity(pairRsaBad);
            session.addPublicKeyIdentity(pairRsa);
            session.auth().verify(5L, TimeUnit.SECONDS);

            assertEquals("Mismatched authentication invocations count", 2, count.size());

            String fpBad = KeyUtils.getFingerPrint(pairRsaBad.getPublic());
            String fpGood = KeyUtils.getFingerPrint(pairRsa.getPublic());
            assertTrue("Missing bad public key", count.containsKey(fpBad));
            assertTrue("Missing good public key", count.containsKey(fpGood));
            assertEquals("Mismatched bad key authentication attempts", 1, count.get(fpBad).get());
            assertEquals("Mismatched good key authentication attempts", 1, count.get(fpGood).get());
        } finally {
            client.stop();
        }
    }

    Thread.sleep(100L);
    assertTrue("Cache not empty", auth.getCache().isEmpty());
}
 

@Test   // see SSHD-620
public void testHostBasedAuthentication() throws Exception {
    final String hostClienUser = getClass().getSimpleName();
    final String hostClientName = SshdSocketAddress.toAddressString(SshdSocketAddress.getFirstExternalNetwork4Address());
    final KeyPair hostClientKey = Utils.generateKeyPair(KeyUtils.RSA_ALGORITHM, 1024);
    final AtomicInteger invocationCount = new AtomicInteger(0);
    sshd.setHostBasedAuthenticator(new HostBasedAuthenticator() {
        @Override
        public boolean authenticate(ServerSession session, String username,
                PublicKey clientHostKey, String clientHostName, String clientUsername, List<X509Certificate> certificates) {
            invocationCount.incrementAndGet();
            return hostClienUser.equals(clientUsername)
                && hostClientName.equals(clientHostName)
                && KeyUtils.compareKeys(hostClientKey.getPublic(), clientHostKey);
        }
    });
    sshd.setPasswordAuthenticator(RejectAllPasswordAuthenticator.INSTANCE);
    sshd.setKeyboardInteractiveAuthenticator(KeyboardInteractiveAuthenticator.NONE);
    sshd.setPublickeyAuthenticator(RejectAllPublickeyAuthenticator.INSTANCE);
    sshd.setUserAuthFactories(
            Collections.<NamedFactory<org.apache.sshd.server.auth.UserAuth>>singletonList(
                    org.apache.sshd.server.auth.hostbased.UserAuthHostBasedFactory.INSTANCE));

    try (SshClient client = setupTestClient()) {
        org.apache.sshd.client.auth.hostbased.UserAuthHostBasedFactory factory =
                new org.apache.sshd.client.auth.hostbased.UserAuthHostBasedFactory();
        // TODO factory.setClientHostname(CLIENT_HOSTNAME);
        factory.setClientUsername(hostClienUser);
        factory.setClientHostKeys(HostKeyIdentityProvider.Utils.wrap(hostClientKey));

        client.setUserAuthFactories(Collections.<NamedFactory<org.apache.sshd.client.auth.UserAuth>>singletonList(factory));
        client.start();
        try (ClientSession s = client.connect(getCurrentTestName(), TEST_LOCALHOST, port).verify(7L, TimeUnit.SECONDS).getSession()) {
            s.auth().verify(11L, TimeUnit.SECONDS);
            assertEquals("Mismatched authenticator invocation count", 1, invocationCount.get());
        } finally {
            client.stop();
        }
    }
}
 

public static KeyPair generateKeyPair(String algorithm, int keySize) throws GeneralSecurityException {
    KeyPairGenerator gen = SecurityUtils.getKeyPairGenerator(algorithm);
    if (KeyUtils.EC_ALGORITHM.equalsIgnoreCase(algorithm)) {
        ECCurves curve = ECCurves.fromCurveSize(keySize);
        if (curve == null) {
            throw new InvalidKeySpecException("Unknown curve for key size=" + keySize);
        }
        gen.initialize(curve.getParameters());
    } else {
        gen.initialize(keySize);
    }

    return gen.generateKeyPair();
}
 

@Test
public void testPublicKeyAuthWithoutCache() throws Exception {
    final ConcurrentHashMap<String, AtomicInteger> count = new ConcurrentHashMap<String, AtomicInteger>();
    delegate = new PublickeyAuthenticator() {
        @SuppressWarnings("synthetic-access")
        @Override
        public boolean authenticate(String username, PublicKey key, ServerSession session) {
            String fp = KeyUtils.getFingerPrint(key);
            count.putIfAbsent(fp, new AtomicInteger());
            count.get(fp).incrementAndGet();
            return key.equals(pairRsa.getPublic());
        }
    };

    try (SshClient client = setupTestClient()) {
        client.start();

        try (ClientSession session = client.connect(getCurrentTestName(), TEST_LOCALHOST, port).verify(7L, TimeUnit.SECONDS).getSession()) {
            session.addPublicKeyIdentity(pairRsaBad);
            session.addPublicKeyIdentity(pairRsa);

            AuthFuture auth = session.auth();
            assertTrue("Failed to authenticate on time", auth.await(5L, TimeUnit.SECONDS));
            assertTrue("Authentication failed", auth.isSuccess());
        } finally {
            client.stop();
        }
    }

    assertEquals("Mismatched attempted keys count", 2, count.size());

    String badFingerPrint = KeyUtils.getFingerPrint(pairRsaBad.getPublic());
    Number badIndex = count.get(badFingerPrint);
    assertNotNull("Missing bad RSA key", badIndex);
    assertEquals("Mismatched attempt index for bad key", 1, badIndex.intValue());

    String goodFingerPrint = KeyUtils.getFingerPrint(pairRsa.getPublic());
    Number goodIndex = count.get(goodFingerPrint);
    assertNotNull("Missing good RSA key", goodIndex);
    assertEquals("Mismatched attempt index for good key", 2, goodIndex.intValue());
}
 

public SinglePublicKeyAuthTest() {
    SimpleGeneratorHostKeyProvider provider = new SimpleGeneratorHostKeyProvider();
    provider.setAlgorithm(KeyUtils.RSA_ALGORITHM);
    pairRsaBad = provider.loadKey(KeyPairProvider.SSH_RSA);
}
 

@Test   // see SSHD-624
public void testMismatchedUserAuthPkOkData() throws Exception {
    final AtomicInteger challengeCounter = new AtomicInteger(0);
    sshd.setUserAuthFactories(Collections.<NamedFactory<org.apache.sshd.server.auth.UserAuth>>singletonList(
            new org.apache.sshd.server.auth.pubkey.UserAuthPublicKeyFactory() {
                @Override
                public org.apache.sshd.server.auth.pubkey.UserAuthPublicKey create() {
                    return new org.apache.sshd.server.auth.pubkey.UserAuthPublicKey() {
                        @Override
                        protected void sendPublicKeyResponse(ServerSession session, String username, String alg, PublicKey key,
                                byte[] keyBlob, int offset, int blobLen, Buffer buffer) throws Exception {
                            int count = challengeCounter.incrementAndGet();
                            outputDebugMessage("sendPublicKeyChallenge(%s)[%s]: count=%d", session, alg, count);
                            if (count == 1) {
                                // send wrong key type
                                super.sendPublicKeyResponse(session, username, KeyPairProvider.SSH_DSS, key, keyBlob, offset, blobLen, buffer);
                            } else if (count == 2) {
                                // send another key
                                KeyPair otherPair = org.apache.sshd.util.test.Utils.generateKeyPair(KeyUtils.RSA_ALGORITHM, 1024);
                                PublicKey otherKey = otherPair.getPublic();
                                Buffer buf = session.createBuffer(SshConstants.SSH_MSG_USERAUTH_PK_OK, blobLen + alg.length() + Long.SIZE);
                                buf.putString(alg);
                                buf.putPublicKey(otherKey);
                                session.writePacket(buf);
                            } else {
                                super.sendPublicKeyResponse(session, username, alg, key, keyBlob, offset, blobLen, buffer);
                            }
                        }
                    };
                }

    }));

    try (SshClient client = setupTestClient()) {
        KeyPair clientIdentity = Utils.generateKeyPair(KeyUtils.RSA_ALGORITHM, 1024);
        client.start();

        try {
            for (int index = 1; index <= 4; index++) {
                try (ClientSession s = client.connect(getCurrentTestName(), TEST_LOCALHOST, port).verify(7L, TimeUnit.SECONDS).getSession()) {
                    s.addPublicKeyIdentity(clientIdentity);
                    s.auth().verify(17L, TimeUnit.SECONDS);
                    assertEquals("Mismatched number of challenges", 3, challengeCounter.get());
                    break;
                } catch (SshException e) {   // expected
                    outputDebugMessage("%s on retry #%d: %s", e.getClass().getSimpleName(), index, e.getMessage());

                    Throwable t = e.getCause();
                    assertObjectInstanceOf("Unexpected failure cause at retry #" + index, InvalidKeySpecException.class, t);
                }
            }
        } finally {
            client.stop();
        }
    }
}
 

private SshServer setUpSshServer() throws IOException {
	return setUpSshServer(KeyUtils.RSA_ALGORITHM);
}
 

@Test
public void testPublicKeyAuthWithoutCache() throws Exception {
    final ConcurrentHashMap<String, AtomicInteger> count = new ConcurrentHashMap<String, AtomicInteger>();
    delegate = new PublickeyAuthenticator() {
        @SuppressWarnings("synthetic-access")
        @Override
        public boolean authenticate(String username, PublicKey key, ServerSession session) {
            String fp = KeyUtils.getFingerPrint(key);
            count.putIfAbsent(fp, new AtomicInteger());
            count.get(fp).incrementAndGet();
            return key.equals(pairRsa.getPublic());
        }
    };

    try (SshClient client = setupTestClient()) {
        client.start();

        try (ClientSession session = client.connect(getCurrentTestName(), TEST_LOCALHOST, port).verify(7L, TimeUnit.SECONDS).getSession()) {
            session.addPublicKeyIdentity(pairRsaBad);
            session.addPublicKeyIdentity(pairRsa);

            AuthFuture auth = session.auth();
            assertTrue("Failed to authenticate on time", auth.await(5L, TimeUnit.SECONDS));
            assertTrue("Authentication failed", auth.isSuccess());
        } finally {
            client.stop();
        }
    }

    assertEquals("Mismatched attempted keys count", 2, count.size());

    String badFingerPrint = KeyUtils.getFingerPrint(pairRsaBad.getPublic());
    Number badIndex = count.get(badFingerPrint);
    assertNotNull("Missing bad RSA key", badIndex);
    assertEquals("Mismatched attempt index for bad key", 1, badIndex.intValue());

    String goodFingerPrint = KeyUtils.getFingerPrint(pairRsa.getPublic());
    Number goodIndex = count.get(goodFingerPrint);
    assertNotNull("Missing good RSA key", goodIndex);
    assertEquals("Mismatched attempt index for good key", 2, goodIndex.intValue());
}
 

public SinglePublicKeyAuthTest() {
    SimpleGeneratorHostKeyProvider provider = new SimpleGeneratorHostKeyProvider();
    provider.setAlgorithm(KeyUtils.RSA_ALGORITHM);
    pairRsaBad = provider.loadKey(KeyPairProvider.SSH_RSA);
}
 

@Test   // see SSHD-624
public void testMismatchedUserAuthPkOkData() throws Exception {
    final AtomicInteger challengeCounter = new AtomicInteger(0);
    sshd.setUserAuthFactories(Collections.<NamedFactory<org.apache.sshd.server.auth.UserAuth>>singletonList(
            new org.apache.sshd.server.auth.pubkey.UserAuthPublicKeyFactory() {
                @Override
                public org.apache.sshd.server.auth.pubkey.UserAuthPublicKey create() {
                    return new org.apache.sshd.server.auth.pubkey.UserAuthPublicKey() {
                        @Override
                        protected void sendPublicKeyResponse(ServerSession session, String username, String alg, PublicKey key,
                                byte[] keyBlob, int offset, int blobLen, Buffer buffer) throws Exception {
                            int count = challengeCounter.incrementAndGet();
                            outputDebugMessage("sendPublicKeyChallenge(%s)[%s]: count=%d", session, alg, count);
                            if (count == 1) {
                                // send wrong key type
                                super.sendPublicKeyResponse(session, username, KeyPairProvider.SSH_DSS, key, keyBlob, offset, blobLen, buffer);
                            } else if (count == 2) {
                                // send another key
                                KeyPair otherPair = org.apache.sshd.util.test.Utils.generateKeyPair(KeyUtils.RSA_ALGORITHM, 1024);
                                PublicKey otherKey = otherPair.getPublic();
                                Buffer buf = session.createBuffer(SshConstants.SSH_MSG_USERAUTH_PK_OK, blobLen + alg.length() + Long.SIZE);
                                buf.putString(alg);
                                buf.putPublicKey(otherKey);
                                session.writePacket(buf);
                            } else {
                                super.sendPublicKeyResponse(session, username, alg, key, keyBlob, offset, blobLen, buffer);
                            }
                        }
                    };
                }

    }));

    try (SshClient client = setupTestClient()) {
        KeyPair clientIdentity = Utils.generateKeyPair(KeyUtils.RSA_ALGORITHM, 1024);
        client.start();

        try {
            for (int index = 1; index <= 4; index++) {
                try (ClientSession s = client.connect(getCurrentTestName(), TEST_LOCALHOST, port).verify(7L, TimeUnit.SECONDS).getSession()) {
                    s.addPublicKeyIdentity(clientIdentity);
                    s.auth().verify(17L, TimeUnit.SECONDS);
                    assertEquals("Mismatched number of challenges", 3, challengeCounter.get());
                    break;
                } catch (SshException e) {   // expected
                    outputDebugMessage("%s on retry #%d: %s", e.getClass().getSimpleName(), index, e.getMessage());

                    Throwable t = e.getCause();
                    assertObjectInstanceOf("Unexpected failure cause at retry #" + index, InvalidKeySpecException.class, t);
                }
            }
        } finally {
            client.stop();
        }
    }
}