From 344457a5088564c984f9e19ef9a161bf4425bca9 Mon Sep 17 00:00:00 2001
From: Eric Myhre <hash@exultant.us>
Date: Thu, 24 Sep 2020 23:47:37 +0200
Subject: [PATCH] Implement resource budgets in dagcbor parsing.

---
 codec/dagcbor/unmarshal.go | 65 +++++++++++++++++++++++++++++++++++---
 1 file changed, 60 insertions(+), 5 deletions(-)

diff --git a/codec/dagcbor/unmarshal.go b/codec/dagcbor/unmarshal.go
index d22d40f9..324fca29 100644
--- a/codec/dagcbor/unmarshal.go
+++ b/codec/dagcbor/unmarshal.go
@@ -14,7 +14,13 @@ import (
 )
 
 var (
-	ErrInvalidMultibase = errors.New("invalid multibase on IPLD link")
+	ErrInvalidMultibase         = errors.New("invalid multibase on IPLD link")
+	ErrAllocationBudgetExceeded = errors.New("message structure demanded too many resources to process")
+)
+
+const (
+	mapEntryGasScore  = 8
+	listEntryGasScore = 4
 )
 
 // This should be identical to the general feature in the parent package,
@@ -22,6 +28,15 @@ var (
 // which has dag-cbor's special sauce for detecting schemafree links.
 
 func Unmarshal(na ipld.NodeAssembler, tokSrc shared.TokenSource) error {
+	// Have a gas budget, which will be decremented as we allocate memory, and an error returned when execeeded (or about to be exceeded).
+	//  This is a DoS defense mechanism.
+	//  It's *roughly* in units of bytes (but only very, VERY roughly) -- it also treats words as 1 in many cases.
+	// FUTURE: this ought be configurable somehow.  (How, and at what granularity though?)
+	var gas int = 1048576 * 10
+	return unmarshal1(na, tokSrc, &gas)
+}
+
+func unmarshal1(na ipld.NodeAssembler, tokSrc shared.TokenSource, gas *int) error {
 	var tk tok.Token
 	done, err := tokSrc.Step(&tk)
 	if err != nil {
@@ -30,12 +45,12 @@ func Unmarshal(na ipld.NodeAssembler, tokSrc shared.TokenSource) error {
 	if done && !tk.Type.IsValue() {
 		return fmt.Errorf("unexpected eof")
 	}
-	return unmarshal(na, tokSrc, &tk)
+	return unmarshal2(na, tokSrc, &tk, gas)
 }
 
 // starts with the first token already primed.  Necessary to get recursion
 //  to flow right without a peek+unpeek system.
-func unmarshal(na ipld.NodeAssembler, tokSrc shared.TokenSource, tk *tok.Token) error {
+func unmarshal2(na ipld.NodeAssembler, tokSrc shared.TokenSource, tk *tok.Token, gas *int) error {
 	// FUTURE: check for schema.TypedNodeBuilder that's going to parse a Link (they can slurp any token kind they want).
 	switch tk.Type {
 	case tok.TMapOpen:
@@ -44,6 +59,10 @@ func unmarshal(na ipld.NodeAssembler, tokSrc shared.TokenSource, tk *tok.Token)
 		if tk.Length == -1 {
 			expectLen = math.MaxInt32
 			allocLen = 0
+		} else {
+			if *gas-allocLen < 0 { // halt early if this will clearly demand too many resources
+				return ErrAllocationBudgetExceeded
+			}
 		}
 		ma, err := na.BeginMap(allocLen)
 		if err != nil {
@@ -62,6 +81,10 @@ func unmarshal(na ipld.NodeAssembler, tokSrc shared.TokenSource, tk *tok.Token)
 				}
 				return ma.Finish()
 			case tok.TString:
+				*gas -= len(tk.Str) + mapEntryGasScore
+				if *gas < 0 {
+					return ErrAllocationBudgetExceeded
+				}
 				// continue
 			default:
 				return fmt.Errorf("unexpected %s token while expecting map key", tk.Type)
@@ -74,7 +97,7 @@ func unmarshal(na ipld.NodeAssembler, tokSrc shared.TokenSource, tk *tok.Token)
 			if err != nil { // return in error if the key was rejected
 				return err
 			}
-			err = Unmarshal(mva, tokSrc)
+			err = unmarshal1(mva, tokSrc, gas)
 			if err != nil { // return in error if some part of the recursion errored
 				return err
 			}
@@ -87,6 +110,10 @@ func unmarshal(na ipld.NodeAssembler, tokSrc shared.TokenSource, tk *tok.Token)
 		if tk.Length == -1 {
 			expectLen = math.MaxInt32
 			allocLen = 0
+		} else {
+			if *gas-allocLen < 0 { // halt early if this will clearly demand too many resources
+				return ErrAllocationBudgetExceeded
+			}
 		}
 		la, err := na.BeginList(allocLen)
 		if err != nil {
@@ -105,11 +132,15 @@ func unmarshal(na ipld.NodeAssembler, tokSrc shared.TokenSource, tk *tok.Token)
 				}
 				return la.Finish()
 			default:
+				*gas -= listEntryGasScore
+				if *gas < 0 {
+					return ErrAllocationBudgetExceeded
+				}
 				observedLen++
 				if observedLen > expectLen {
 					return fmt.Errorf("unexpected continuation of array elements beyond declared length")
 				}
-				err := unmarshal(la.AssembleValue(), tokSrc, tk)
+				err := unmarshal2(la.AssembleValue(), tokSrc, tk, gas)
 				if err != nil { // return in error if some part of the recursion errored
 					return err
 				}
@@ -120,8 +151,16 @@ func unmarshal(na ipld.NodeAssembler, tokSrc shared.TokenSource, tk *tok.Token)
 	case tok.TNull:
 		return na.AssignNull()
 	case tok.TString:
+		*gas -= len(tk.Str)
+		if *gas < 0 {
+			return ErrAllocationBudgetExceeded
+		}
 		return na.AssignString(tk.Str)
 	case tok.TBytes:
+		*gas -= len(tk.Bytes)
+		if *gas < 0 {
+			return ErrAllocationBudgetExceeded
+		}
 		if !tk.Tagged {
 			return na.AssignBytes(tk.Bytes)
 		}
@@ -139,12 +178,28 @@ func unmarshal(na ipld.NodeAssembler, tokSrc shared.TokenSource, tk *tok.Token)
 			return fmt.Errorf("unhandled cbor tag %d", tk.Tag)
 		}
 	case tok.TBool:
+		*gas -= 1
+		if *gas < 0 {
+			return ErrAllocationBudgetExceeded
+		}
 		return na.AssignBool(tk.Bool)
 	case tok.TInt:
+		*gas -= 1
+		if *gas < 0 {
+			return ErrAllocationBudgetExceeded
+		}
 		return na.AssignInt(int(tk.Int)) // FIXME overflow check
 	case tok.TUint:
+		*gas -= 1
+		if *gas < 0 {
+			return ErrAllocationBudgetExceeded
+		}
 		return na.AssignInt(int(tk.Uint)) // FIXME overflow check
 	case tok.TFloat64:
+		*gas -= 1
+		if *gas < 0 {
+			return ErrAllocationBudgetExceeded
+		}
 		return na.AssignFloat(tk.Float64)
 	default:
 		panic("unreachable")