From ad12fa30fe8ee0e2673a1305a64bf39786919521 Mon Sep 17 00:00:00 2001
From: Jeromy <why@ipfs.io>
Date: Wed, 4 May 2016 13:26:50 -0700
Subject: [PATCH] respond to CR feedback

---
 codec.go          |  5 +++++
 multiaddr.go      | 22 ++++++++++++++++++++--
 multiaddr_test.go | 43 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 68 insertions(+), 2 deletions(-)

diff --git a/codec.go b/codec.go
index 747f160..554ed4e 100644
--- a/codec.go
+++ b/codec.go
@@ -57,6 +57,10 @@ func stringToBytes(s string) ([]byte, error) {
 func validateBytes(b []byte) (err error) {
 	for len(b) > 0 {
 		code, n, err := ReadVarintCode(b)
+		if err != nil {
+			return err
+		}
+
 		b = b[n:]
 		p := ProtocolWithCode(code)
 		if p.Code == 0 {
@@ -81,6 +85,7 @@ func validateBytes(b []byte) (err error) {
 
 	return nil
 }
+
 func bytesToString(b []byte) (ret string, err error) {
 	s := ""
 
diff --git a/multiaddr.go b/multiaddr.go
index be04a7f..430cb1d 100644
--- a/multiaddr.go
+++ b/multiaddr.go
@@ -3,6 +3,7 @@ package multiaddr
 import (
 	"bytes"
 	"fmt"
+	"log"
 	"strings"
 )
 
@@ -12,7 +13,13 @@ type multiaddr struct {
 }
 
 // NewMultiaddr parses and validates an input string, returning a *Multiaddr
-func NewMultiaddr(s string) (Multiaddr, error) {
+func NewMultiaddr(s string) (a Multiaddr, err error) {
+	defer func() {
+		if e := recover(); e != nil {
+			log.Printf("Panic in NewMultiaddr on input %q: %s", s, e)
+			err = fmt.Errorf("%v", e)
+		}
+	}()
 	b, err := stringToBytes(s)
 	if err != nil {
 		return nil, err
@@ -22,7 +29,14 @@ func NewMultiaddr(s string) (Multiaddr, error) {
 
 // NewMultiaddrBytes initializes a Multiaddr from a byte representation.
 // It validates it as an input string.
-func NewMultiaddrBytes(b []byte) (Multiaddr, error) {
+func NewMultiaddrBytes(b []byte) (a Multiaddr, err error) {
+	defer func() {
+		if e := recover(); e != nil {
+			log.Printf("Panic in NewMultiaddrBytes on input %q: %s", b, e)
+			err = fmt.Errorf("%v", e)
+		}
+	}()
+
 	if err := validateBytes(b); err != nil {
 		return nil, err
 	}
@@ -82,6 +96,10 @@ func (m *multiaddr) Protocols() []Protocol {
 		b = b[n:]
 
 		size, err := sizeForAddr(p, b)
+		if err != nil {
+			panic(err)
+		}
+
 		b = b[size:]
 	}
 	return ps
diff --git a/multiaddr_test.go b/multiaddr_test.go
index ac81ac2..d469451 100644
--- a/multiaddr_test.go
+++ b/multiaddr_test.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/hex"
 	"math/rand"
+	"strings"
 	"testing"
 	"time"
 )
@@ -140,6 +141,10 @@ func TestStringToBytes(t *testing.T) {
 		if !bytes.Equal(b1, b2) {
 			t.Error("failed to convert", s, "to", b1, "got", b2)
 		}
+
+		if err := validateBytes(b2); err != nil {
+			t.Error(err)
+		}
 	}
 
 	testString("/ip4/127.0.0.1/udp/1234", "047f0000011104d2")
@@ -155,6 +160,10 @@ func TestBytesToString(t *testing.T) {
 			t.Error("failed to decode hex", h)
 		}
 
+		if err := validateBytes(b); err != nil {
+			t.Error(err)
+		}
+
 		s2, err := bytesToString(b)
 		if err != nil {
 			t.Error("failed to convert", b)
@@ -362,3 +371,37 @@ func TestFuzzBytes(t *testing.T) {
 		}
 	}
 }
+
+func randMaddrString() string {
+	good_corpus := []string{"tcp", "ip", "udp", "ipfs", "0.0.0.0", "127.0.0.1", "12345", "QmbHVEEepCi7rn7VL7Exxpd2Ci9NNB6ifvqwhsrbRMgQFP"}
+
+	size := rand.Intn(256)
+	parts := make([]string, 0, size)
+	for i := 0; i < size; i++ {
+		switch rand.Intn(5) {
+		case 0, 1, 2:
+			parts = append(parts, good_corpus[rand.Intn(len(good_corpus))])
+		default:
+			badbuf := make([]byte, rand.Intn(256))
+			rand.Read(badbuf)
+			parts = append(parts, string(badbuf))
+		}
+	}
+
+	return "/" + strings.Join(parts, "/")
+}
+
+func TestFuzzString(t *testing.T) {
+	rand.Seed(time.Now().UnixNano())
+	// Bump up these numbers if you want to stress this
+	for i := 0; i < 2000; i++ {
+
+		// just checking that it doesnt panic
+		ma, err := NewMultiaddr(randMaddrString())
+		if err == nil {
+			// for any valid multiaddrs, make sure these calls don't panic
+			ma.String()
+			ma.Protocols()
+		}
+	}
+}