From c99b6451948f8fc4443136f75afe4d502445cd41 Mon Sep 17 00:00:00 2001
From: Stephan Renatus <stephan.renatus@gmail.com>
Date: Mon, 8 Nov 2021 14:10:32 +0100
Subject: [PATCH] ast: fuzzed parser bug, two codeql issues (#3988)

* storage/path.Ref: parse int64 into ast.Number
* ast/PtrRef: guard against giant paths

I don't believe this limit is every going to be reached. But CodeQL had flagged
this, and it's not entirely wrong. Let's error our on giant wonky inputs instead
of seeing what'll happen with it eventually.

* ast/parser: fix bad import alias var

The fuzzer came up with that!

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>
---
 ast/parser.go        | 15 ++++++++-------
 ast/parser_test.go   |  1 +
 ast/term.go          |  4 ++++
 storage/path.go      |  4 ++--
 storage/path_test.go |  2 ++
 5 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/ast/parser.go b/ast/parser.go
index 076d880af9..0990a6a539 100644
--- a/ast/parser.go
+++ b/ast/parser.go
@@ -472,14 +472,15 @@ func (p *Parser) parseImport() *Import {
 			return nil
 		}
 
-		alias := p.parseTerm()
-
-		v, ok := alias.Value.(Var)
-		if !ok {
-			p.illegal("expected var")
-			return nil
+		if alias := p.parseTerm(); alias != nil {
+			v, ok := alias.Value.(Var)
+			if ok {
+				imp.Alias = v
+				return &imp
+			}
 		}
-		imp.Alias = v
+		p.illegal("expected var")
+		return nil
 	}
 
 	return &imp
diff --git a/ast/parser_test.go b/ast/parser_test.go
index 7960d6af3e..56fc775427 100644
--- a/ast/parser_test.go
+++ b/ast/parser_test.go
@@ -1091,6 +1091,7 @@ func TestImport(t *testing.T) {
 	assertParseErrorContains(t, "non-ground ref", "import data.foo[x]", "rego_parse_error: unexpected var token: expecting string")
 	assertParseErrorContains(t, "non-string", "import input.foo[0]", "rego_parse_error: unexpected number token: expecting string")
 	assertParseErrorContains(t, "unknown root", "import foo.bar", "rego_parse_error: unexpected import path, must begin with one of: {data, future, input}, got: foo")
+	assertParseErrorContains(t, "bad variable term", "import input as A(", "rego_parse_error: unexpected eof token: expected var")
 
 	_, _, err := ParseStatements("", "package foo\nimport bar.data\ndefault foo=1")
 	if err == nil {
diff --git a/ast/term.go b/ast/term.go
index e9b1576236..f3937a32c7 100644
--- a/ast/term.go
+++ b/ast/term.go
@@ -10,6 +10,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"math"
 	"math/big"
 	"net/url"
 	"regexp"
@@ -843,6 +844,9 @@ func PtrRef(head *Term, s string) (Ref, error) {
 		return Ref{head}, nil
 	}
 	parts := strings.Split(s, "/")
+	if max := math.MaxInt32; len(parts) >= max {
+		return nil, fmt.Errorf("path too long: %s, %d > %d (max)", s, len(parts), max)
+	}
 	ref := make(Ref, uint(len(parts))+1)
 	ref[0] = head
 	for i := 0; i < len(parts); i++ {
diff --git a/storage/path.go b/storage/path.go
index c3d32f2e66..02ef4cab40 100644
--- a/storage/path.go
+++ b/storage/path.go
@@ -125,9 +125,9 @@ func (p Path) Ref(head *ast.Term) (ref ast.Ref) {
 	ref = make(ast.Ref, len(p)+1)
 	ref[0] = head
 	for i := range p {
-		idx, err := strconv.ParseInt(p[i], 10, 32)
+		idx, err := strconv.ParseInt(p[i], 10, 64)
 		if err == nil {
-			ref[i+1] = ast.IntNumberTerm(int(idx))
+			ref[i+1] = ast.UIntNumberTerm(uint64(idx))
 		} else {
 			ref[i+1] = ast.StringTerm(p[i])
 		}
diff --git a/storage/path_test.go b/storage/path_test.go
index e2bee3daea..0a7729c370 100644
--- a/storage/path_test.go
+++ b/storage/path_test.go
@@ -5,6 +5,7 @@
 package storage
 
 import (
+	"math"
 	"reflect"
 	"testing"
 
@@ -179,6 +180,7 @@ func TestPathRef(t *testing.T) {
 		{"/", "data", "data"},
 		{"/foo/bar", "data", "data.foo.bar"},
 		{"/foo/bar/3", "data", "data.foo.bar[3]"},
+		{fmt.Sprintf("/foo/bar/%d", math.MaxInt64), "data", fmt.Sprintf("data.foo.bar[%d]", math.MaxInt64)},
 	}
 	for _, tc := range tests {
 		path := MustParsePath(tc.path)