add support for custom types

Author: lunixbochs

binary.go

@@ -0,0 +1,52 @@
+package struc
+
+import (
+	"encoding/binary"
+	"io"
+	"reflect"
+)
+
+type byteWriter struct {
+	buf []byte
+	pos int
+}
+
+func (b byteWriter) Write(p []byte) (int, error) {
+	capacity := len(b.buf) - b.pos
+	if capacity < len(p) {
+		p = p[:capacity]
+	}
+	if len(p) > 0 {
+		copy(b.buf[b.pos:], p)
+		b.pos += len(p)
+	}
+	return len(p), nil
+}
+
+type binaryFallback reflect.Value
+
+func (b binaryFallback) String() string {
+	return b.String()
+}
+
+func (b binaryFallback) Sizeof(val reflect.Value, options *Options) int {
+	return binary.Size(val.Interface())
+}
+
+func (b binaryFallback) Pack(buf []byte, val reflect.Value, options *Options) (int, error) {
+	tmp := byteWriter{buf: buf}
+	var order binary.ByteOrder = binary.BigEndian
+	if options.Order != nil {
+		order = options.Order
+	}
+	err := binary.Write(tmp, order, val.Interface())
+	return tmp.pos, err
+}
+
+func (b binaryFallback) Unpack(r io.Reader, val reflect.Value, options *Options) error {
+	var order binary.ByteOrder = binary.BigEndian
+	if options.Order != nil {
+		order = options.Order
+	}
+	return binary.Read(r, order, val.Interface())
+}

custom.go

@@ -0,0 +1,33 @@
+package struc
+
+import (
+	"io"
+	"reflect"
+)
+
+type Custom interface {
+	Pack(p []byte, opt *Options) (int, error)
+	Unpack(r io.Reader, length int, opt *Options) error
+	Size(opt *Options) int
+	String() string
+}
+
+type customFallback struct {
+	custom Custom
+}
+
+func (c customFallback) Pack(p []byte, val reflect.Value, opt *Options) (int, error) {
+	return c.custom.Pack(p, opt)
+}
+
+func (c customFallback) Unpack(r io.Reader, val reflect.Value, opt *Options) error {
+	return c.custom.Unpack(r, 1, opt)
+}
+
+func (c customFallback) Sizeof(val reflect.Value, opt *Options) int {
+	return c.custom.Size(opt)
+}
+
+func (c customFallback) String() string {
+	return c.custom.String()
+}

custom_test.go

@@ -0,0 +1,97 @@
+package struc
+
+import (
+	"bytes"
+	"encoding/binary"
+	"io"
+	"strconv"
+	"testing"
+)
+
+type Int3 uint32
+
+func (i *Int3) Pack(p []byte, opt *Options) (int, error) {
+	var tmp [4]byte
+	binary.BigEndian.PutUint32(tmp[:], uint32(*i))
+	copy(p, tmp[1:])
+	return 3, nil
+}
+func (i *Int3) Unpack(r io.Reader, length int, opt *Options) error {
+	var tmp [4]byte
+	if _, err := r.Read(tmp[1:]); err != nil {
+		return err
+	}
+	*i = Int3(binary.BigEndian.Uint32(tmp[:]))
+	return nil
+}
+func (i *Int3) Size(opt *Options) int {
+	return 3
+}
+func (i *Int3) String() string {
+	return strconv.FormatUint(uint64(*i), 10)
+}
+
+func TestCustom(t *testing.T) {
+	var buf bytes.Buffer
+	var i Int3 = 3
+	if err := Pack(&buf, &i); err != nil {
+		t.Fatal(err)
+	}
+	if !bytes.Equal(buf.Bytes(), []byte{0, 0, 3}) {
+		t.Fatal("error packing custom int")
+	}
+	var i2 Int3
+	if err := Unpack(&buf, &i2); err != nil {
+		t.Fatal(err)
+	}
+	if i2 != 3 {
+		t.Fatal("error unpacking custom int")
+	}
+}
+
+type Int3Struct struct {
+	I Int3
+}
+
+func TestCustomStruct(t *testing.T) {
+	var buf bytes.Buffer
+	i := Int3Struct{3}
+	if err := Pack(&buf, &i); err != nil {
+		t.Fatal(err)
+	}
+	if !bytes.Equal(buf.Bytes(), []byte{0, 0, 3}) {
+		t.Fatal("error packing custom int struct")
+	}
+	var i2 Int3Struct
+	if err := Unpack(&buf, &i2); err != nil {
+		t.Fatal(err)
+	}
+	if i2.I != 3 {
+		t.Fatal("error unpacking custom int struct")
+	}
+}
+
+// TODO: slices of custom types don't work yet
+/*
+type Int3SliceStruct struct {
+	I [2]Int3
+}
+
+func TestCustomSliceStruct(t *testing.T) {
+	var buf bytes.Buffer
+	i := Int3SliceStruct{[2]Int3{3, 4}}
+	if err := Pack(&buf, &i); err != nil {
+		t.Fatal(err)
+	}
+	if !bytes.Equal(buf.Bytes(), []byte{0, 0, 3}) {
+		t.Fatal("error packing custom int struct")
+	}
+	var i2 Int3SliceStruct
+	if err := Unpack(&buf, &i2); err != nil {
+		t.Fatal(err)
+	}
+	if i2.I[0] != 3 && i2.I[1] != 4 {
+		t.Fatal("error unpacking custom int struct")
+	}
+}
+*/

field.go

@@ -29,8 +29,7 @@ func (f *Field) String() string {
 	if f.Type == Pad {
 		return fmt.Sprintf("{type: Pad, len: %d}", f.Len)
 	} else {
-		typeName := typeNames[f.Type]
-		out = fmt.Sprintf("type: %s, order: %v", typeName, f.Order)
+		out = fmt.Sprintf("type: %s, order: %v", f.Type.String(), f.Order)
 	}
 	if f.Sizefrom != nil {
 		out += fmt.Sprintf(", sizefrom: %v", f.Sizefrom)
@@ -65,6 +64,8 @@ func (f *Field) Size(val reflect.Value, options *Options) int {
 			length = f.Len
 		}
 		size = length * typ.Size()
+	} else if typ == CustomType {
+		return val.Addr().Interface().(Custom).Size(options)
 	} else {
 		size = typ.Size()
 	}
@@ -137,6 +138,10 @@ func (f *Field) packVal(buf []byte, val reflect.Value, length int, options *Opti
 			size = val.Len()
 			copy(buf, val.Bytes())
 		}
+	case CustomType:
+		return val.Addr().Interface().(Custom).Pack(buf, options)
+	default:
+		panic(fmt.Sprintf("no pack handler for type: %s", typ))
 	}
 	return
 }
@@ -214,6 +219,8 @@ func (f *Field) unpackVal(buf []byte, val reflect.Value, length int, options *Op
 		default:
 			val.SetUint(n)
 		}
+	default:
+		panic(fmt.Sprintf("no unpack handler for type: %s", typ))
 	}
 	return nil
 }

fields.go

@@ -137,18 +137,23 @@ func (f Fields) Unpack(r io.Reader, val reflect.Value, options *Options) error {
 			}
 			continue
 		} else {
-			size := length * field.Type.Resolve(options).Size()
-			if size < 8 {
-				buf = tmp[:size]
+			typ := field.Type.Resolve(options)
+			if typ == CustomType {
+				return v.Addr().Interface().(Custom).Unpack(r, length, options)
 			} else {
-				buf = make([]byte, size)
-			}
-			if _, err := io.ReadFull(r, buf); err != nil {
-				return err
-			}
-			err := field.Unpack(buf[:size], v, length, options)
-			if err != nil {
-				return err
+				size := length * field.Type.Resolve(options).Size()
+				if size < 8 {
+					buf = tmp[:size]
+				} else {
+					buf = make([]byte, size)
+				}
+				if _, err := io.ReadFull(r, buf); err != nil {
+					return err
+				}
+				err := field.Unpack(buf[:size], v, length, options)
+				if err != nil {
+					return err
+				}
 			}
 		}
 	}

packable.go

@@ -0,0 +1,13 @@
+package struc
+
+import (
+	"io"
+	"reflect"
+)
+
+type Packer interface {
+	Pack(buf []byte, val reflect.Value, options *Options) (int, error)
+	Unpack(r io.Reader, val reflect.Value, options *Options) error
+	Sizeof(val reflect.Value, options *Options) int
+	String() string
+}

packer.go

@@ -73,6 +73,12 @@ func parseField(f reflect.StructField) (fd *Field, err error) {
 		fd.Ptr = true
 		fd.kind = f.Type.Elem().Kind()
 	}
+	// check for custom types
+	tmp := reflect.New(f.Type)
+	if _, ok := tmp.Interface().(Custom); ok {
+		fd.Type = CustomType
+		return
+	}
 	var defTypeOk bool
 	fd.defType, defTypeOk = reflectTypeMap[fd.kind]
 	// find a type in the struct tag