From 75b5f2f41263b3ed6e5168c1be2b85818d5d4a81 Mon Sep 17 00:00:00 2001
From: Tobie Morgan Hitchcock <tobie@abcum.com>
Date: Tue, 22 Nov 2016 19:10:42 +0000
Subject: [PATCH] Improve SQL DEFINE TOKEN statements

---
 sql/exprs.go  | 44 +++++++++++++++++++++++++++++++++++++++++++-
 sql/token.go  | 37 +++++++++++++++++++++++++++++++++++--
 sql/tokens.go |  2 ++
 3 files changed, 80 insertions(+), 3 deletions(-)

diff --git a/sql/exprs.go b/sql/exprs.go
index d8ba4c3d..7bc43785 100644
--- a/sql/exprs.go
+++ b/sql/exprs.go
@@ -240,7 +240,7 @@ func (p *parser) parseString() (string, error) {
 
 func (p *parser) parseRegion() (string, error) {
 
-	tok, lit, err := p.shouldBe(IDENT, STRING, REGION)
+	tok, lit, err := p.shouldBe(STRING, REGION)
 	if err != nil {
 		return string(""), &ParseError{Found: lit, Expected: []string{"string"}}
 	}
@@ -251,6 +251,19 @@ func (p *parser) parseRegion() (string, error) {
 
 }
 
+func (p *parser) parseBinary() ([]byte, error) {
+
+	tok, lit, err := p.shouldBe(STRING, REGION)
+	if err != nil {
+		return nil, &ParseError{Found: lit, Expected: []string{"string"}}
+	}
+
+	val, err := p.declare(tok, lit)
+
+	return []byte(val.(string)), err
+
+}
+
 func (p *parser) parseScript() (string, error) {
 
 	tok, lit, err := p.shouldBe(STRING, REGION)
@@ -319,6 +332,35 @@ func (p *parser) parseBcrypt() ([]byte, error) {
 
 }
 
+func (p *parser) parseAlgorithm() (string, error) {
+
+	expected := []string{
+		"ES256", "ES384", "ES512",
+		"HS256", "HS384", "HS512",
+		"PS256", "PS384", "PS512",
+		"RS256", "RS384", "RS512",
+	}
+
+	_, lit, err := p.shouldBe(IDENT, STRING)
+	if err != nil {
+		return string(""), &ParseError{Found: lit, Expected: expected}
+	}
+
+	switch lit {
+	case "ES256", "ES384", "ES512":
+	case "HS256", "HS384", "HS512":
+	case "PS256", "PS384", "PS512":
+	case "RS256", "RS384", "RS512":
+	default:
+		return string(""), &ParseError{Found: lit, Expected: expected}
+	}
+
+	val, err := p.declare(STRING, lit)
+
+	return val.(string), err
+
+}
+
 func (p *parser) parseExpr() (exp Expr, err error) {
 
 	// Create the root binary expression tree.
diff --git a/sql/token.go b/sql/token.go
index 7496ae57..0c5d0022 100644
--- a/sql/token.go
+++ b/sql/token.go
@@ -18,6 +18,10 @@ func (p *parser) parseDefineTokenStatement() (stmt *DefineTokenStatement, err er
 
 	stmt = &DefineTokenStatement{}
 
+	if stmt.Name, err = p.parseName(); err != nil {
+		return nil, err
+	}
+
 	if _, _, err = p.shouldBe(ON); err != nil {
 		return nil, err
 	}
@@ -44,8 +48,33 @@ func (p *parser) parseDefineTokenStatement() (stmt *DefineTokenStatement, err er
 		}
 	}
 
-	if stmt.Code, err = p.parseRand(); err != nil {
-		return nil, err
+	for {
+
+		tok, _, exi := p.mightBe(TYPE, VALUE)
+		if !exi {
+			break
+		}
+
+		if p.is(tok, TYPE) {
+			if stmt.Type, err = p.parseAlgorithm(); err != nil {
+				return nil, err
+			}
+		}
+
+		if p.is(tok, VALUE) {
+			if stmt.Code, err = p.parseBinary(); err != nil {
+				return nil, err
+			}
+		}
+
+	}
+
+	if stmt.Type == "" {
+		return nil, &ParseError{Found: ";", Expected: []string{"TYPE"}}
+	}
+
+	if stmt.Code == nil {
+		return nil, &ParseError{Found: ";", Expected: []string{"VALUE"}}
 	}
 
 	if _, _, err = p.shouldBe(EOF, SEMICOLON); err != nil {
@@ -60,6 +89,10 @@ func (p *parser) parseRemoveTokenStatement() (stmt *RemoveTokenStatement, err er
 
 	stmt = &RemoveTokenStatement{}
 
+	if stmt.Name, err = p.parseName(); err != nil {
+		return nil, err
+	}
+
 	if _, _, err = p.shouldBe(ON); err != nil {
 		return nil, err
 	}
diff --git a/sql/tokens.go b/sql/tokens.go
index 4bb198ee..dff90b27 100644
--- a/sql/tokens.go
+++ b/sql/tokens.go
@@ -184,6 +184,7 @@ const (
 	UPSERT
 	USE
 	VALIDATE
+	VALUE
 	VERSION
 	VIEW
 	VOID
@@ -346,6 +347,7 @@ var tokens = [...]string{
 	UPSERT:          "UPSERT",
 	USE:             "USE",
 	VALIDATE:        "VALIDATE",
+	VALUE:           "VALUE",
 	VERSION:         "VERSION",
 	VIEW:            "VIEW",
 	VOID:            "VOID",