.

go/neo/proto.go

@@ -243,15 +243,22 @@ func (a *Address) neoEncode(b []byte) int {
 	return n
 }
 
-func (a *Address) neoDecode(b []byte) int {
-	n := string_neoDecode(&a.Host, b)
+func (a *Address) neoDecode(b []byte) (int, bool) {
+	n, ok := string_neoDecode(&a.Host, b)
+	if !ok {
+		return 0, false
+	}
 	if a.Host != "" {
+		b = b[n:]
+		if len(b) < 2 {
+			return 0, false
+		}
 		a.Port = binary.BigEndian.Uint16(b[n:])
 		n += 2
 	} else {
 		a.Port = 0
 	}
-	return n
+	return n, true
 }
 
 // Checksum is a SHA1 hash.
@@ -1010,7 +1017,16 @@ type Truncate struct {
 	// XXX _answer = Error
 }
 
-// ---- runtime support for protogen and custom encodings ----
+// ---- runtime support for protogen and custom codecs ----
+
+// customCodec is the interface that is implemented by types with custom encodings.
+//
+// its semantic is very similar to Msg.
+type customCodec interface {
+	neoEncodedLen() int
+	neoEncode(buf []byte) (nwrote int)
+	neoDecode(data []byte) (nread int, ok bool)
+}
 
 func byte2bool(b byte) bool {
 	return b != 0
@@ -1042,3 +1058,32 @@ func float64_NEODecode(b []byte) float64 {
 	fu := binary.BigEndian.Uint64(b)
 	return math.Float64frombits(fu)
 }
+
+
+// XXX we need string_neo* only for Address
+// XXX dup of genSlice1 in protogen.go
+
+func string_neoEncodedLen(s string) int {
+	return 4 + len(s)
+}
+
+func string_neoEncode(s string, data []byte) int {
+	l := len(s)
+	binary.BigEndian.PutUint32(data, uint32(l))
+	copy(data[4:4+l], s) // NOTE [:l] to catch data overflow as copy copies minimal len
+	return 4 + l
+}
+
+func string_neoDecode(sp *string, data []byte) (nread int, ok bool) {
+	if len(data) < 4 {
+		return 0, false
+	}
+	l := binary.BigEndian.Uint32(data)
+	data = data[4:]
+	if uint32(len(data)) < l {
+		return 0, false
+	}
+
+	*sp = string(data[:l])
+	return 4 + int(l), true
+}

go/neo/protogen.go

@@ -438,9 +438,9 @@ type CodeGenerator interface {
 	genArray1(path string, typ *types.Array)
 	genSlice1(path string, typ types.Type)
 
-	// generate code for a custom type which implements its own encoding/decodeing
-	// XXX review text
-	genCustom(path string, typ types.Type)	// XXX + obj?
+	// generate code for a custom type which implements its own
+	// encoding/decoding via implementing neo.customCodec interface.
+	genCustom(path string)
 
 	// get generated code.
 	generatedCode() string
@@ -1068,25 +1068,32 @@ func (d *decoder) genMap(assignto string, typ *types.Map, obj types.Object) {
 }
 
 // emit code to size/encode/decode custom type
-// XXX typ not needed
-func (s *sizer) genCustom(path string, typ *types.Type) {
+func (s *sizer) genCustom(path string) {
 	s.size.AddExpr("%s.neoEncodedLen()", path)
 }
 
-// XXX typ unused
-func (e *encoder) genCustom(path string, typ *types.Type) {
+func (e *encoder) genCustom(path string) {
 	e.emit("{")
-	e.emit("n := %s.neoEncodedLen()", path)
-	e.emit("%s.neoEncode(data[%v:])", e.n)
+	e.emit("n := %s.neoEncode(data[%v:])", e.n)
 	e.emit("data = data[%v + n:]", e.n)
 	e.emit("}")
 	e.n = 0
 }
 
-func (e *decoder) genCustom(path string) {
+func (d *decoder) genCustom(path string) {
+	// make sure we check for overflow previous-code before proceeding to custom decoder.
+	d.overflowCheck()
+	d.resetPos()
+
 	d.emit("{")
-	d.emit("n := %s.neoDecode(data[%v ...")	// XXX error!
+	d.emit("n, ok := %s.neoDecode(data)")
+	d.emit("if !ok { goto overflow }")
+	d.emit("data = data[n:]")
 	d.emit("}")
+
+	// insert overflow checkpoint after custom decoder so that overflow
+	// checks for following code are inserted after custom decoder call.
+	d.overflowCheck()
 }
 
 // top-level driver for emitting size/encode/decode code for a type
@@ -1095,7 +1102,7 @@ func (e *decoder) genCustom(path string) {
 // we can point to source location for where it happened)
 func codegenType(path string, typ types.Type, obj types.Object, codegen CodeGenerator) {
 	if types.Implements(typ, neoCustomXXX) {
-		codegen.genCustom(path, typ)
+		codegen.genCustom(path)
 		return
 	}