Optimize calls to UnmarshalCBOR() for RawTag, etc.

Currently, unreleased changes in PR #636 and #645 cause the
input data to be checked twice when UnmarshalCBOR() is
called internally by Unmarshal() for:
- ByteString
- RawTag
- SimpleValue

UnmarshalCBOR() checks input data because it can be called by
user apps providing bad data. However, the codec already checks
input data before internally calling UnmarshalCBOR() so the
2nd check is redundant.

This commit avoids redundant check on the input data by having
Unmarshal() call the private unmarshalCBOR() if implemented
by ByteString, RawTag, SimpleValue, etc.:
- Internally, the codec calls the private unmarshalCBOR() to
  avoid the redundant check on input data.
- Externally, UnmarshalCBOR() is available as a wrapper that
  checks input data before calling the private unmarshalCBOR().

UnmarshalCBOR() for ByteString, RawTag, and SimpleValue are marked
as deprecated and Unmarshal() should be used instead.

Author: fxamacker

bytestring.go

@@ -38,24 +38,18 @@ func (bs ByteString) MarshalCBOR() ([]byte, error) {
 
 // UnmarshalCBOR decodes CBOR byte string (major type 2) to ByteString.
 // Decoding CBOR null and CBOR undefined sets ByteString to be empty.
+//
+// Deprecated: No longer used by this codec; kept for compatibility
+// with user apps that directly call this function.
 func (bs *ByteString) UnmarshalCBOR(data []byte) error {
 	if bs == nil {
 		return errors.New("cbor.ByteString: UnmarshalCBOR on nil pointer")
 	}
 
-	// Decoding CBOR null and CBOR undefined to ByteString resets data.
-	// This behavior is similar to decoding CBOR null and CBOR undefined to []byte.
-	if len(data) == 1 && (data[0] == 0xf6 || data[0] == 0xf7) {
-		*bs = ""
-		return nil
-	}
-
 	d := decoder{data: data, dm: defaultDecMode}
 
 	// Check well-formedness of CBOR data item.
-	// NOTE: well-formedness check here is redundant when
-	// Unmarshal() invokes ByteString.UnmarshalCBOR().
-	// However, ByteString.UnmarshalCBOR() is exported, so
+	// ByteString.UnmarshalCBOR() is exported, so
 	// the codec needs to support same behavior for:
 	// - Unmarshal(data, *ByteString)
 	// - ByteString.UnmarshalCBOR(data)
@@ -64,8 +58,26 @@ func (bs *ByteString) UnmarshalCBOR(data []byte) error {
 		return err
 	}
 
-	// Restore decoder offset after well-formedness check.
-	d.off = 0
+	return bs.unmarshalCBOR(data)
+}
+
+// unmarshalCBOR decodes CBOR byte string (major type 2) to ByteString.
+// Decoding CBOR null and CBOR undefined sets ByteString to be empty.
+// This function assumes data is well-formed, and does not perform bounds checking.
+// This function is called by Unmarshal().
+func (bs *ByteString) unmarshalCBOR(data []byte) error {
+	if bs == nil {
+		return errors.New("cbor.ByteString: UnmarshalCBOR on nil pointer")
+	}
+
+	// Decoding CBOR null and CBOR undefined to ByteString resets data.
+	// This behavior is similar to decoding CBOR null and CBOR undefined to []byte.
+	if len(data) == 1 && (data[0] == 0xf6 || data[0] == 0xf7) {
+		*bs = ""
+		return nil
+	}
+
+	d := decoder{data: data, dm: defaultDecMode}
 
 	// Check if CBOR data type is byte string
 	if typ := d.nextCBORType(); typ != cborTypeByteString {

bytestring_test.go

@@ -195,16 +195,16 @@ func TestUnmarshalByteStringOnBadData(t *testing.T) {
 					t.Errorf("UnmarshalCBOR(%x) returned error %q, want %q", tc.data, err.Error(), tc.errMsg)
 				}
 			}
-			// Test Unmarshal(data, *ByteString), which calls ByteString.UnmarshalCBOR() under the hood
+			// Test Unmarshal(data, *ByteString), which calls ByteString.unmarshalCBOR() under the hood
 			{
 				var v ByteString
 
 				err := Unmarshal(tc.data, &v)
 				if err == nil {
-					t.Errorf("UnmarshalCBOR(%x) didn't return error", tc.data)
+					t.Errorf("Unmarshal(%x) didn't return error", tc.data)
 				}
 				if !strings.HasPrefix(err.Error(), tc.errMsg) {
-					t.Errorf("UnmarshalCBOR(%x) returned error %q, want %q", tc.data, err.Error(), tc.errMsg)
+					t.Errorf("Unmarshal(%x) returned error %q, want %q", tc.data, err.Error(), tc.errMsg)
 				}
 			}
 		})

cache.go

@@ -32,6 +32,7 @@ type specialType int
 const (
 	specialTypeNone specialType = iota
 	specialTypeUnmarshalerIface
+	specialTypeUnexportedUnmarshalerIface
 	specialTypeEmptyIface
 	specialTypeIface
 	specialTypeTag
@@ -70,6 +71,8 @@ func newTypeInfo(t reflect.Type) *typeInfo {
 		tInfo.spclType = specialTypeTag
 	} else if t == typeTime {
 		tInfo.spclType = specialTypeTime
+	} else if reflect.PointerTo(t).Implements(typeUnexportedUnmarshaler) {
+		tInfo.spclType = specialTypeUnexportedUnmarshalerIface
 	} else if reflect.PointerTo(t).Implements(typeUnmarshaler) {
 		tInfo.spclType = specialTypeUnmarshalerIface
 	}

decode.go

@@ -151,6 +151,10 @@ type Unmarshaler interface {
 	UnmarshalCBOR([]byte) error
 }
 
+type unmarshaler interface {
+	unmarshalCBOR([]byte) error
+}
+
 // InvalidUnmarshalError describes an invalid argument passed to Unmarshal.
 type InvalidUnmarshalError struct {
 	s string
@@ -1460,6 +1464,9 @@ func (d *decoder) parseToValue(v reflect.Value, tInfo *typeInfo) error { //nolin
 
 		case specialTypeUnmarshalerIface:
 			return d.parseToUnmarshaler(v)
+
+		case specialTypeUnexportedUnmarshalerIface:
+			return d.parseToUnexportedUnmarshaler(v)
 		}
 	}
 
@@ -1805,6 +1812,26 @@ func (d *decoder) parseToUnmarshaler(v reflect.Value) error {
 	return errors.New("cbor: failed to assert " + v.Type().String() + " as cbor.Unmarshaler")
 }
 
+// parseToUnexportedUnmarshaler parses CBOR data to value implementing unmarshaler interface.
+// It assumes data is well-formed, and does not perform bounds checking.
+func (d *decoder) parseToUnexportedUnmarshaler(v reflect.Value) error {
+	if d.nextCBORNil() && v.Kind() == reflect.Pointer && v.IsNil() {
+		d.skip()
+		return nil
+	}
+
+	if v.Kind() != reflect.Pointer && v.CanAddr() {
+		v = v.Addr()
+	}
+	if u, ok := v.Interface().(unmarshaler); ok {
+		start := d.off
+		d.skip()
+		return u.unmarshalCBOR(d.data[start:d.off])
+	}
+	d.skip()
+	return errors.New("cbor: failed to assert " + v.Type().String() + " as cbor.unmarshaler")
+}
+
 // parse parses CBOR data and returns value in default Go type.
 // It assumes data is well-formed, and does not perform bounds checking.
 func (d *decoder) parse(skipSelfDescribedTag bool) (any, error) { //nolint:gocyclo
@@ -2969,13 +2996,14 @@ func (d *decoder) nextCBORNil() bool {
 }
 
 var (
-	typeIntf              = reflect.TypeOf([]any(nil)).Elem()
-	typeTime              = reflect.TypeOf(time.Time{})
-	typeBigInt            = reflect.TypeOf(big.Int{})
-	typeUnmarshaler       = reflect.TypeOf((*Unmarshaler)(nil)).Elem()
-	typeBinaryUnmarshaler = reflect.TypeOf((*encoding.BinaryUnmarshaler)(nil)).Elem()
-	typeString            = reflect.TypeOf("")
-	typeByteSlice         = reflect.TypeOf([]byte(nil))
+	typeIntf                  = reflect.TypeOf([]any(nil)).Elem()
+	typeTime                  = reflect.TypeOf(time.Time{})
+	typeBigInt                = reflect.TypeOf(big.Int{})
+	typeUnmarshaler           = reflect.TypeOf((*Unmarshaler)(nil)).Elem()
+	typeUnexportedUnmarshaler = reflect.TypeOf((*unmarshaler)(nil)).Elem()
+	typeBinaryUnmarshaler     = reflect.TypeOf((*encoding.BinaryUnmarshaler)(nil)).Elem()
+	typeString                = reflect.TypeOf("")
+	typeByteSlice             = reflect.TypeOf([]byte(nil))
 )
 
 func fillNil(_ cborType, v reflect.Value) error {

simplevalue.go

@@ -48,6 +48,9 @@ func (sv SimpleValue) MarshalCBOR() ([]byte, error) {
 }
 
 // UnmarshalCBOR decodes CBOR simple value (major type 7) to SimpleValue.
+//
+// Deprecated: No longer used by this codec; kept for compatibility
+// with user apps that directly call this function.
 func (sv *SimpleValue) UnmarshalCBOR(data []byte) error {
 	if sv == nil {
 		return errors.New("cbor.SimpleValue: UnmarshalCBOR on nil pointer")
@@ -56,9 +59,7 @@ func (sv *SimpleValue) UnmarshalCBOR(data []byte) error {
 	d := decoder{data: data, dm: defaultDecMode}
 
 	// Check well-formedness of CBOR data item.
-	// NOTE: well-formedness check here is redundant when
-	// Unmarshal() invokes SimpleValue.UnmarshalCBOR().
-	// However, SimpleValue.UnmarshalCBOR() is exported, so
+	// SimpleValue.UnmarshalCBOR() is exported, so
 	// the codec needs to support same behavior for:
 	// - Unmarshal(data, *SimpleValue)
 	// - SimpleValue.UnmarshalCBOR(data)
@@ -67,8 +68,18 @@ func (sv *SimpleValue) UnmarshalCBOR(data []byte) error {
 		return err
 	}
 
-	// Restore decoder offset after well-formedness check.
-	d.off = 0
+	return sv.unmarshalCBOR(data)
+}
+
+// unmarshalCBOR decodes CBOR simple value (major type 7) to SimpleValue.
+// This function assumes data is well-formed, and does not perform bounds checking.
+// This function is called by Unmarshal().
+func (sv *SimpleValue) unmarshalCBOR(data []byte) error {
+	if sv == nil {
+		return errors.New("cbor.SimpleValue: UnmarshalCBOR on nil pointer")
+	}
+
+	d := decoder{data: data, dm: defaultDecMode}
 
 	typ, ai, val := d.getHead()
 

simplevalue_test.go

@@ -156,16 +156,16 @@ func TestUnmarshalSimpleValueOnBadData(t *testing.T) {
 					t.Errorf("UnmarshalCBOR(%x) returned error %q, want %q", tc.data, err.Error(), tc.errMsg)
 				}
 			}
-			// Test Unmarshal(data, *SimpleValue), which calls SimpleValue.UnmarshalCBOR() under the hood
+			// Test Unmarshal(data, *SimpleValue), which calls SimpleValue.unmarshalCBOR() under the hood
 			{
 				var v SimpleValue
 
 				err := Unmarshal(tc.data, &v)
 				if err == nil {
-					t.Errorf("UnmarshalCBOR(%x) didn't return error", tc.data)
+					t.Errorf("Unmarshal(%x) didn't return error", tc.data)
 				}
 				if !strings.HasPrefix(err.Error(), tc.errMsg) {
-					t.Errorf("UnmarshalCBOR(%x) returned error %q, want %q", tc.data, err.Error(), tc.errMsg)
+					t.Errorf("Unmarshal(%x) returned error %q, want %q", tc.data, err.Error(), tc.errMsg)
 				}
 			}
 		})

tag.go

@@ -26,22 +26,18 @@ type RawTag struct {
 }
 
 // UnmarshalCBOR sets *t with tag number and raw tag content copied from data.
+//
+// Deprecated: No longer used by this codec; kept for compatibility
+// with user apps that directly call this function.
 func (t *RawTag) UnmarshalCBOR(data []byte) error {
 	if t == nil {
 		return errors.New("cbor.RawTag: UnmarshalCBOR on nil pointer")
 	}
 
-	// Decoding CBOR null and undefined to cbor.RawTag is no-op.
-	if len(data) == 1 && (data[0] == 0xf6 || data[0] == 0xf7) {
-		return nil
-	}
-
 	d := decoder{data: data, dm: defaultDecMode}
 
 	// Check if data is a well-formed CBOR data item.
-	// NOTE: well-formedness check here is redundant when
-	// Unmarshal() invokes RawTag.UnmarshalCBOR().
-	// However, RawTag.UnmarshalCBOR() is exported, so
+	// RawTag.UnmarshalCBOR() is exported, so
 	// the codec needs to support same behavior for:
 	// - Unmarshal(data, *RawTag)
 	// - RawTag.UnmarshalCBOR(data)
@@ -50,8 +46,23 @@ func (t *RawTag) UnmarshalCBOR(data []byte) error {
 		return err
 	}
 
-	// Restore decoder offset after well-formedness check.
-	d.off = 0
+	return t.unmarshalCBOR(data)
+}
+
+// unmarshalCBOR sets *t with tag number and raw tag content copied from data.
+// This function assumes data is well-formed, and does not perform bounds checking.
+// This function is called by Unmarshal().
+func (t *RawTag) unmarshalCBOR(data []byte) error {
+	if t == nil {
+		return errors.New("cbor.RawTag: UnmarshalCBOR on nil pointer")
+	}
+
+	// Decoding CBOR null and undefined to cbor.RawTag is no-op.
+	if len(data) == 1 && (data[0] == 0xf6 || data[0] == 0xf7) {
+		return nil
+	}
+
+	d := decoder{data: data, dm: defaultDecMode}
 
 	// Unmarshal tag number.
 	typ, _, num := d.getHead()

tag_test.go

@@ -1646,16 +1646,16 @@ func TestUnmarshalRawTagOnBadData(t *testing.T) {
 					t.Errorf("UnmarshalCBOR(%x) returned error %q, want %q", tc.data, err.Error(), tc.errMsg)
 				}
 			}
-			// Test Unmarshal(data, *RawTag), which calls RawTag.UnmarshalCBOR() under the hood
+			// Test Unmarshal(data, *RawTag), which calls RawTag.unmarshalCBOR() under the hood
 			{
 				var v RawTag
 
 				err := Unmarshal(tc.data, &v)
 				if err == nil {
-					t.Errorf("UnmarshalCBOR(%x) didn't return error", tc.data)
+					t.Errorf("Unmarshal(%x) didn't return error", tc.data)
 				}
 				if !strings.HasPrefix(err.Error(), tc.errMsg) {
-					t.Errorf("UnmarshalCBOR(%x) returned error %q, want %q", tc.data, err.Error(), tc.errMsg)
+					t.Errorf("Unmarshal(%x) returned error %q, want %q", tc.data, err.Error(), tc.errMsg)
 				}
 			}
 		})