From cee4ab8ab891f473a8215d0405eba63def717eae Mon Sep 17 00:00:00 2001
From: Ravishankar <ravishankar.gnanaprakasam@gmail.com>
Date: Wed, 4 Dec 2024 19:14:36 +0530
Subject: [PATCH] feat: Add new stdlib functions (#2199)

* feat: Add new stdlib functions

* Added from_URLbase64 and review comments

* Update docs/sources/reference/stdlib/encoding.md

Co-authored-by: Clayton Cornell <131809008+clayton-cornell@users.noreply.github.com>

* Update docs/sources/reference/stdlib/encoding.md

Co-authored-by: Clayton Cornell <131809008+clayton-cornell@users.noreply.github.com>

* Update docs/sources/reference/stdlib/encoding.md

Co-authored-by: Clayton Cornell <131809008+clayton-cornell@users.noreply.github.com>

* Update CHANGELOG.md

---------

Co-authored-by: Clayton Cornell <131809008+clayton-cornell@users.noreply.github.com>
---
 CHANGELOG.md                              |  2 ++
 docs/sources/reference/stdlib/encoding.md | 37 ++++++++++++++++++++++-
 syntax/internal/stdlib/stdlib.go          | 27 +++++++++++++++--
 syntax/vm/vm_stdlib_test.go               |  3 ++
 4 files changed, 65 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1791a758bd..b574a7896c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -44,6 +44,8 @@ Main (unreleased)
   - Change processlist query to support ONLY_FULL_GROUP_BY sql_mode
   - Add perf_schema quantile columns to collector
 
+- Add three new stdlib functions to_base64, from_URLbase64 and to_URLbase64 (@ravishankar15)
+
 ### Bugfixes
 
 - Fixed issue with reloading configuration and prometheus metrics duplication in `prometheus.write.queue`. (@mattdurham)
diff --git a/docs/sources/reference/stdlib/encoding.md b/docs/sources/reference/stdlib/encoding.md
index 9bc8fcaed5..7e3aa84d16 100644
--- a/docs/sources/reference/stdlib/encoding.md
+++ b/docs/sources/reference/stdlib/encoding.md
@@ -19,13 +19,48 @@ The `encoding.from_base64` function decodes a RFC4648-compliant Base64-encoded s
 
 `encoding.from_base64` fails if the provided string argument contains invalid Base64 data.
 
-### Examples
+### Example
 
 ```text
 > encoding.from_base64("dGFuZ2VyaW5l")
 tangerine
 ```
 
+## encoding.from_URLbase64
+
+The `encoding.from_URLbase64` function decodes a RFC4648-compliant Base64 URL safe encoded string into the original string.
+
+`encoding.from_URLbase64` fails if the provided string argument contains invalid Base64 data.
+
+### Example
+
+```
+> encoding.from_URLbase64("c3RyaW5nMTIzIT8kKiYoKSctPUB-")
+string123!?$*&()'-=@~
+```
+
+## encoding.to_base64
+
+The `encoding.to_base64` function encodes the original string into RFC4648-compliant Base64 encoded string.
+
+### Example
+
+```
+> encoding.to_base64("string123!?$*&()'-=@~")
+c3RyaW5nMTIzIT8kKiYoKSctPUB+
+```
+
+## encoding.to_URLbase64
+
+The `encoding.to_base64` function encodes the original string into RFC4648-compliant URL safe Base64 encoded string.
+
+### Example
+
+```
+> encoding.to_URLbase64("string123!?$*&()'-=@~")
+c3RyaW5nMTIzIT8kKiYoKSctPUB-
+```
+
 ## encoding.from_json
 
 The `encoding.from_json` function decodes a string representing JSON into an {{< param "PRODUCT_NAME" >}} value.
diff --git a/syntax/internal/stdlib/stdlib.go b/syntax/internal/stdlib/stdlib.go
index 85dadd9c67..e81856ae30 100644
--- a/syntax/internal/stdlib/stdlib.go
+++ b/syntax/internal/stdlib/stdlib.go
@@ -75,9 +75,12 @@ var file = map[string]interface{}{
 }
 
 var encoding = map[string]interface{}{
-	"from_json":   jsonDecode,
-	"from_yaml":   yamlDecode,
-	"from_base64": base64Decode,
+	"from_json":      jsonDecode,
+	"from_yaml":      yamlDecode,
+	"from_base64":    base64Decode,
+	"from_URLbase64": base64URLDecode,
+	"to_base64":      base64Encode,
+	"to_URLbase64":   base64URLEncode,
 }
 
 var str = map[string]interface{}{
@@ -313,6 +316,24 @@ func base64Decode(in string) (interface{}, error) {
 	return decoded, nil
 }
 
+func base64URLDecode(in string) (interface{}, error) {
+	decoded, err := base64.URLEncoding.DecodeString(in)
+	if err != nil {
+		return nil, err
+	}
+	return decoded, nil
+}
+
+func base64URLEncode(in string) (interface{}, error) {
+	encoded := base64.URLEncoding.EncodeToString([]byte(in))
+	return encoded, nil
+}
+
+func base64Encode(in string) (interface{}, error) {
+	encoded := base64.StdEncoding.EncodeToString([]byte(in))
+	return encoded, nil
+}
+
 func jsonPath(jsonString string, path string) (interface{}, error) {
 	jsonPathExpr, err := jp.ParseString(path)
 	if err != nil {
diff --git a/syntax/vm/vm_stdlib_test.go b/syntax/vm/vm_stdlib_test.go
index 36a74c29ab..70dc2711c6 100644
--- a/syntax/vm/vm_stdlib_test.go
+++ b/syntax/vm/vm_stdlib_test.go
@@ -39,6 +39,9 @@ func TestVM_Stdlib(t *testing.T) {
 		{"encoding.from_yaml nil field", "encoding.from_yaml(`foo: null`)", map[string]interface{}{"foo": nil}},
 		{"encoding.from_yaml nil array element", `encoding.from_yaml("[0, null]")`, []interface{}{0, nil}},
 		{"encoding.from_base64", `encoding.from_base64("Zm9vYmFyMTIzIT8kKiYoKSctPUB+")`, string(`foobar123!?$*&()'-=@~`)},
+		{"encoding.from_URLbase64", `encoding.from_URLbase64("c3RyaW5nMTIzIT8kKiYoKSctPUB-")`, string(`string123!?$*&()'-=@~`)},
+		{"encoding.to_base64", `encoding.to_base64("string123!?$*&()'-=@~")`, string(`c3RyaW5nMTIzIT8kKiYoKSctPUB+`)},
+		{"encoding.to_URLbase64", `encoding.to_URLbase64("string123!?$*&()'-=@~")`, string(`c3RyaW5nMTIzIT8kKiYoKSctPUB-`)},
 
 		// Map tests
 		{