Experimenting with URLConnection (keepalive, SSL)

JavaURLConnection.java

import static org.junit.Assert.assertFalse;

import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.ServerSocket;
import java.net.Socket;
import java.net.SocketException;
import java.net.URL;
import java.net.URLConnection;
import java.security.InvalidKeyException;
import java.security.KeyFactory;
import java.security.KeyManagementException;
import java.security.KeyStore;
import java.security.KeyStoreException;
import java.security.NoSuchAlgorithmException;
import java.security.PrivateKey;
import java.security.SecureRandom;
import java.security.UnrecoverableKeyException;
import java.security.cert.Certificate;
import java.security.cert.CertificateException;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;
import java.security.spec.InvalidKeySpecException;
import java.security.spec.PKCS8EncodedKeySpec;
import java.util.concurrent.Semaphore;

import javax.net.ssl.HttpsURLConnection;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.SSLHandshakeException;
import javax.net.ssl.SSLServerSocketFactory;
import javax.net.ssl.SSLSocketFactory;
import javax.net.ssl.TrustManager;
import javax.net.ssl.TrustManagerFactory;
import javax.net.ssl.X509TrustManager;

import org.apache.commons.io.IOUtils;
import org.apache.commons.io.input.AutoCloseInputStream;
import org.apache.commons.lang.CharEncoding;
import org.junit.Test;

import com.sun.xml.internal.messaging.saaj.util.ByteOutputStream;

/**
 * 
 * How all this cruft (keystores, keys, certs) was generated:
 * <pre>
 * # server certificate keystore to check when connecting as client
 * openssl genrsa -out server.private.key 1024
 * openssl req -out server.csr -key server.private.key -new
 * openssl x509 -req -in server.csr -out server.crt -signkey server.private.key
 * keytool -importcert -v -trustcacerts -alias server_cert -file server.crt -keystore client.jks -storepass changeit
 * 
 * # keystore with private key for debug server
 * # java can read only pkcs8 private keys on default
 * # so we convert PEM encoded RSA private key to PKCS8
 * # and import it with importPEMPrivate() method
 * openssl pkcs8 -in server.private.key -out server.private.pkcs8 -topk8 -nocryp
t -outform der
 * 
 * # replacement server certificate to check when connecting as client
 * openssl genrsa -out repl.private.key 1024
 * openssl req -out repl.csr -key repl.private.key -new
 * openssl x509 -req -in repl.csr -out repl.crt -signkey repl.private.key
 * keytool -importcert -v -trustcacerts -alias repl -file repl.crt -keystore rep
l.jks -storepass changeit
 * </pre>
 * 
 * Actually, we could generate all this right here, but it's not that fun
 *
 */
public class JavaURLConnection {
	
	Semaphore s = new Semaphore(0);
	
	@Test
	public void importPEMPrivate() throws KeyStoreException, NoSuchAlgorithmException, CertificateException, IOException, InvalidKeyException, InvalidKeySpecException {
		KeyStore ks = KeyStore.getInstance(KeyStore.getDefaultType());
		ks.load(null, null);
		
		InputStream pkIs = new AutoCloseInputStream(new FileInputStream(new File("server.private.pkcs8")));
		ByteOutputStream baos = new ByteOutputStream();
		IOUtils.copy(pkIs, baos);
		
		PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(baos.getBytes());
		PrivateKey key = KeyFactory.getInstance("RSA").generatePrivate(spec);
		
		CertificateFactory certFac = CertificateFactory.getInstance("X.509");
		Certificate cert = certFac.generateCertificate(new FileInputStream(new File("server.crt")));
		
		ks.setKeyEntry("server_key", key, "changeit".toCharArray(), new Certificate[] { cert });
		
		
		FileOutputStream fos = new FileOutputStream(new File("server.priv.jks"));
		ks.store(fos, "changeit".toCharArray());
		fos.close();
	}
	
	public void runDebugServer(boolean isSsl, boolean ka) throws NoSuchAlgorithmException, IOException, KeyManagementException, KeyStoreException, CertificateException, UnrecoverableKeyException {
		ServerSocket ss = null;
		if (isSsl) {
			SSLContext context = SSLContext.getInstance("TLS");
			KeyManagerFactory kmf = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
			KeyStore ks = KeyStore.getInstance(KeyStore.getDefaultType());
			ks.load(new FileInputStream(new File("server.priv.jks")), "changeit".toCharArray());
			kmf.init(ks, "changeit".toCharArray());
			context.init(kmf.getKeyManagers(), null, new SecureRandom());
			SSLServerSocketFactory ssf = context.getServerSocketFactory();
			ss = ssf.createServerSocket(10000);
		} else {
			ss = new ServerSocket(10000);
		}
		
		
		s.release();
		
		int conn = 1;
		
		while(true) {
			Socket cs = ss.accept();
			System.out.println("----- Connect #" + conn + "\n\n");
			conn++;
			
			try {
				InputStream is = cs.getInputStream();
				OutputStream os = cs.getOutputStream();
	
				boolean closeConn = false;
				
				int state = 0;
				// 0 - in content
				// 1 - \r occured
				// 2 - \n occured
				// 3 - \r second time
				// 4 - \n second time
				while (!closeConn) {
				
					byte [] buff = new byte[100];
					// yack, ugly, why am I reading byte buffer and then examine each byte - who knows
					// better to use BufferedIS and use simple byte-by-byte read()
					int read = is.read(buff);
					while (read > 0) {
						String a = new String(buff, 0, read, CharEncoding.US_ASCII);
						System.out.println("PACKET: " + a.replaceAll("\n", "N").replace("\r", "R"));
						for (int i = 0; i < read; i++) {
							byte b = buff[i];
							
							String s = new String(new byte[] { b }, "ascii");
							if (b == '\n') {
								s = "N";
							} else if (b == '\r') {
								s = "R";
							}
							System.out.print(s);
							
							switch (b) {
							case '\n':
								if (state == 1) {
									state = 2;
								} else if (state == 3) {
									state = 4;
								}
								break;
							case '\r':
								if (state == 0) {
									state = 1;
								} else if (state == 2) {
									state = 3;
								}
								break;
							default:
								state = 0;
							}
						}
						if (state == 4) {
							break;
						}
						System.out.println();
						read = is.read(buff);
					}
					
					
					String headers = "HTTP/1.1 200 OK\r\n" + 
						"Content-Length: %d\r\n" + 
						"Keep-Alive: timeout=60, max=99\r\n" +
						"Connection: Keep-Alive\r\n" +
						"Content-Type: text/html\r\n\r\n";
					String body = "<html><body><h1>It works!</h1></body></html>";
					byte[] bytes = body.getBytes(CharEncoding.UTF_8);
					
					os.write(String.format(headers, bytes.length).getBytes(CharEncoding.UTF_8));
					os.write(bytes);
					
					// if testing ka disabled 
					closeConn = !ka;
				}
				System.out.println();
				
				// close only if we don't keep alive
				os.close();
				is.close();
			} catch (SSLHandshakeException e) {
				System.out.println("handshake: " + e.getMessage());
			} catch (SocketException e) {
				System.out.println("failed to handle con " + conn + ": " + e.getMessage());
			}
		}
	}
	
	public String getContent(String u, SSLSocketFactory sf) throws IOException {
		URL url = new URL(u);
		URLConnection conn = url.openConnection();
		conn.setDoInput(true);
//		conn.setDoOutput(true);
		if ("https".equals(url.getProtocol())) {
			HttpsURLConnection sconn = (HttpsURLConnection) conn;
		
			sconn.setSSLSocketFactory(sf);
		}
		
		//sconn.getOutputStream().write("test".getBytes("utf-8"));
		
		byte buff [] = new byte[255];
		int read = conn.getInputStream().read(buff);
		
		
		ByteArrayOutputStream baos = new ByteArrayOutputStream(8012);
		while (read > 0) {
			baos.write(buff, 0, read);
			read = conn.getInputStream().read(buff);
		}
		//sconn.getOutputStream().close();
		conn.getInputStream().close();
		
		return baos.toString("utf-8");
	}
	
	private static class RuntimeReplacingTrustManager implements X509TrustManager {
		X509TrustManager m;
		String tmName;
		
		public RuntimeReplacingTrustManager(X509TrustManager m, String name) {
			this.m = m;
			tmName = name;
		}
		
		public void setTM(X509TrustManager m, String name) {
			this.m = m;
			tmName = name;
		}
		
		@Override
		public void checkClientTrusted(X509Certificate[] chain, String authType)
				throws CertificateException {
			System.out.println("Checking whether client trusted with " + tmName);
			this.m.checkClientTrusted(chain, authType);
		}

		@Override
		public void checkServerTrusted(X509Certificate[] chain, String authType)
				throws CertificateException {
			System.out.println("Checking if server trusted with " + tmName);
			this.m.checkServerTrusted(chain, authType);
		}

		@Override
		public X509Certificate[] getAcceptedIssuers() {
			System.out.println("Getting accepted issuers with " + tmName);
			return this.m.getAcceptedIssuers();
		}
		
	}
	
	private static class LoggingTrustManager implements X509TrustManager {
		X509TrustManager m;
		public LoggingTrustManager(X509TrustManager m) {
			this. m = m;
		}
		@Override
		public void checkClientTrusted(X509Certificate[] chain, String authType)
				throws CertificateException {
			System.out.println("check client trusted");
			this.m.checkClientTrusted(chain, authType);
		}
		@Override
		public void checkServerTrusted(X509Certificate[] chain, String authType)
				throws CertificateException {
			System.out.println("check server trusted");
			this.m.checkServerTrusted(chain, authType);
		}
		@Override
		public X509Certificate[] getAcceptedIssuers() {
			System.out.println("get accepted issuers");
			return this.m.getAcceptedIssuers();
		}
		
	}
	
	@Test
	public void test() throws IOException, KeyManagementException, NoSuchAlgorithmException, KeyStoreException, CertificateException, UnrecoverableKeyException, InterruptedException {
		
		final boolean isSsl = true;
		final boolean ka = false;
		
		SSLContext context = SSLContext.getInstance("TLS");
		new Thread(new Runnable() {
			
			@Override
			public void run() {
				try {
					runDebugServer(isSsl, ka);
				} catch (Exception e) {
					e.printStackTrace();
				}
			}
			
		}, "http-server").start();
		
		
		TrustManager[] tms = getTrustManagerFromKeyStore(new File("client.jks"), "changeit");
		TrustManager[] tms2 = getTrustManagerFromKeyStore(new File("repl.jks"), "changeit");
		
		RuntimeReplacingTrustManager ltms = new RuntimeReplacingTrustManager((X509TrustManager) tms[0], "initial");
		
		System.out.println("TMS: " + tms.length);
		
		context.init(
				null,
				new TrustManager[] { ltms },
				SecureRandom.getInstance("SHA1PRNG"));
		
		
		
		System.setProperty("http.keepAlive", Boolean.toString(ka));
		System.setProperty("http.maxConnections", "1");
		
		s.acquire();
		
		SSLSocketFactory sf = context.getSocketFactory();
		SSLSocketFactory sf2 = context.getSocketFactory();
		
		assertFalse(sf.equals(sf2));
		
		String scheme = isSsl ? "https" : "http";
		
		String c = getContent(scheme + "://test.local:10000/", sf);
		System.out.println("Content: " + c);

		/*
		Enumeration<byte[]> ids = context.getClientSessionContext().getIds();
		while (ids.hasMoreElements()) {
			byte[] id = ids.nextElement();
			SSLSession session = context.getClientSessionContext().getSession(id);
			session.invalidate();
			System.out.println("Invalidated ssl session");
		}
		*/

		ltms.setTM((X509TrustManager) tms2[0], "replacement-1");
		
		String c2 = getContent(scheme + "://test.local:10000/bbb", sf);
		System.out.println("Content 2: " + c2);
		
	}
	
	public TrustManager[] getTrustManagerFromKeyStore(File ks, String password) throws NoSuchAlgorithmException, KeyStoreException, CertificateException, IOException {
		TrustManagerFactory tmf = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
		KeyStore tks = KeyStore.getInstance(KeyStore.getDefaultType());
		
		AutoCloseInputStream fs = new AutoCloseInputStream(new FileInputStream(ks));
		
		tks.load(fs, password.toCharArray());
		tmf.init(tks);
		
		TrustManager[] tms = tmf.getTrustManagers();
		
		return tms;
	}
	
}