diff --git a/src/main/java/com/profiler/dto/Header.java b/src/main/java/com/profiler/dto/Header.java index 27029062f..cde7e33ee 100644 --- a/src/main/java/com/profiler/dto/Header.java +++ b/src/main/java/com/profiler/dto/Header.java @@ -2,7 +2,9 @@ package com.profiler.dto; public class Header { - private byte signature = (byte) 0xef; + public static final byte SIGNATURE = (byte) 0xef; + + private byte signature = SIGNATURE; private byte version = 0x10; private short type = 0; @@ -38,5 +40,14 @@ public class Header { public void setType(short type) { this.type = type; } + + @Override + public String toString() { + return "Header{" + + "signature=" + signature + + ", version=" + version + + ", type=" + type + + '}'; + } } diff --git a/src/main/java/com/profiler/util/DefaultTBaseLocator.java b/src/main/java/com/profiler/util/DefaultTBaseLocator.java new file mode 100644 index 000000000..10bab8ec1 --- /dev/null +++ b/src/main/java/com/profiler/util/DefaultTBaseLocator.java @@ -0,0 +1,20 @@ +package com.profiler.util; + +import com.profiler.dto.*; +import org.apache.thrift.TBase; + +public class DefaultTBaseLocator implements TBaseLocator{ + @Override + public TBase lookup(Header header) { + short type = header.getType(); + switch (type) { + case 10: + return new JVMInfoThriftDTO(); + case 20: + return new RequestDataListThriftDTO(); + case 30: + return new RequestThriftDTO(); + } + throw new IllegalArgumentException("Unsupported type:" + type + " " + header); + } +} diff --git a/src/main/java/com/profiler/util/HeaderTBaseDeserializer.java b/src/main/java/com/profiler/util/HeaderTBaseDeserializer.java index e0c17ccab..5a68b3644 100644 --- a/src/main/java/com/profiler/util/HeaderTBaseDeserializer.java +++ b/src/main/java/com/profiler/util/HeaderTBaseDeserializer.java @@ -50,15 +50,17 @@ public class HeaderTBaseDeserializer { } private void validate(Header header) { - //To change body of created methods use File | Settings | File Templates. + boolean accepted = HeaderUtil.validateSignature(header.getSignature()); + if(!accepted) { + throw new IllegalArgumentException("Invalid Signature:" + header); + } } private Header readHeader() throws TException { byte signature = protocol_.readByte(); byte version = protocol_.readByte(); short type = protocol_.readI16(); - Header header = new Header(signature, version, type); - return header; + return new Header(signature, version, type); } /** diff --git a/src/main/java/com/profiler/util/HeaderTBaseSerializer.java b/src/main/java/com/profiler/util/HeaderTBaseSerializer.java index d429d64e1..ecd78ab9d 100644 --- a/src/main/java/com/profiler/util/HeaderTBaseSerializer.java +++ b/src/main/java/com/profiler/util/HeaderTBaseSerializer.java @@ -60,14 +60,18 @@ public class HeaderTBaseSerializer { */ public byte[] serialize(Header header, TBase base) throws TException { baos_.reset(); - protocol_.writeByte(header.getSignature()); - protocol_.writeByte(header.getVersion()); - protocol_.writeI16(header.getType()); + writeHeader(header); base.write(protocol_); return baos_.toByteArray(); } - /** + private void writeHeader(Header header) throws TException { + protocol_.writeByte(header.getSignature()); + protocol_.writeByte(header.getVersion()); + protocol_.writeI16(header.getType()); + } + + /** * Serialize the Thrift object into a Java string, using a specified * character set for encoding. * diff --git a/src/main/java/com/profiler/util/HeaderUtil.java b/src/main/java/com/profiler/util/HeaderUtil.java index cba6cc10a..2672e64cb 100644 --- a/src/main/java/com/profiler/util/HeaderUtil.java +++ b/src/main/java/com/profiler/util/HeaderUtil.java @@ -1,4 +1,10 @@ package com.profiler.util; + +import com.profiler.dto.Header; + public class HeaderUtil { + public static boolean validateSignature(byte signature){ + return Header.SIGNATURE == signature; + } }