From d5403f7fe9cf2d2919bf5f8674389e596fbbf9d7 Mon Sep 17 00:00:00 2001
From: Jim Kalafut <jim@kalafut.net>
Date: Wed, 15 Aug 2018 13:16:17 -0700
Subject: [PATCH] Add check of truncated length (#5109)

---
 helper/base62/base62.go      | 26 +++++++++++++++++---------
 helper/base62/base62_test.go | 13 +++++++------
 2 files changed, 24 insertions(+), 15 deletions(-)

diff --git a/helper/base62/base62.go b/helper/base62/base62.go
index c5d2accede..56890b0c45 100644
--- a/helper/base62/base62.go
+++ b/helper/base62/base62.go
@@ -9,6 +9,9 @@ import (
 )
 
 // Encode converts buf into a base62 string
+//
+// Note: this should only be used for reducing a string's character set range.
+// It is not for use with arbitrary data since leading 0 bytes will be dropped.
 func Encode(buf []byte) string {
 	var encoder big.Int
 
@@ -29,15 +32,20 @@ func Decode(input string) []byte {
 // If truncate is true, the result will be a string of the requested length.
 // Otherwise, it will be the encoded result of length bytes of random data.
 func Random(length int, truncate bool) (string, error) {
-	buf, err := uuid.GenerateRandomBytes(length)
-	if err != nil {
-		return "", err
-	}
+	for {
+		buf, err := uuid.GenerateRandomBytes(length)
+		if err != nil {
+			return "", err
+		}
 
-	result := Encode(buf)
-	if truncate {
-		result = result[:length]
-	}
+		result := Encode(buf)
 
-	return result, nil
+		if truncate {
+			if len(result) < length {
+				continue
+			}
+			result = result[:length]
+		}
+		return result, nil
+	}
 }
diff --git a/helper/base62/base62_test.go b/helper/base62/base62_test.go
index d5e5744f01..19b5f8b053 100644
--- a/helper/base62/base62_test.go
+++ b/helper/base62/base62_test.go
@@ -96,14 +96,15 @@ func TestRandom(t *testing.T) {
 		t.Fatalf("Expected different random values. Got duplicate: %s", a)
 	}
 
-	c, _ := Random(4738, true)
-	if len(c) != 4738 {
-		t.Fatalf("Expected length 4738, got: %d", len(c))
+	for i := 0; i < 3000; i++ {
+		c, _ := Random(i, true)
+		if len(c) != i {
+			t.Fatalf("Expected length %d, got: %d", i, len(c))
+		}
 	}
 
 	d, _ := Random(100, false)
-	if len(d) < 134 || len(d) > 135 {
-		t.Fatalf("Expected length 134 or 135, got: %d", len(d))
+	if len(d) < 133 || len(d) > 135 {
+		t.Fatalf("Expected length 133-135, got: %d", len(d))
 	}
-
 }