WebSocketHandshake.java

/*
 *  Copyright (C) 2011 Roderick Baier
 *  
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *  
 *  	http://www.apache.org/licenses/LICENSE-2.0
 *  
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License. 
 */

package de.roderick.weberknecht;

import java.net.URI;
import java.nio.ByteBuffer;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Arrays;
import java.util.HashMap;


public class WebSocketHandshake
{	
	private String key1 = null;
	private String key2 = null;
	private byte[] key3 = null;
	private byte[] expectedServerResponse = null;
	
	private URI url = null;
	private String origin = null;
	private String protocol = null;
	
	
	public WebSocketHandshake(URI url, String protocol)
	{
		this.url = url;
		this.protocol = null;
		generateKeys();
	}
	
	
	public byte[] getHandshake()
	{
		String path = url.getPath();
		String host = url.getHost();
		origin = "http://" + host;
		
		String handshake = "GET " + path + " HTTP/1.1\r\n" +
				"Host: " + host + "\r\n" +
				"Connection: Upgrade\r\n" +
				"Sec-WebSocket-Key2: " + key2 + "\r\n";
		
		if (protocol != null) {
			handshake += "Sec-WebSocket-Protocol: " + protocol + "\r\n";
		}
		
		handshake += "Upgrade: WebSocket\r\n" +
				"Sec-WebSocket-Key1: " + key1 + "\r\n" +
				"Origin: " + origin + "\r\n" +
				"\r\n";
		
		byte[] handshakeBytes = new byte[handshake.getBytes().length + 8];
		System.arraycopy(handshake.getBytes(), 0, handshakeBytes, 0, handshake.getBytes().length);
		System.arraycopy(key3, 0, handshakeBytes, handshake.getBytes().length, 8);		
		
		return handshakeBytes;
	}
	
	
	public void verifyServerResponse(byte[] bytes)
		throws WebSocketException
	{
		if (!Arrays.equals(bytes, expectedServerResponse)) {
			throw new WebSocketException("not a WebSocket Server");
		}
	}
	
	
	public void verifyServerStatusLine(String statusLine)
		throws WebSocketException
	{
		int statusCode = Integer.valueOf(statusLine.substring(9, 12));
		
		if (statusCode == 407) {
			throw new WebSocketException("connection failed: proxy authentication not supported");
		}
		else if (statusCode == 404) {
			throw new WebSocketException("connection failed: 404 not found");
		}
		else if (statusCode != 101) {
			throw new WebSocketException("connection failed: unknown status code " + statusCode);
		}
	}
	
	
	public void verifyServerHandshakeHeaders(HashMap<String, String> headers)
		throws WebSocketException
	{
		if (!headers.get("Upgrade").equals("WebSocket")) {
			throw new WebSocketException("connection failed: missing header field in server handshake: Upgrade");
		}
		else if (!headers.get("Connection").equals("Upgrade")) {
			throw new WebSocketException("connection failed: missing header field in server handshake: Connection");
		}
		else if (!headers.get("Sec-WebSocket-Origin").equals(origin)) {
			throw new WebSocketException("connection failed: missing header field in server handshake: Sec-WebSocket-Origin");
		}
		
		// TODO see 4.1. step 41
//		else if (!headers.get("Sec-WebSocket-Location").equals(url.toASCIIString())) {
//			System.out.println("location: " + url.toASCIIString());
//		}
//		else if protocol
	}
		
	
	private void generateKeys()
	{
		int spaces1 = rand(1,12);
		int spaces2 = rand(1,12);
		
		int max1 = Integer.MAX_VALUE / spaces1;
		int max2 = Integer.MAX_VALUE / spaces2;
		
		int number1 = rand(0, max1);
		int number2 = rand(0, max2);
		
		int product1 = number1 * spaces1;
		int product2 = number2 * spaces2;
		
		key1 = Integer.toString(product1);
		key2 = Integer.toString(product2);
		
		key1 = insertRandomCharacters(key1);
		key2 = insertRandomCharacters(key2);
		
		key1 = insertSpaces(key1, spaces1);
		key2 = insertSpaces(key2, spaces2);
		
		key3 = createRandomBytes();
		
		ByteBuffer buffer = ByteBuffer.allocate(4);
		buffer.putInt(number1);
		byte[] number1Array = buffer.array();
		buffer = ByteBuffer.allocate(4);
		buffer.putInt(number2);
		byte[] number2Array = buffer.array();
		
		byte[] challenge = new byte[16];				
		System.arraycopy(number1Array, 0, challenge, 0, 4);
		System.arraycopy(number2Array, 0, challenge, 4, 4);
		System.arraycopy(key3, 0, challenge, 8, 8);

		expectedServerResponse = md5(challenge);
	}
	
	
	private String insertRandomCharacters(String key)
	{
		int count = rand(1, 12);
		
		char[] randomChars = new char[count];
		int randCount = 0;
		while (randCount < count) {
			int rand = (int) (Math.random() * 0x7e + 0x21);
			if (((0x21 < rand) && (rand < 0x2f)) || ((0x3a < rand) && (rand < 0x7e))) {
				randomChars[randCount] = (char) rand;
				randCount += 1;
			}
		}
		
		for (int i = 0; i < count; i++) {
			int split = rand(0, key.length());
			String part1 = key.substring(0, split);
			String part2 = key.substring(split);
			key = part1 + randomChars[i] + part2;
		}
		
		return key;
	}
	
	
	private String insertSpaces(String key, int spaces)
	{
		for (int i = 0; i < spaces; i++) {
			int split = rand(1, key.length()-1);
			String part1 = key.substring(0, split);
			String part2 = key.substring(split);
			key = part1 + " " + part2;
		}
		
		return key;
	}
	
	
	private byte[] createRandomBytes()
	{
		byte[] bytes = new byte[8];
		
		for (int i = 0; i < 8; i++) {
			bytes[i] = (byte) rand(0, 255);
		}
		
		return bytes;
	}
	
	
	private byte[] md5(byte[] bytes)
	{
		try {
			MessageDigest md = MessageDigest.getInstance("MD5");
			return md.digest(bytes);
		}
		catch (NoSuchAlgorithmException e) {
			return null;
		}
	}
	
	
	private int rand(int min, int max)
	{
		int rand = (int) (Math.random() * max + min);
		return rand;
	}
}