From be378e682fa8d546521712ebd4f5e84d93905cf3 Mon Sep 17 00:00:00 2001
From: Hiroki Yorimitsu <52403055+hyorimitsu@users.noreply.github.com>
Date: Wed, 6 Mar 2024 09:29:57 +0900
Subject: [PATCH] Add support for math/rand/v2 added in Go 1.22

---
 import_tracker.go         | 10 +++++
 rules/rand.go             | 24 ++++++-----
 testutils/g404_samples.go | 85 +++++++++++++++++++++++++++++++++++++++
 3 files changed, 109 insertions(+), 10 deletions(-)

diff --git a/import_tracker.go b/import_tracker.go
index 7984e99..cf6c392 100644
--- a/import_tracker.go
+++ b/import_tracker.go
@@ -15,9 +15,14 @@ package gosec
 import (
 	"go/ast"
 	"go/types"
+	"regexp"
 	"strings"
 )
 
+var (
+	versioningPackagePattern = regexp.MustCompile(`v[0-9]+$`)
+)
+
 // ImportTracker is used to normalize the packages that have been imported
 // by a source file. It is able to differentiate between plain imports, aliased
 // imports and init only imports.
@@ -66,5 +71,10 @@ func importName(importPath string) string {
 	if len(parts) > 0 {
 		name = parts[len(parts)-1]
 	}
+	// If the last segment of the path is version information, consider the second to last segment as the package name.
+	// (e.g., `math/rand/v2` would be `rand`)
+	if len(parts) > 1 && versioningPackagePattern.MatchString(name) {
+		name = parts[len(parts)-2]
+	}
 	return name
 }
diff --git a/rules/rand.go b/rules/rand.go
index 4491fd9..fe34ca9 100644
--- a/rules/rand.go
+++ b/rules/rand.go
@@ -23,8 +23,7 @@ import (
 
 type weakRand struct {
 	issue.MetaData
-	funcNames   []string
-	packagePath string
+	blocklist map[string][]string
 }
 
 func (w *weakRand) ID() string {
@@ -32,8 +31,8 @@ func (w *weakRand) ID() string {
 }
 
 func (w *weakRand) Match(n ast.Node, c *gosec.Context) (*issue.Issue, error) {
-	for _, funcName := range w.funcNames {
-		if _, matched := gosec.MatchCallByPackage(n, c, w.packagePath, funcName); matched {
+	for pkg, funcs := range w.blocklist {
+		if _, matched := gosec.MatchCallByPackage(n, c, pkg, funcs...); matched {
 			return c.NewIssue(n, w.ID(), w.What, w.Severity, w.Confidence), nil
 		}
 	}
@@ -43,17 +42,22 @@ func (w *weakRand) Match(n ast.Node, c *gosec.Context) (*issue.Issue, error) {
 
 // NewWeakRandCheck detects the use of random number generator that isn't cryptographically secure
 func NewWeakRandCheck(id string, _ gosec.Config) (gosec.Rule, []ast.Node) {
+	calls := make(map[string][]string)
+	calls["math/rand"] = []string{
+		"New", "Read", "Float32", "Float64", "Int", "Int31", "Int31n",
+		"Int63", "Int63n", "Intn", "NormFloat64", "Uint32", "Uint64",
+	}
+	calls["math/rand/v2"] = []string{
+		"New", "Float32", "Float64", "Int", "Int32", "Int32N",
+		"Int64", "Int64N", "IntN", "N", "NormFloat64", "Uint32", "Uint32N", "Uint64", "Uint64N", "UintN",
+	}
 	return &weakRand{
-		funcNames: []string{
-			"New", "Read", "Float32", "Float64", "Int", "Int31",
-			"Int31n", "Int63", "Int63n", "Intn", "NormalFloat64", "Uint32", "Uint64",
-		},
-		packagePath: "math/rand",
+		blocklist: calls,
 		MetaData: issue.MetaData{
 			ID:         id,
 			Severity:   issue.High,
 			Confidence: issue.Medium,
-			What:       "Use of weak random number generator (math/rand instead of crypto/rand)",
+			What:       "Use of weak random number generator (math/rand or math/rand/v2 instead of crypto/rand)",
 		},
 	}, []ast.Node{(*ast.CallExpr)(nil)}
 }
diff --git a/testutils/g404_samples.go b/testutils/g404_samples.go
index cc8c9b8..72e1034 100644
--- a/testutils/g404_samples.go
+++ b/testutils/g404_samples.go
@@ -27,6 +27,16 @@ func main() {
 	{[]string{`
 package main
 
+import "math/rand/v2"
+
+func main() {
+	bad := rand.Int()
+	println(bad)
+}
+`}, 1, gosec.NewConfig()},
+	{[]string{`
+package main
+
 import (
 	"crypto/rand"
 	mrand "math/rand"
@@ -42,6 +52,21 @@ func main() {
 	{[]string{`
 package main
 
+import (
+	"crypto/rand"
+	mrand "math/rand/v2"
+)
+
+func main() {
+	good, _ := rand.Read(nil)
+	println(good)
+	bad := mrand.Int32()
+	println(bad)
+}
+`}, 1, gosec.NewConfig()},
+	{[]string{`
+package main
+
 import (
 	"math/rand"
 )
@@ -55,6 +80,19 @@ func main() {
 	{[]string{`
 package main
 
+import (
+	"math/rand/v2"
+)
+
+func main() {
+	gen := rand.New(rand.NewPCG(1, 2))
+	bad := gen.Int()
+	println(bad)
+}
+`}, 1, gosec.NewConfig()},
+	{[]string{`
+package main
+
 import (
 	"math/rand"
 )
@@ -67,6 +105,18 @@ func main() {
 	{[]string{`
 package main
 
+import (
+	"math/rand/v2"
+)
+
+func main() {
+	bad := rand.IntN(10)
+	println(bad)
+}
+`}, 1, gosec.NewConfig()},
+	{[]string{`
+package main
+
 import (
 	"crypto/rand"
 	"math/big"
@@ -83,6 +133,22 @@ func main() {
 	{[]string{`
 package main
 
+import (
+	"crypto/rand"
+	"math/big"
+	rnd "math/rand/v2"
+)
+
+func main() {
+	good, _ := rand.Int(rand.Reader, big.NewInt(int64(2)))
+	println(good)
+	bad := rnd.IntN(2)
+	println(bad)
+}
+`}, 1, gosec.NewConfig()},
+	{[]string{`
+package main
+
 import (
 	crand "crypto/rand"
 	"math/big"
@@ -98,5 +164,24 @@ func main() {
 	_ = rand2.Intn(2)  // bad
 	_ = rand3.Intn(2)  // bad
 }
+`}, 3, gosec.NewConfig()},
+	{[]string{`
+package main
+
+import (
+	crand "crypto/rand"
+	"math/big"
+	"math/rand/v2"
+	rand2 "math/rand/v2"
+	rand3 "math/rand/v2"
+)
+
+func main() {
+	_, _ = crand.Int(crand.Reader, big.NewInt(int64(2))) // good
+
+	_ = rand.IntN(2) // bad
+	_ = rand2.IntN(2)  // bad
+	_ = rand3.IntN(2)  // bad
+}
 `}, 3, gosec.NewConfig()},
 }