Add interface for customizing packet headers.

example/org/spacehq/packetlib/test/TestProtocol.java

@@ -5,6 +5,8 @@ import org.spacehq.packetlib.Server;
 import org.spacehq.packetlib.Session;
 import org.spacehq.packetlib.crypt.AESEncryption;
 import org.spacehq.packetlib.crypt.PacketEncryption;
+import org.spacehq.packetlib.packet.DefaultPacketHeader;
+import org.spacehq.packetlib.packet.PacketHeader;
 import org.spacehq.packetlib.packet.PacketProtocol;
 
 import javax.crypto.SecretKey;
@@ -12,6 +14,7 @@ import java.security.GeneralSecurityException;
 
 public class TestProtocol extends PacketProtocol {
 
+	private PacketHeader header = new DefaultPacketHeader();
 	private AESEncryption encrypt;
 
 	@SuppressWarnings("unused")
@@ -36,6 +39,11 @@ public class TestProtocol extends PacketProtocol {
 		return this.encrypt;
 	}
 
+	@Override
+	public PacketHeader getPacketHeader() {
+		return this.header;
+	}
+
 	@Override
 	public void newClientSession(Client client, Session session) {
 		session.addListener(new ClientSessionListener());

src/main/java/org/spacehq/packetlib/packet/DefaultPacketHeader.java

@@ -0,0 +1,59 @@
+package org.spacehq.packetlib.packet;
+
+import org.spacehq.packetlib.io.NetInput;
+import org.spacehq.packetlib.io.NetOutput;
+
+import java.io.IOException;
+
+public class DefaultPacketHeader implements PacketHeader {
+
+	@Override
+	public boolean isLengthVariable() {
+		return true;
+	}
+
+	@Override
+	public int getLengthSize() {
+		return 5;
+	}
+
+	@Override
+	public int getLengthSize(int length) {
+		return varintLength(length);
+	}
+
+	@Override
+	public int readLength(NetInput in, int available) throws IOException {
+		return in.readVarInt();
+	}
+
+	@Override
+	public void writeLength(NetOutput out, int length) throws IOException {
+		out.writeVarInt(length);
+	}
+
+	@Override
+	public int readPacketId(NetInput in) throws IOException {
+		return in.readVarInt();
+	}
+
+	@Override
+	public void writePacketId(NetOutput out, int packetId) throws IOException {
+		out.writeVarInt(packetId);
+	}
+
+	private static int varintLength(int i) {
+		if((i & -128) == 0) {
+			return 1;
+		} else if((i & -16384) == 0) {
+			return 2;
+		} else if((i & -2097152) == 0) {
+			return 3;
+		} else if((i & -268435456) == 0) {
+			return 4;
+		} else {
+			return 5;
+		}
+	}
+
+}

src/main/java/org/spacehq/packetlib/packet/PacketHeader.java

@@ -0,0 +1,72 @@
+package org.spacehq.packetlib.packet;
+
+import org.spacehq.packetlib.io.NetInput;
+import org.spacehq.packetlib.io.NetOutput;
+
+import java.io.IOException;
+
+/**
+ * The header of a protocol's packets.
+ */
+public interface PacketHeader {
+
+	/**
+	 * Gets whether the header's length value can vary in size.
+	 *
+	 * @return Whether the header's length value can vary in size.
+	 */
+	public boolean isLengthVariable();
+
+	/**
+	 * Gets the size of the header's length value.
+	 *
+	 * @return The length value's size.
+	 */
+	public int getLengthSize();
+
+	/**
+	 * Gets the size of the header's length value.
+	 *
+	 * @param length Length value to get the size of.
+	 * @return The length value's size.
+	 */
+	public int getLengthSize(int length);
+
+	/**
+	 * Reads the length of a packet from the given input.
+	 *
+	 * @param in Input to read from.
+	 * @param available Number of available bytes in the NetInput.
+	 * @return The resulting packet length.
+	 * @throws java.io.IOException If an I/O error occurs.
+	 */
+	public int readLength(NetInput in, int available) throws IOException;
+
+	/**
+	 * Writes the length of a packet to the given output.
+	 *
+	 * @param out Output to write to.
+	 * @param length Length to write.
+	 * @throws java.io.IOException If an I/O error occurs.
+	 */
+	public void writeLength(NetOutput out, int length) throws IOException;
+
+	/**
+	 * Reads the ID of a packet from the given input.
+	 *
+	 * @param in Input to read from.
+	 * @return The resulting packet ID.
+	 * @throws java.io.IOException If an I/O error occurs.
+	 */
+	public int readPacketId(NetInput in) throws IOException;
+
+	/**
+	 * Writes the ID of a packet to the given output.
+	 *
+	 * @param out Output to write to.
+	 * @param packetId Packet ID to write.
+	 * @throws java.io.IOException If an I/O error occurs.
+	 */
+	public void writePacketId(NetOutput out, int packetId) throws IOException;
+
+}

src/main/java/org/spacehq/packetlib/packet/PacketProtocol.java

@@ -18,6 +18,13 @@ public abstract class PacketProtocol {
 	private final Map<Integer, Class<? extends Packet>> incoming = new HashMap<Integer, Class<? extends Packet>>();
 	private final Map<Class<? extends Packet>, Integer> outgoing = new HashMap<Class<? extends Packet>, Integer>();
 
+	/**
+	 * Gets the packet header of this protocol.
+	 *
+	 * @return The protocol's packet header.
+	 */
+	public abstract PacketHeader getPacketHeader();
+
 	/**
 	 * Called when a client is created with this protocol.
 	 *

src/main/java/org/spacehq/packetlib/tcp/TcpPacketCodec.java

@@ -24,14 +24,14 @@ public class TcpPacketCodec extends ByteToMessageCodec<Packet> {
 	@Override
 	public void encode(ChannelHandlerContext ctx, Packet packet, ByteBuf buf) throws Exception {
 		NetOutput out = new ByteBufNetOutput(buf);
-		out.writeVarInt(this.session.getPacketProtocol().getOutgoingId(packet.getClass()));
+		this.session.getPacketProtocol().getPacketHeader().writePacketId(out, this.session.getPacketProtocol().getOutgoingId(packet.getClass()));
 		packet.write(out);
 	}
 
 	@Override
 	protected void decode(ChannelHandlerContext ctx, ByteBuf buf, List<Object> out) throws Exception {
 		NetInput in = new ByteBufNetInput(buf);
-		int id = in.readVarInt();
+		int id = this.session.getPacketProtocol().getPacketHeader().readPacketId(in);
 		Packet packet = this.session.getPacketProtocol().createIncomingPacket(id);
 		packet.read(in);
 		if(packet.isPriority()) {

src/main/java/org/spacehq/packetlib/tcp/TcpPacketSizer.java

@@ -5,6 +5,7 @@ import io.netty.buffer.Unpooled;
 import io.netty.channel.ChannelHandlerContext;
 import io.netty.handler.codec.ByteToMessageCodec;
 import io.netty.handler.codec.CorruptedFrameException;
+import org.spacehq.packetlib.Session;
 import org.spacehq.packetlib.tcp.io.ByteBufNetInput;
 import org.spacehq.packetlib.tcp.io.ByteBufNetOutput;
 
@@ -12,51 +13,48 @@ import java.util.List;
 
 public class TcpPacketSizer extends ByteToMessageCodec<ByteBuf> {
 
+	private Session session;
+
+	public TcpPacketSizer(Session session) {
+		this.session = session;
+	}
+
 	@Override
 	public void encode(ChannelHandlerContext ctx, ByteBuf in, ByteBuf out) throws Exception {
 		int length = in.readableBytes();
-		out.ensureWritable(varintLength(length) + length);
-		new ByteBufNetOutput(out).writeVarInt(length);
+		out.ensureWritable(this.session.getPacketProtocol().getPacketHeader().getLengthSize(length) + length);
+		this.session.getPacketProtocol().getPacketHeader().writeLength(new ByteBufNetOutput(out), length);
 		out.writeBytes(in);
 	}
 
 	@Override
 	protected void decode(ChannelHandlerContext ctx, ByteBuf buf, List<Object> out) throws Exception {
-		buf.markReaderIndex();
-		byte[] lengthBytes = new byte[5];
-		for(int index = 0; index < lengthBytes.length; index++) {
-			if(!buf.isReadable()) {
-				buf.resetReaderIndex();
-				return;
-			}
-
-			lengthBytes[index] = buf.readByte();
-			if(lengthBytes[index] >= 0) {
-				int length = new ByteBufNetInput(Unpooled.wrappedBuffer(lengthBytes)).readVarInt();
-				if(buf.readableBytes() < length) {
+		int size = this.session.getPacketProtocol().getPacketHeader().getLengthSize();
+		if(size > 0) {
+			buf.markReaderIndex();
+			byte[] lengthBytes = new byte[size];
+			for(int index = 0; index < lengthBytes.length; index++) {
+				if(!buf.isReadable()) {
 					buf.resetReaderIndex();
 					return;
 				}
 
-				out.add(buf.readBytes(length));
-				return;
+				lengthBytes[index] = buf.readByte();
+				if((this.session.getPacketProtocol().getPacketHeader().isLengthVariable() && lengthBytes[index] >= 0) || index == size - 1) {
+					int length = this.session.getPacketProtocol().getPacketHeader().readLength(new ByteBufNetInput(Unpooled.wrappedBuffer(lengthBytes)), buf.readableBytes());
+					if(buf.readableBytes() < length) {
+						buf.resetReaderIndex();
+						return;
+					}
+
+					out.add(buf.readBytes(length));
+					return;
+				}
 			}
-		}
 
-		throw new CorruptedFrameException("Length is too long.");
-	}
-
-	private static int varintLength(int i) {
-		if((i & -128) == 0) {
-			return 1;
-		} else if((i & -16384) == 0) {
-			return 2;
-		} else if((i & -2097152) == 0) {
-			return 3;
-		} else if((i & -268435456) == 0) {
-			return 4;
+			throw new CorruptedFrameException("Length is too long.");
 		} else {
-			return 5;
+			out.add(buf);
 		}
 	}
 

src/main/java/org/spacehq/packetlib/tcp/TcpSessionFactory.java

@@ -50,7 +50,7 @@ public class TcpSessionFactory implements SessionFactory {
 				ch.pipeline()
 						.addLast("timer", new ReadTimeoutHandler(30))
 						.addLast("encryption", new TcpPacketEncryptor(session))
-						.addLast("sizer", new TcpPacketSizer())
+						.addLast("sizer", new TcpPacketSizer(session))
 						.addLast("codec", new TcpPacketCodec(session))
 						.addLast("manager", session);
 			}
@@ -73,7 +73,7 @@ public class TcpSessionFactory implements SessionFactory {
 				ch.pipeline()
 						.addLast("timer", new ReadTimeoutHandler(30))
 						.addLast("encryption", new TcpPacketEncryptor(session))
-						.addLast("sizer", new TcpPacketSizer())
+						.addLast("sizer", new TcpPacketSizer(session))
 						.addLast("codec", new TcpPacketCodec(session))
 						.addLast("manager", session);
 				server.addSession(session);