infer compress type before decode NBT fail

nbt/nbt.go

@@ -26,7 +26,7 @@ const (
 
 type Decoder struct {
 	r interface {
-		io.ByteReader
+		io.ByteScanner
 		io.Reader
 	}
 }
@@ -34,7 +34,7 @@ type Decoder struct {
 func NewDecoder(r io.Reader) *Decoder {
 	d := new(Decoder)
 	if br, ok := r.(interface {
-		io.ByteReader
+		io.ByteScanner
 		io.Reader
 	}); ok {
 		d.r = br

nbt/read.go

@@ -18,6 +18,17 @@ func (d *Decoder) Decode(v interface{}) error {
 	if val.Kind() != reflect.Ptr {
 		return errors.New("non-pointer passed to Unmarshal")
 	}
+
+	// check the head
+	compress, err := d.checkCompressed()
+	if err != nil {
+		return fmt.Errorf("check compressed fail: %v", err)
+	}
+	if compress != "" {
+		return fmt.Errorf("data may compressed with %s", compress)
+	}
+
+	//start read NBT
 	tagType, tagName, err := d.readTag()
 	if err != nil {
 		return err
@@ -25,6 +36,29 @@ func (d *Decoder) Decode(v interface{}) error {
 	return d.unmarshal(val.Elem(), tagType, tagName)
 }
 
+// check the first byte and return if it use compress
+func (d *Decoder) checkCompressed() (compress string, err error) {
+	var head byte
+	head, err = d.r.ReadByte()
+	if err != nil {
+		return
+	}
+
+	if head <= 12 { //NBT
+		compress = ""
+	} else if head == 0x1f { //gzip
+		compress = "gzip"
+	} else if head == 0x78 { //zlib
+		compress = "zlib"
+	} else {
+		compress = "unknown"
+	}
+
+	err = d.r.UnreadByte()
+
+	return
+}
+
 // ErrEND error will be returned when reading a NBT with only Tag_End
 var ErrEND = errors.New("unexpected TAG_End")