Repository: ysugimoto/falco
Pull request #57 (2022-04-24 20:39:06)

feat(lint): Adds support for subs with return type
Extends falco to support subroutines with return values

Fixes https://github.com/ysugimoto/falco/issues/8

Signed-off-by: Sotiris Nanopoulos <sotirisnan@gmail.com>

davinci26: It might have some rough edges left and right, will compare it tomorrow with our VCL corpus which is fairly big so it can expose bugs
davinci26: darn linter, I will fix this tomorrow
davinci26: I tried this PR on a fairly large corpus of VCL and it went ok. @ysugimoto wanna take a look? Honestly not 100% that I got the abstraction/code patterns all right but at least it works
ysugimoto: @davinci26 Thank you for your PR! I'll take a look in a few days 😃 
davinci26: I have been travelling the past few days, I will address your comments today/tomorrow
davinci26: Regarding the big comment:

```
sub get_int INTEGER {
  declare local var.tmp INTEGER;
  set var.tmp = 10;
  return var.tmp  // invalid, could not specify variable
}
```

This is missing a semicolon but otherwise it should compile, https://fiddle.fastlydemo.net/fiddle/11869fd6

---
```
sub get_int INTEGER {
  return (100); // invalid, literal could not enclose in parenthesis
}
```

Already covered in https://github.com/ysugimoto/falco/blob/be9da0c50e8a8d6fa9ae029559abffeb5465a018/linter/linter_test.go#L1287

---
```
sub get_str STRING {
  declare local var.tmp STRING;
  set var.tmp = "foo";
  return (var.tmp "bar"); // invalid, return value must be a single literal (even enclosed in parenthesis or not)
}
```

This was a bit tricky to get it right, I implemented it in https://github.com/ysugimoto/falco/pull/57/commits/a80d30a36d6e618a3ac8bfa1752bca6b0b7b6254

---


Patch:
diff --git a/.gitignore b/.gitignore
index b75848c..14eb5f3 100644
--- a/.gitignore
+++ b/.gitignore
@@ -16,3 +16,5 @@
 falco
 dist/*
 playground
+
+.vscode
diff --git a/__generator__/builtin.yml b/__generator__/builtin.yml
index 7016cee..8b2b4b6 100644
--- a/__generator__/builtin.yml
+++ b/__generator__/builtin.yml
@@ -1221,14 +1221,14 @@ ratelimit.check_rate:
   reference: "https://developer.fastly.com/reference/vcl/functions/rate-limiting/ratelimit-check-rate/"
   on: [RECV, HASH, HIT, MISS, PASS, FETCH, ERROR, DELIVER, LOG]
   arguments:
-    - [STRING, ID, INTEGER, INTEGER, INTEGER, ID, TIME]
+    - [STRING, ID, INTEGER, INTEGER, INTEGER, ID, RTIME]
   return: BOOL
 
 ratelimit.check_rates:
   reference: "https://developer.fastly.com/reference/vcl/functions/rate-limiting/ratelimit-check-rates/"
   on: [RECV, HASH, HIT, MISS, PASS, FETCH, ERROR, DELIVER, LOG]
   arguments:
-    - [STRING, ID, INTEGER, INTEGER, INTEGER, ID, INTEGER, INTEGER, INTEGER, ID, TIME]
+    - [STRING, ID, INTEGER, INTEGER, INTEGER, ID, INTEGER, INTEGER, INTEGER, ID, RTIME]
   return: BOOL
 
 ratelimit.penaltybox_add:
diff --git a/__generator__/predefined.yml b/__generator__/predefined.yml
index 6038398..02f912d 100644
--- a/__generator__/predefined.yml
+++ b/__generator__/predefined.yml
@@ -2093,3 +2093,48 @@ LF:
   reference: "https://developer.fastly.com/reference/vcl/types/string/"
   on: [RECV, HASH, HIT, MISS, PASS, FETCH, ERROR, DELIVER, LOG]
   get: STRING
+
+ratecounter.%any%.bucket.10s:
+  reference: "https://developer.fastly.com/reference/vcl/variables/rate-limiting/ratecounter-bucket-10s/"
+  on: [RECV, HASH, HIT, MISS, PASS, FETCH, ERROR, DELIVER, LOG]
+  get: INTEGER
+
+ratecounter.%any%.bucket.20s:
+  reference: "https://developer.fastly.com/reference/vcl/variables/rate-limiting/ratecounter-bucket-20s/"
+  on: [RECV, HASH, HIT, MISS, PASS, FETCH, ERROR, DELIVER, LOG]
+  get: INTEGER
+
+ratecounter.%any%.bucket.30s:
+  reference: "https://developer.fastly.com/reference/vcl/variables/rate-limiting/ratecounter-bucket-30s/"
+  on: [RECV, HASH, HIT, MISS, PASS, FETCH, ERROR, DELIVER, LOG]
+  get: INTEGER
+
+ratecounter.%any%.bucket.40s:
+  reference: "https://developer.fastly.com/reference/vcl/variables/rate-limiting/ratecounter-bucket-40s/"
+  on: [RECV, HASH, HIT, MISS, PASS, FETCH, ERROR, DELIVER, LOG]
+  get: INTEGER
+
+ratecounter.%any%.bucket.50s:
+  reference: "https://developer.fastly.com/reference/vcl/variables/rate-limiting/ratecounter-bucket-50s/"
+  on: [RECV, HASH, HIT, MISS, PASS, FETCH, ERROR, DELIVER, LOG]
+  get: INTEGER
+
+ratecounter.%any%.bucket.60s:
+  reference: "https://developer.fastly.com/reference/vcl/variables/rate-limiting/ratecounter-bucket-60s/"
+  on: [RECV, HASH, HIT, MISS, PASS, FETCH, ERROR, DELIVER, LOG]
+  get: INTEGER
+
+ratecounter.%any%.rate.1s:
+  reference: "https://developer.fastly.com/reference/vcl/variables/rate-limiting/ratecounter-rate-1s/"
+  on: [RECV, HASH, HIT, MISS, PASS, FETCH, ERROR, DELIVER, LOG]
+  get: FLOAT
+
+ratecounter.%any%.rate.10s:
+  reference: "https://developer.fastly.com/reference/vcl/variables/rate-limiting/ratecounter-rate-10s/"
+  on: [RECV, HASH, HIT, MISS, PASS, FETCH, ERROR, DELIVER, LOG]
+  get: FLOAT
+
+ratecounter.%any%.rate.60s:
+  reference: "https://developer.fastly.com/reference/vcl/variables/rate-limiting/ratecounter-rate-60s/"
+  on: [RECV, HASH, HIT, MISS, PASS, FETCH, ERROR, DELIVER, LOG]
+  get: FLOAT
diff --git a/ast/return_statement.go b/ast/return_statement.go
index b3b3374..44453b3 100644
--- a/ast/return_statement.go
+++ b/ast/return_statement.go
@@ -6,7 +6,8 @@ import (
 
 type ReturnStatement struct {
 	*Meta
-	Ident *Ident
+	ReturnExpression *Expression
+	HasParenthesis   bool
 }
 
 func (r *ReturnStatement) statement()     {}
@@ -15,7 +16,11 @@ func (r *ReturnStatement) String() string {
 	var buf bytes.Buffer
 
 	buf.WriteString(r.LeadingComment())
-	buf.WriteString(indent(r.Nest) + "return(" + r.Ident.String() + ");")
+	if r.ReturnExpression != nil {
+		buf.WriteString(indent(r.Nest) + "return(" + (*r.ReturnExpression).String() + ");")
+	} else {
+		buf.WriteString(indent(r.Nest) + "return;")
+	}
 	buf.WriteString(r.TrailingComment())
 	buf.WriteString("\n")
 
diff --git a/ast/subroutine_declaration.go b/ast/subroutine_declaration.go
index 0b89614..907c031 100644
--- a/ast/subroutine_declaration.go
+++ b/ast/subroutine_declaration.go
@@ -6,8 +6,9 @@ import (
 
 type SubroutineDeclaration struct {
 	*Meta
-	Name  *Ident
-	Block *BlockStatement
+	Name       *Ident
+	Block      *BlockStatement
+	ReturnType *Ident
 }
 
 func (s *SubroutineDeclaration) statement()     {}
diff --git a/context/builtin.go b/context/builtin.go
index 7c7d5c2..9b5f77a 100644
--- a/context/builtin.go
+++ b/context/builtin.go
@@ -1204,7 +1204,7 @@ func builtinFunctions() Functions {
 					Value: &BuiltinFunction{
 						Return: types.BoolType,
 						Arguments: [][]types.Type{
-							[]types.Type{types.StringType, types.IDType, types.IntegerType, types.IntegerType, types.IntegerType, types.IDType, types.TimeType},
+							[]types.Type{types.StringType, types.IDType, types.IntegerType, types.IntegerType, types.IntegerType, types.IDType, types.RTimeType},
 						},
 						Scopes:    RECV | HASH | HIT | MISS | PASS | FETCH | ERROR | DELIVER | LOG,
 						Reference: "https://developer.fastly.com/reference/vcl/functions/rate-limiting/ratelimit-check-rate/",
@@ -1215,7 +1215,7 @@ func builtinFunctions() Functions {
 					Value: &BuiltinFunction{
 						Return: types.BoolType,
 						Arguments: [][]types.Type{
-							[]types.Type{types.StringType, types.IDType, types.IntegerType, types.IntegerType, types.IntegerType, types.IDType, types.IntegerType, types.IntegerType, types.IntegerType, types.IDType, types.TimeType},
+							[]types.Type{types.StringType, types.IDType, types.IntegerType, types.IntegerType, types.IntegerType, types.IDType, types.IntegerType, types.IntegerType, types.IntegerType, types.IDType, types.RTimeType},
 						},
 						Scopes:    RECV | HASH | HIT | MISS | PASS | FETCH | ERROR | DELIVER | LOG,
 						Reference: "https://developer.fastly.com/reference/vcl/functions/rate-limiting/ratelimit-check-rates/",
diff --git a/context/context.go b/context/context.go
index 1c45cbe..d175bc5 100644
--- a/context/context.go
+++ b/context/context.go
@@ -82,6 +82,7 @@ type Accessor struct {
 type Context struct {
 	curMode        int
 	prevMode       int
+	curName        string
 	Acls           map[string]*types.Acl
 	Backends       map[string]*types.Backend
 	Tables         map[string]*types.Table
@@ -93,11 +94,13 @@ type Context struct {
 	functions      Functions
 	Variables      Variables
 	RegexVariables map[string]int
+	ReturnType     *types.Type
 }
 
 func New() *Context {
 	return &Context{
 		curMode:        RECV,
+		curName:        "vcl_recv",
 		Acls:           make(map[string]*types.Acl),
 		Backends:       make(map[string]*types.Backend),
 		Tables:         make(map[string]*types.Table),
@@ -116,6 +119,10 @@ func (c *Context) Mode() int {
 	return c.curMode
 }
 
+func (c *Context) CurrentFunction() string {
+	return c.curName
+}
+
 func (c *Context) Restore() *Context {
 	c.curMode = c.prevMode
 	c.prevMode = 0
@@ -134,6 +141,16 @@ func (c *Context) Scope(mode int) *Context {
 	return c
 }
 
+func (c *Context) UserDefinedFunctionScope(name string, mode int, returnType types.Type) *Context {
+	c.prevMode = c.curMode
+	c.curMode = mode
+	// reset regex matched value
+	c.RegexVariables = newRegexMatchedValues()
+	c.ReturnType = &returnType
+	c.curName = name
+	return c
+}
+
 // Regex variables, "re.group.N" is a special variable in VCL.
 // These variables could use if(else) block statement when condition has regex operator like "~" or "!~".
 // Note that group matched variable has potential of making bugs due to its spec:
@@ -185,6 +202,31 @@ func (c *Context) GetRegexGroupVariable(name string) (types.Type, error) {
 	return types.StringType, nil
 }
 
+// Get ratecounter variable
+func (c *Context) GetRatecounterVariable(name string) (types.Type, error) {
+	// Ratecounter variables have the shape: ratecounter.{Variable Name}.[bucket/rate].Time
+	nameComponents := strings.Split(name, ".")
+	if len(nameComponents) != 4 {
+		return types.NullType, fmt.Errorf(`undefined variable "%s"`, name)
+	}
+
+	ratecounterVariableName := nameComponents[1]
+	// Check first if this ratecounter is defined in the first place.
+	if _, ok := c.Ratecounters[ratecounterVariableName]; !ok {
+		return types.NullType, fmt.Errorf(`undefined variable "%s"`, name)
+	}
+
+	// nameComponents[0] should be "ratecounter"
+	// nameComponents[1] should be the variable name
+	// nameComponents[2] should be either "bucket" or "rate"
+	// nameComponents[3] should be the time (10s, 60s, etc)
+	if v, ok := c.Variables[nameComponents[0]].Items["%any%"].Items[nameComponents[2]].Items[nameComponents[3]]; ok {
+		return v.Value.Get, nil
+	}
+
+	return types.NullType, fmt.Errorf(`undefined variable "%s"`, name)
+}
+
 func (c *Context) AddAcl(name string, acl *types.Acl) error {
 	// check existence
 	if _, duplicated := c.Acls[name]; duplicated {
@@ -203,11 +245,6 @@ func (c *Context) AddBackend(name string, backend *types.Backend) error {
 
 	// Additionally, assign some backend name related predefined variable
 	c.Variables["backend"].Items[name] = dynamicBackend()
-	c.Variables["ratecounter"] = &Object{
-		Items: map[string]*Object{
-			name: dynamicRateCounter(),
-		},
-	}
 
 	return nil
 }
@@ -262,13 +299,45 @@ func (c *Context) AddDirector(name string, director *types.Director) error {
 
 func (c *Context) AddSubroutine(name string, subroutine *types.Subroutine) error {
 	// check existence
+	if _, duplicated := c.functions[name]; duplicated {
+		if !IsFastlySubroutine(name) {
+			return fmt.Errorf(`duplicate definition of subroutine "%s"`, name)
+		}
+	}
+
 	if _, duplicated := c.Subroutines[name]; duplicated {
 		if !IsFastlySubroutine(name) {
 			return fmt.Errorf(`duplicate definition of subroutine "%s"`, name)
 		}
-	} else {
-		c.Subroutines[name] = subroutine
 	}
+
+	c.Subroutines[name] = subroutine
+	return nil
+}
+
+func (c *Context) AddUserDefinedFunction(name string, scopes int, returnType types.Type) error {
+	// check existence
+	if _, duplicated := c.functions[name]; duplicated {
+		if !IsFastlySubroutine(name) {
+			return fmt.Errorf(`duplication definition of subroutine "%s"`, name)
+		}
+	}
+
+	if _, duplicated := c.Subroutines[name]; duplicated {
+		if !IsFastlySubroutine(name) {
+			return fmt.Errorf(`duplication definition of subroutine "%s"`, name)
+		}
+	}
+
+	c.functions[name] = &FunctionSpec{
+		Items: map[string]*FunctionSpec{},
+		Value: &BuiltinFunction{
+			Return:    returnType,
+			Arguments: [][]types.Type{},
+			Scopes:    scopes,
+		},
+	}
+
 	return nil
 }
 
@@ -301,6 +370,12 @@ func (c *Context) Get(name string) (types.Type, error) {
 		return c.GetRegexGroupVariable(name)
 	}
 
+	// If program want to access to ratecounter variables like "ratecounter.{Name}.bucket.10s",
+	// proxy to dedicated getter.
+	if first == "ratecounter" {
+		return c.GetRatecounterVariable(name)
+	}
+
 	obj, ok := c.Variables[first]
 	if !ok {
 		return types.NullType, fmt.Errorf(`undefined variable "%s"`, name)
diff --git a/context/predefined.go b/context/predefined.go
index f14d9d1..503aaa0 100644
--- a/context/predefined.go
+++ b/context/predefined.go
@@ -2748,6 +2748,112 @@ func predefinedVariables() Variables {
 				},
 			},
 		},
+		"ratecounter": &Object{
+			Items: map[string]*Object{
+				"%any%": &Object{
+					Items: map[string]*Object{
+						"bucket": &Object{
+							Items: map[string]*Object{
+								"10s": &Object{
+									Items: map[string]*Object{},
+									Value: &Accessor{
+										Get:       types.IntegerType,
+										Set:       types.NeverType,
+										Unset:     false,
+										Scopes:    RECV | HASH | HIT | MISS | PASS | FETCH | ERROR | DELIVER | LOG,
+										Reference: "https://developer.fastly.com/reference/vcl/variables/rate-limiting/ratecounter-bucket-10s/",
+									},
+								},
+								"20s": &Object{
+									Items: map[string]*Object{},
+									Value: &Accessor{
+										Get:       types.IntegerType,
+										Set:       types.NeverType,
+										Unset:     false,
+										Scopes:    RECV | HASH | HIT | MISS | PASS | FETCH | ERROR | DELIVER | LOG,
+										Reference: "https://developer.fastly.com/reference/vcl/variables/rate-limiting/ratecounter-bucket-20s/",
+									},
+								},
+								"30s": &Object{
+									Items: map[string]*Object{},
+									Value: &Accessor{
+										Get:       types.IntegerType,
+										Set:       types.NeverType,
+										Unset:     false,
+										Scopes:    RECV | HASH | HIT | MISS | PASS | FETCH | ERROR | DELIVER | LOG,
+										Reference: "https://developer.fastly.com/reference/vcl/variables/rate-limiting/ratecounter-bucket-30s/",
+									},
+								},
+								"40s": &Object{
+									Items: map[string]*Object{},
+									Value: &Accessor{
+										Get:       types.IntegerType,
+										Set:       types.NeverType,
+										Unset:     false,
+										Scopes:    RECV | HASH | HIT | MISS | PASS | FETCH | ERROR | DELIVER | LOG,
+										Reference: "https://developer.fastly.com/reference/vcl/variables/rate-limiting/ratecounter-bucket-40s/",
+									},
+								},
+								"50s": &Object{
+									Items: map[string]*Object{},
+									Value: &Accessor{
+										Get:       types.IntegerType,
+										Set:       types.NeverType,
+										Unset:     false,
+										Scopes:    RECV | HASH | HIT | MISS | PASS | FETCH | ERROR | DELIVER | LOG,
+										Reference: "https://developer.fastly.com/reference/vcl/variables/rate-limiting/ratecounter-bucket-50s/",
+									},
+								},
+								"60s": &Object{
+									Items: map[string]*Object{},
+									Value: &Accessor{
+										Get:       types.IntegerType,
+										Set:       types.NeverType,
+										Unset:     false,
+										Scopes:    RECV | HASH | HIT | MISS | PASS | FETCH | ERROR | DELIVER | LOG,
+										Reference: "https://developer.fastly.com/reference/vcl/variables/rate-limiting/ratecounter-bucket-60s/",
+									},
+								},
+							},
+						},
+						"rate": &Object{
+							Items: map[string]*Object{
+								"10s": &Object{
+									Items: map[string]*Object{},
+									Value: &Accessor{
+										Get:       types.FloatType,
+										Set:       types.NeverType,
+										Unset:     false,
+										Scopes:    RECV | HASH | HIT | MISS | PASS | FETCH | ERROR | DELIVER | LOG,
+										Reference: "https://developer.fastly.com/reference/vcl/variables/rate-limiting/ratecounter-rate-10s/",
+									},
+								},
+								"1s": &Object{
+									Items: map[string]*Object{},
+									Value: &Accessor{
+										Get:       types.FloatType,
+										Set:       types.NeverType,
+										Unset:     false,
+										Scopes:    RECV | HASH | HIT | MISS | PASS | FETCH | ERROR | DELIVER | LOG,
+										Reference: "https://developer.fastly.com/reference/vcl/variables/rate-limiting/ratecounter-rate-1s/",
+									},
+								},
+								"60s": &Object{
+									Items: map[string]*Object{},
+									Value: &Accessor{
+										Get:       types.FloatType,
+										Set:       types.NeverType,
+										Unset:     false,
+										Scopes:    RECV | HASH | HIT | MISS | PASS | FETCH | ERROR | DELIVER | LOG,
+										Reference: "https://developer.fastly.com/reference/vcl/variables/rate-limiting/ratecounter-rate-60s/",
+									},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
 		"req": &Object{
 			Items: map[string]*Object{
 				"backend": &Object{
diff --git a/linter/helper.go b/linter/helper.go
index bac58ac..58bdf0c 100644
--- a/linter/helper.go
+++ b/linter/helper.go
@@ -184,6 +184,23 @@ func isValidStatmentExpression(exp ast.Expression) error {
 	return nil
 }
 
+func isValidReturnExpression(exp ast.Expression) error {
+	switch t := exp.(type) {
+	case *ast.PrefixExpression:
+		if t.Operator != "!" {
+			return fmt.Errorf("expected variable or function")
+		}
+		return isValidReturnExpression(t.Right)
+	case *ast.InfixExpression:
+		// Only comparison and boolean operators are allowed
+		if !(isBooleanOperator(t.Operator) || isComparisonOperator(t.Operator)) {
+			return fmt.Errorf("expected ;, got %s", t.Operator)
+		}
+		return isValidReturnExpression(t.Left)
+	}
+	return nil
+}
+
 func pushRegexGroupVars(exp ast.Expression, ctx *context.Context) error {
 	switch t := exp.(type) {
 	case *ast.PrefixExpression:
@@ -201,6 +218,44 @@ func pushRegexGroupVars(exp ast.Expression, ctx *context.Context) error {
 	return nil
 }
 
+func isBooleanOperator(operator string) bool {
+	switch operator {
+	case "(":
+		return true
+	case ")":
+		return true
+	case "!":
+		return true
+	case "&&":
+		return true
+	case "||":
+		return true
+	}
+	return false
+}
+
+func isComparisonOperator(operator string) bool {
+	switch operator {
+	case "==":
+		return true
+	case "!=":
+		return true
+	case "~":
+		return true
+	case "!~":
+		return true
+	case ">":
+		return true
+	case "<":
+		return true
+	case ">=":
+		return true
+	case "<=":
+		return true
+	}
+	return false
+}
+
 func isConstantExpression(exp ast.Expression) bool {
 	switch exp.(type) {
 	case *ast.Float:
diff --git a/linter/linter.go b/linter/linter.go
index f7fcd0e..c629473 100644
--- a/linter/linter.go
+++ b/linter/linter.go
@@ -47,6 +47,8 @@ func (l *Linter) Lint(node ast.Node, ctx *context.Context, isMain bool) types.Ty
 		l.lintUnusedAcls(ctx)
 		l.lintUnusedBackends(ctx)
 		l.lintUnusedSubroutines(ctx)
+		l.lintUnusedPenaltyboxes(ctx)
+		l.lintUnusedRatecounters(ctx)
 	}
 
 	return types.NeverType
@@ -96,6 +98,24 @@ func (l *Linter) lintUnusedSubroutines(ctx *context.Context) {
 	}
 }
 
+func (l *Linter) lintUnusedPenaltyboxes(ctx *context.Context) {
+	for _, p := range ctx.Penaltyboxes {
+		if p.IsUsed {
+			continue
+		}
+		l.Error(UnusedDeclaration(p.Decl.GetMeta(), p.Decl.Name.Value, "penaltybox").Match(UNUSED_DECLARATION))
+	}
+}
+
+func (l *Linter) lintUnusedRatecounters(ctx *context.Context) {
+	for _, rc := range ctx.Ratecounters {
+		if rc.IsUsed {
+			continue
+		}
+		l.Error(UnusedDeclaration(rc.Decl.GetMeta(), rc.Decl.Name.Value, "ratecounter").Match(UNUSED_DECLARATION))
+	}
+}
+
 func (l *Linter) lintUnusedVariables(ctx *context.Context) {
 	v, ok := ctx.Variables["var"]
 	if !ok {
@@ -283,17 +303,48 @@ func (l *Linter) factoryRootStatements(vcl *ast.VCL, ctx *context.Context) []ast
 			}
 			statements = append(statements, stmt)
 		case *ast.SubroutineDeclaration:
-			if err := ctx.AddSubroutine(t.Name.Value, &types.Subroutine{Decl: t, Body: t.Block}); err != nil {
-				e := &LintError{
-					Severity: ERROR,
-					Token:    t.Name.GetMeta().Token,
-					Message:  err.Error(),
+			if t.ReturnType != nil {
+				if context.IsFastlySubroutine(t.Name.Value) {
+					err := &LintError{
+						Severity: ERROR,
+						Token:    t.ReturnType.GetMeta().Token,
+						Message:  fmt.Sprintf("State-machine method %s may not have a return type", t.Name.Value),
+					}
+					l.Error(err.Match(SUBROUTINE_INVALID_RETURN_TYPE))
+				}
+
+				returnType, ok := ValueTypeMap[t.ReturnType.Value]
+				if !ok {
+					err := &LintError{
+						Severity: ERROR,
+						Token:    t.ReturnType.GetMeta().Token,
+						Message:  fmt.Sprintf("Unexpected variable type found: %s", t.ReturnType.Value),
+					}
+					l.Error(err.Match(SUBROUTINE_INVALID_RETURN_TYPE))
+				}
+
+				if err := ctx.AddUserDefinedFunction(t.Name.Value, getSubroutineCallScope(t), returnType); err != nil {
+					err := &LintError{
+						Severity: ERROR,
+						Token:    t.Name.GetMeta().Token,
+						Message:  err.Error(),
+					}
+					l.Error(err.Match(SUBROUTINE_DUPLICATED))
+				}
+			} else {
+				err := ctx.AddSubroutine(t.Name.Value, &types.Subroutine{Decl: t, Body: t.Block})
+				if err != nil {
+					e := &LintError{
+						Severity: ERROR,
+						Token:    t.Name.GetMeta().Token,
+						Message:  err.Error(),
+					}
+					l.Error(e.Match(SUBROUTINE_DUPLICATED))
 				}
-				l.Error(e.Match(SUBROUTINE_DUPLICATED))
 			}
 			statements = append(statements, stmt)
 		case *ast.PenaltyboxDeclaration:
-			if err := ctx.AddPenaltybox(t.Name.Value, &types.Penaltybox{Penaltybox: t}); err != nil {
+			if err := ctx.AddPenaltybox(t.Name.Value, &types.Penaltybox{Decl: t}); err != nil {
 				e := &LintError{
 					Severity: ERROR,
 					Token:    t.Name.GetMeta().Token,
@@ -303,7 +354,7 @@ func (l *Linter) factoryRootStatements(vcl *ast.VCL, ctx *context.Context) []ast
 			}
 			statements = append(statements, stmt)
 		case *ast.RatecounterDeclaration:
-			if err := ctx.AddRatecounter(t.Name.Value, &types.Ratecounter{Ratecounter: t}); err != nil {
+			if err := ctx.AddRatecounter(t.Name.Value, &types.Ratecounter{Decl: t}); err != nil {
 				e := &LintError{
 					Severity: ERROR,
 					Token:    t.Name.GetMeta().Token,
@@ -588,7 +639,15 @@ func (l *Linter) lintSubRoutineDeclaration(decl *ast.SubroutineDeclaration, ctx
 		l.Error(InvalidName(decl.Name.GetMeta(), decl.Name.Value, "sub").Match(SUBROUTINE_SYNTAX))
 	}
 
-	cc := ctx.Scope(getSubroutineCallScope(decl))
+	scope := getSubroutineCallScope(decl)
+	var cc *context.Context
+	if decl.ReturnType != nil {
+		returnType := ValueTypeMap[decl.ReturnType.Value]
+		cc = ctx.UserDefinedFunctionScope(decl.Name.Value, scope, returnType)
+	} else {
+		cc = ctx.Scope(scope)
+	}
+
 	// Switch context mode which corredponds to call scope and restore after linting block statements
 	defer func() {
 		// Lint declared variables are used
@@ -596,6 +655,9 @@ func (l *Linter) lintSubRoutineDeclaration(decl *ast.SubroutineDeclaration, ctx
 		cc.Restore()
 	}()
 	l.lint(decl.Block, cc)
+	// We are done linting inside the previous scope so
+	// we dont need the return type anymore
+	cc.ReturnType = nil
 
 	// If fastly reserved subroutine name (e.g vcl_recv, vcl_fetch, etc),
 	// validate fastly specific boilerplate macro is embedded like "FASTLY recv"
@@ -727,7 +789,7 @@ func (l *Linter) lintSetStatement(stmt *ast.SetStatement, ctx *context.Context)
 	// Fastly has various assignment operators and required correspond types for each operator
 	// https://developer.fastly.com/reference/vcl/operators/#assignment-operators
 	//
-	// Above document is not enough to explain for other types... actually more complex type coparison may occur.
+	// Above document is not enough to explain for other types... actually more complex type comparison may occur.
 	// We investigated type comparison and summarized.
 	// See: https://docs.google.com/spreadsheets/d/16xRPugw9ubKA1nXHIc5ysVZKokLLhysI-jAu3qbOFJ8/edit#gid=0
 	switch stmt.Operator.Operator {
@@ -964,6 +1026,49 @@ func (l *Linter) lintLogStatement(stmt *ast.LogStatement, ctx *context.Context)
 }
 
 func (l *Linter) lintReturnStatement(stmt *ast.ReturnStatement, ctx *context.Context) types.Type {
+	if ctx.ReturnType != nil {
+		if stmt.HasParenthesis {
+			l.Error(&LintError{
+				Severity: ERROR,
+				Token:    stmt.Token,
+				Message:  fmt.Sprintf("function %s: only actions should be enclosed in ()", ctx.CurrentFunction()),
+			})
+		}
+
+		if stmt.ReturnExpression == nil {
+			lintErr := &LintError{
+				Severity: ERROR,
+				Token:    stmt.Token,
+				Message:  fmt.Sprintf("function %s must have a return value", ctx.CurrentFunction()),
+			}
+			l.Error(lintErr.Match(SUBROUTINE_INVALID_RETURN_TYPE))
+			return types.NeverType
+		}
+
+		err := isValidReturnExpression(*stmt.ReturnExpression)
+		if err != nil {
+			lintErr := &LintError{
+				Severity: ERROR,
+				Token:    (*stmt.ReturnExpression).GetMeta().Token,
+				Message:  err.Error(),
+			}
+			l.Error(lintErr.Match(SUBROUTINE_INVALID_RETURN_TYPE))
+			return types.NeverType
+		}
+
+		cc := l.lint(*stmt.ReturnExpression, ctx)
+		// Condition expression return type must be BOOL or STRING
+		if !expectType(cc, *ctx.ReturnType) {
+			lintErr := &LintError{
+				Severity: ERROR,
+				Token:    (*stmt.ReturnExpression).GetMeta().Token,
+				Message:  fmt.Sprintf("function %s return type is incompatible with type %s", ctx.CurrentFunction(), cc.String()),
+			}
+			l.Error(lintErr.Match(SUBROUTINE_INVALID_RETURN_TYPE))
+		}
+		return types.NeverType
+	}
+
 	// legal return actions are different in subroutine.
 	// https://developer.fastly.com/learning/vcl/using/#the-vcl-request-lifecycle
 	expects := make([]string, 0, 3)
@@ -999,13 +1104,13 @@ func (l *Linter) lintReturnStatement(stmt *ast.ReturnStatement, ctx *context.Con
 	}
 
 	// return statement may not have arguments, then stop linting
-	if stmt.Ident == nil {
+	if stmt.ReturnExpression == nil {
 		return types.NeverType
 	}
 
-	if !expectState(stmt.Ident.Value, expects...) {
+	if !expectState((*stmt.ReturnExpression).String(), expects...) {
 		l.Error(InvalidReturnState(
-			stmt.Ident.GetMeta(), context.ScopeString(ctx.Mode()), stmt.Ident.Value, expects...,
+			(*stmt.ReturnExpression).GetMeta(), context.ScopeString(ctx.Mode()), (*stmt.ReturnExpression).String(), expects...,
 		).Match(RESTART_STATEMENT_SCOPE))
 	}
 	return types.NeverType
@@ -1041,6 +1146,16 @@ func (l *Linter) lintIdent(exp *ast.Ident, ctx *context.Context) types.Type {
 			// mark table is used
 			t.IsUsed = true
 			return types.TableType
+		} else if p, ok := ctx.Penaltyboxes[exp.Value]; ok {
+			// mark penaltybox is used
+			p.IsUsed = true
+			// Fastly treats these variables as type IDs
+			return types.IDType
+		} else if rc, ok := ctx.Ratecounters[exp.Value]; ok {
+			// mark ratecounter is used
+			rc.IsUsed = true
+			// Fastly treats these variables as type IDs
+			return types.IDType
 		} else if _, ok := ctx.Identifiers[exp.Value]; ok {
 			return types.IDType
 		}
diff --git a/linter/rules.go b/linter/rules.go
index f9e3195..2daff1b 100644
--- a/linter/rules.go
+++ b/linter/rules.go
@@ -30,6 +30,7 @@ const (
 	SUBROUTINE_SYNTAX                    = "subroutine/syntax"
 	SUBROUTINE_BOILERPLATE_MACRO         = "subroutine/boilerplate-macro"
 	SUBROUTINE_DUPLICATED                = "subroutine/duplicated"
+	SUBROUTINE_INVALID_RETURN_TYPE       = "subroutine/invalid-return-type"
 	PENALTYBOX_SYNTAX                    = "penaltybox/syntax"
 	PENALTYBOX_DUPLICATED                = "penaltybox/duplicated"
 	PENALTYBOX_NONEMPTY_BLOCK            = "penaltybox/nonempty-block"
diff --git a/parser/declaration_parser.go b/parser/declaration_parser.go
index 5eeb2de..b7c0000 100644
--- a/parser/declaration_parser.go
+++ b/parser/declaration_parser.go
@@ -431,6 +431,13 @@ func (p *Parser) parseSubroutineDeclaration() (*ast.SubroutineDeclaration, error
 	}
 	s.Name = p.parseIdent()
 
+	// Custom subroutines might be returning a type
+	// https://developer.fastly.com/reference/vcl/subroutines/
+	// we dont need to validate the type here, linter will do that later.
+	if p.expectPeek(token.IDENT) {
+		s.ReturnType = p.parseIdent()
+	}
+
 	if !p.expectPeek(token.LEFT_BRACE) {
 		return nil, errors.WithStack(UnexpectedToken(p.peekToken, "LEFT_BRACE"))
 	}
diff --git a/parser/statement_parser.go b/parser/statement_parser.go
index 0c137fa..863400b 100644
--- a/parser/statement_parser.go
+++ b/parser/statement_parser.go
@@ -372,7 +372,8 @@ func (p *Parser) parseRestartStatement() (*ast.RestartStatement, error) {
 
 func (p *Parser) parseReturnStatement() (*ast.ReturnStatement, error) {
 	stmt := &ast.ReturnStatement{
-		Meta: p.curToken,
+		Meta:           p.curToken,
+		HasParenthesis: false,
 	}
 
 	// return statement may not have argument
@@ -383,19 +384,31 @@ func (p *Parser) parseReturnStatement() (*ast.ReturnStatement, error) {
 		return stmt, nil
 	}
 
-	if !p.expectPeek(token.LEFT_PAREN) {
-		return nil, errors.WithStack(UnexpectedToken(p.peekToken, "LEFT_PAREN"))
+	hasLeftParen := p.peekTokenIs(token.LEFT_PAREN)
+	if hasLeftParen {
+		stmt.HasParenthesis = true
+		p.nextToken() // point to expression
 	}
+	p.nextToken()
 
-	if !p.expectPeek(token.IDENT) {
-		return nil, errors.WithStack(UnexpectedToken(p.peekToken, "IDENT"))
+	expression, err := p.parseExpression(LOWEST)
+	if err != nil {
+		return nil, errors.WithStack(err)
 	}
-	stmt.Ident = p.parseIdent()
+	stmt.ReturnExpression = &expression
 
-	if !p.expectPeek(token.RIGHT_PAREN) {
-		return nil, errors.WithStack(UnexpectedToken(p.peekToken, "RIGHT_PAREN"))
+	hasRightParen := p.peekTokenIs(token.RIGHT_PAREN)
+	if hasRightParen {
+		p.nextToken() // point to condition expression
 	}
-	swapLeadingTrailing(p.curToken, stmt.Ident.Meta)
+	if hasLeftParen != hasRightParen {
+		return nil, errors.WithStack(&ParseError{
+			Token:   p.curToken.Token,
+			Message: "Parenthesis missmatch",
+		})
+	}
+
+	swapLeadingTrailing(p.curToken, (*stmt.ReturnExpression).GetMeta())
 
 	if !p.peekTokenIs(token.SEMICOLON) {
 		return nil, errors.WithStack(MissingSemicolon(p.curToken))
diff --git a/types/types.go b/types/types.go
index 3a958a1..5a9ad09 100644
--- a/types/types.go
+++ b/types/types.go
@@ -226,17 +226,19 @@ func (d *Director) Token() token.Token { return d.Decl.Token }
 func (d *Director) String() string     { return d.Decl.String() }
 
 type Penaltybox struct {
-	Penaltybox *ast.PenaltyboxDeclaration
+	Decl   *ast.PenaltyboxDeclaration
+	IsUsed bool // mark this penaltybox is called at least once
 }
 
 func (pb *Penaltybox) Type() Type         { return PenaltyboxType }
-func (pb *Penaltybox) Token() token.Token { return pb.Penaltybox.Token }
-func (pb *Penaltybox) String() string     { return pb.Penaltybox.String() }
+func (pb *Penaltybox) Token() token.Token { return pb.Decl.Token }
+func (pb *Penaltybox) String() string     { return pb.Decl.String() }
 
 type Ratecounter struct {
-	Ratecounter *ast.RatecounterDeclaration
+	Decl   *ast.RatecounterDeclaration
+	IsUsed bool // mark this ratecounter is called at least once
 }
 
 func (rc *Ratecounter) Type() Type         { return RatecounterType }
-func (rc *Ratecounter) Token() token.Token { return rc.Ratecounter.Token }
-func (rc *Ratecounter) String() string     { return rc.Ratecounter.String() }
+func (rc *Ratecounter) Token() token.Token { return rc.Decl.Token }
+func (rc *Ratecounter) String() string     { return rc.Decl.String() }


Test patch:
diff --git a/ast/return_statement_test.go b/ast/return_statement_test.go
index 82079c4..c32423a 100644
--- a/ast/return_statement_test.go
+++ b/ast/return_statement_test.go
@@ -5,12 +5,14 @@ import (
 )
 
 func TestReturnStatement(t *testing.T) {
+	var rt Expression
+	rt = &Ident{
+		Meta:  New(T, 0),
+		Value: "pass",
+	}
 	r := &ReturnStatement{
-		Meta: New(T, 0, comments("// This is comment"), comments("// This is comment")),
-		Ident: &Ident{
-			Meta:  New(T, 0),
-			Value: "pass",
-		},
+		Meta:             New(T, 0, comments("// This is comment"), comments("// This is comment")),
+		ReturnExpression: &rt,
 	}
 
 	expect := `// This is comment
diff --git a/linter/linter_test.go b/linter/linter_test.go
index bc0a5b2..2b2ab8e 100644
--- a/linter/linter_test.go
+++ b/linter/linter_test.go
@@ -332,6 +332,15 @@ sub foo {
 		input := `
 sub vcl_recv {
 	set req.http.Host = "example.com";
+}`
+		assertError(t, input)
+	})
+
+	t.Run("Fastly reserved subroutine cannot have a return type", func(t *testing.T) {
+		input := `
+sub vcl_recv BOOL {
+	set req.http.Host = "example.com";
+	return true;
 }`
 		assertError(t, input)
 	})
@@ -1169,6 +1178,18 @@ sub foo {
 		assertNoError(t, input)
 	})
 
+	t.Run("pass with user defined sub", func(t *testing.T) {
+		input := `
+sub returns_one INTEGER {
+	return 1;
+}
+
+sub returns_true BOOL {
+	return returns_one() == 1;
+}`
+		assertNoError(t, input)
+	})
+
 	t.Run("function not found", func(t *testing.T) {
 		input := `
 sub foo {
@@ -1229,6 +1250,88 @@ sub vcl_recv {
 }`
 		assertNoError(t, input)
 	})
+
+	t.Run("sub: return correct type", func(t *testing.T) {
+		input := `
+sub custom_sub INTEGER {
+	#Fastly recv
+	return 1;
+}`
+		assertNoError(t, input)
+	})
+
+	t.Run("sub: return empty statement", func(t *testing.T) {
+		input := `
+sub custom_sub INTEGER {
+	return;
+}`
+		assertError(t, input)
+	})
+
+	t.Run("sub: return wrong type", func(t *testing.T) {
+		input := `
+sub custom_sub INTEGER {
+	return (req.http.foo);
+}`
+		assertError(t, input)
+	})
+
+	t.Run("sub: return action", func(t *testing.T) {
+		input := `
+sub custom_sub INTEGER {
+	return (pass);
+}`
+		assertError(t, input)
+	})
+
+	t.Run("sub: return value as action", func(t *testing.T) {
+		input := `
+sub custom_sub INTEGER {
+	return (1);
+}`
+		assertError(t, input)
+	})
+
+	t.Run("sub: return local value", func(t *testing.T) {
+		input := `
+sub custom_sub INTEGER {
+	declare local var.tmp INTEGER;
+	set var.tmp = 10;
+	return var.tmp;
+}`
+		assertNoError(t, input)
+	})
+
+	t.Run("sub: return value contains operations", func(t *testing.T) {
+		input := `
+sub get_str STRING {
+	declare local var.tmp STRING;
+	set var.tmp = "foo";
+	return var.tmp "bar";
+}`
+		assertError(t, input)
+	})
+
+	t.Run("sub: return value contains jibber", func(t *testing.T) {
+		input := `
+sub get_str STRING {
+	declare local var.tmp STRING;
+	set var.tmp = "foo";
+	return +-var.tmp;
+}`
+		assertError(t, input)
+	})
+
+	t.Run("sub: bool return value is allowed to have operations", func(t *testing.T) {
+		input := `
+sub get_bool BOOL {
+	declare local var.tmp STRING;
+	set var.tmp = "foo";
+	return std.strlen(var.tmp) > 5;
+}`
+		assertNoError(t, input)
+	})
+
 }
 
 func TestBlockSyntaxInsideBlockStatement(t *testing.T) {
@@ -1541,6 +1644,45 @@ penaltybox ip_pb {}
 penaltybox ip_pb {
 	set var.bar = "baz";
 }
+`
+		assertError(t, input)
+	})
+
+	t.Run("penaltybox variable should be pass if it is defined", func(t *testing.T) {
+		input := `
+penaltybox ip_pb {}
+ratecounter counter_60 {}
+
+sub test_sub{
+	declare local var.ratelimit_exceeded BOOL;
+	set var.ratelimit_exceeded = ratelimit.check_rate(
+		digest.hash_sha256("123"),
+		counter_60,
+		1,
+		60,
+		135,
+		ip_pb,
+		2m);
+}
+`
+		assertNoError(t, input)
+	})
+
+	t.Run("penaltybox variable should be defined", func(t *testing.T) {
+		input := `
+ratecounter counter_60 {}
+
+sub test_sub{
+	declare local var.ratelimit_exceeded BOOL;
+	set var.ratelimit_exceeded = ratelimit.check_rate(
+		digest.hash_sha256("123"),
+		counter_60,
+		1,
+		60,
+		135,
+		ip_pb,
+		2m);
+}
 `
 		assertError(t, input)
 	})
@@ -1583,6 +1725,78 @@ ratecounter req_counter {}
 ratecounter req_counter {
 	set var.bar = "baz";
 }
+`
+		assertError(t, input)
+	})
+
+	t.Run("ratecounter variable should be pass if it is defined", func(t *testing.T) {
+		input := `
+penaltybox ip_pb {}
+ratecounter counter_60 {}
+
+sub test_sub{
+	declare local var.ratelimit_exceeded BOOL;
+	set var.ratelimit_exceeded = ratelimit.check_rate(
+		digest.hash_sha256("123"),
+		counter_60,
+		1,
+		60,
+		135,
+		ip_pb,
+		2m);
+}
+`
+		assertNoError(t, input)
+	})
+
+	t.Run("ratecounter variable should be defined", func(t *testing.T) {
+		input := `
+penaltybox ip_pb {}
+
+sub test_sub{
+	declare local var.ratelimit_exceeded BOOL;
+	set var.ratelimit_exceeded = ratelimit.check_rate(
+		digest.hash_sha256("123"),
+		counter_60,
+		1,
+		60,
+		135,
+		ip_pb,
+		2m);
+}
+`
+		assertError(t, input)
+	})
+
+	t.Run("ratecounter bucket variables should pass if the ratecounter is defined", func(t *testing.T) {
+		input := `
+ratecounter counter_60 {}
+
+sub test_sub{
+	set req.http.X-ERL:tls_bucket_10s = std.itoa(ratecounter.counter_60.bucket.10s);
+}
+`
+		assertNoError(t, input)
+	})
+
+	t.Run("ratecounter bucket variables should not pass if the ratecounter is not defined", func(t *testing.T) {
+		input := `
+ratecounter counter_60 {}
+
+sub test_sub{
+	set req.http.X-ERL:tls_rate_10s = std.itoa(ratecounter.counter.bucket.10s);
+}
+`
+		assertError(t, input)
+	})
+
+	t.Run("ratecounter bucket variables should exist", func(t *testing.T) {
+		input := `
+ratecounter counter_60 {}
+
+sub test_sub{
+	set req.http.X-ERL:tls_bucket_10s = std.itoa(ratecounter.counter_60.bucket.100s);
+}
 `
 		assertError(t, input)
 	})
diff --git a/parser/statement_parser_test.go b/parser/statement_parser_test.go
index 84e972a..ac0797a 100644
--- a/parser/statement_parser_test.go
+++ b/parser/statement_parser_test.go
@@ -1103,6 +1103,11 @@ sub vcl_recv {
 	// Leading comment
 	return(deliver); // Trailing comment
 }`
+	var rt ast.Expression
+	rt = &ast.Ident{
+		Meta:  ast.New(T, 1),
+		Value: "deliver",
+	}
 	expect := &ast.VCL{
 		Statements: []ast.Statement{
 			&ast.SubroutineDeclaration{
@@ -1115,11 +1120,9 @@ sub vcl_recv {
 					Meta: ast.New(T, 1),
 					Statements: []ast.Statement{
 						&ast.ReturnStatement{
-							Meta: ast.New(T, 1, comments("// Leading comment"), comments("// Trailing comment")),
-							Ident: &ast.Ident{
-								Meta:  ast.New(T, 1),
-								Value: "deliver",
-							},
+							Meta:             ast.New(T, 1, comments("// Leading comment"), comments("// Trailing comment")),
+							ReturnExpression: &rt,
+							HasParenthesis:   true,
 						},
 					},
 				},