Improve error checking on number parsing

lexer.go

@@ -184,6 +184,8 @@ func (l *tomlLexer) lexRvalue() tomlLexStateFn {
 				return l.lexVoid
 			}
 			return l.lexRvalue
+		case '_':
+			return l.errorf("cannot start number with underscore")
 		}
 
 		if l.follow("true") {
@@ -550,7 +552,6 @@ func (l *tomlLexer) lexNumber() tomlLexStateFn {
 		} else if isDigit(next) {
 			digitSeen = true
 		} else if next == '_' {
-			l.pos++
 		} else {
 			l.backup()
 			break

lexer_test.go

@@ -618,3 +618,12 @@ func TestKeyNewline(t *testing.T) {
 		token{Position{1, 1}, tokenError, "keys cannot contain new lines"},
 	})
 }
+
+func TestInvalidFloat(t *testing.T) {
+	testFlow(t, "a=7e1_", []token{
+		token{Position{1, 1}, tokenKey, "a"},
+		token{Position{1, 2}, tokenEqual, "="},
+		token{Position{1, 3}, tokenFloat, "7e1_"},
+		token{Position{1, 7}, tokenEOF, ""},
+	})
+}

parser.go

@@ -5,6 +5,7 @@ package toml
 import (
 	"fmt"
 	"reflect"
+	"regexp"
 	"strconv"
 	"strings"
 	"time"
@@ -211,6 +212,16 @@ func (p *tomlParser) parseAssign() tomlParserStateFn {
 	return p.parseStart
 }
 
+var numberUnderscoreInvalidRegexp *regexp.Regexp
+
+func cleanupNumberToken(value string) (string, error) {
+	if numberUnderscoreInvalidRegexp.MatchString(value) {
+		return "", fmt.Errorf("invalid use of _ in number")
+	}
+	cleanedVal := strings.Replace(value, "_", "", -1)
+	return cleanedVal, nil
+}
+
 func (p *tomlParser) parseRvalue() interface{} {
 	tok := p.getToken()
 	if tok == nil || tok.typ == tokenEOF {
@@ -225,14 +236,20 @@ func (p *tomlParser) parseRvalue() interface{} {
 	case tokenFalse:
 		return false
 	case tokenInteger:
-		cleanedVal := strings.Replace(tok.val, "_", "", -1)
+		cleanedVal, err := cleanupNumberToken(tok.val)
+		if err != nil {
+			p.raiseError(tok, "%s", err)
+		}
 		val, err := strconv.ParseInt(cleanedVal, 10, 64)
 		if err != nil {
 			p.raiseError(tok, "%s", err)
 		}
 		return val
 	case tokenFloat:
-		cleanedVal := strings.Replace(tok.val, "_", "", -1)
+		cleanedVal, err := cleanupNumberToken(tok.val)
+		if err != nil {
+			p.raiseError(tok, "%s", err)
+		}
 		val, err := strconv.ParseFloat(cleanedVal, 64)
 		if err != nil {
 			p.raiseError(tok, "%s", err)
@@ -361,3 +378,7 @@ func parseToml(flow chan token) *TomlTree {
 	parser.run()
 	return result
 }
+
+func init() {
+	numberUnderscoreInvalidRegexp = regexp.MustCompile(`([^\d]_|_[^\d]|_$|^_)`)
+}

parser_test.go

@@ -633,3 +633,25 @@ func TestGroupArrayReassign(t *testing.T) {
 		t.Error("Bad error message:", err.Error())
 	}
 }
+
+func TestInvalidFloatParsing(t *testing.T) {
+	_, err := Load("a=1e_2")
+	if err.Error() != "(1, 3): invalid use of _ in number" {
+		t.Error("Bad error message:", err.Error())
+	}
+
+	_, err = Load("a=1e2_")
+	if err.Error() != "(1, 3): invalid use of _ in number" {
+		t.Error("Bad error message:", err.Error())
+	}
+
+	_, err = Load("a=1__2")
+	if err.Error() != "(1, 3): invalid use of _ in number" {
+		t.Error("Bad error message:", err.Error())
+	}
+
+	_, err = Load("a=_1_2")
+	if err.Error() != "(1, 3): cannot start number with underscore" {
+		t.Error("Bad error message:", err.Error())
+	}
+}