From 0079ca1e395f15ddd9fee75edb98289cc824f0aa Mon Sep 17 00:00:00 2001
From: UgOrange <lichanghao.orange@bytedance.com>
Date: Tue, 21 Nov 2023 12:17:26 +0800
Subject: [PATCH] feat: Add external_url_check custom JMESPath function (#8614)

Signed-off-by: lichanghao.orange <lichanghao.orange@bytedance.com>
Signed-off-by: UgOrange <lichanghao.orange@bytedance.com>
---
 pkg/engine/jmespath/functions.go      | 40 ++++++++++++++++++
 pkg/engine/jmespath/functions_test.go | 58 +++++++++++++++++++++++++++
 2 files changed, 98 insertions(+)

diff --git a/pkg/engine/jmespath/functions.go b/pkg/engine/jmespath/functions.go
index 99c545e6d6..4f82475086 100644
--- a/pkg/engine/jmespath/functions.go
+++ b/pkg/engine/jmespath/functions.go
@@ -12,6 +12,8 @@ import (
 	"errors"
 	"fmt"
 	"math"
+	"net"
+	"net/url"
 	"path/filepath"
 	"reflect"
 	"regexp"
@@ -73,6 +75,7 @@ var (
 	random                 = "random"
 	x509_decode            = "x509_decode"
 	imageNormalize         = "image_normalize"
+	isExternalURL          = "is_external_url"
 )
 
 func GetFunctions(configuration config.Configuration) []FunctionEntry {
@@ -580,6 +583,16 @@ func GetFunctions(configuration config.Configuration) []FunctionEntry {
 		},
 		ReturnType: []jpType{jpString},
 		Note:       "normalizes an image reference",
+	}, {
+		FunctionEntry: gojmespath.FunctionEntry{
+			Name: isExternalURL,
+			Arguments: []argSpec{
+				{Types: []jpType{jpString}},
+			},
+			Handler: jpIsExternalURL,
+		},
+		ReturnType: []jpType{jpBool},
+		Note:       "determine if a URL points to an external network address",
 	}}
 }
 
@@ -1221,3 +1234,30 @@ func jpImageNormalize(configuration config.Configuration) gojmespath.JpFunction
 		}
 	}
 }
+
+func jpIsExternalURL(arguments []interface{}) (interface{}, error) {
+	var err error
+	str, err := validateArg(pathCanonicalize, arguments, 0, reflect.String)
+	if err != nil {
+		return nil, err
+	}
+	parsedURL, err := url.Parse(str.String())
+	if err != nil {
+		return false, err
+	}
+	ip := net.ParseIP(parsedURL.Hostname())
+	if ip != nil {
+		return !(ip.IsLoopback() || ip.IsPrivate()), nil
+	}
+	// If it can't be parsed as an IP, then resolve the domain name
+	ips, err := net.LookupIP(parsedURL.Hostname())
+	if err != nil {
+		return nil, err
+	}
+	for _, ip := range ips {
+		if ip.IsLoopback() || ip.IsPrivate() {
+			return false, nil
+		}
+	}
+	return true, nil
+}
diff --git a/pkg/engine/jmespath/functions_test.go b/pkg/engine/jmespath/functions_test.go
index 6e83391943..0325521094 100644
--- a/pkg/engine/jmespath/functions_test.go
+++ b/pkg/engine/jmespath/functions_test.go
@@ -1675,3 +1675,61 @@ func Test_ImageNormalize(t *testing.T) {
 		})
 	}
 }
+func Test_IsExternalURL(t *testing.T) {
+	testCases := []struct {
+		jmesPath       string
+		expectedResult bool
+	}{
+		{
+			jmesPath:       "is_external_url('http://localhost')",
+			expectedResult: false,
+		}, {
+			jmesPath:       "is_external_url('http://127.0.0.1')",
+			expectedResult: false,
+		}, {
+			jmesPath:       "is_external_url('http://172.16.1.1')",
+			expectedResult: false,
+		}, {
+			jmesPath:       "is_external_url('http://10.1.1.1')",
+			expectedResult: false,
+		}, {
+			jmesPath:       "is_external_url('http://192.168.0.1')",
+			expectedResult: false,
+		}, {
+			jmesPath:       "is_external_url('http://google.com')",
+			expectedResult: true,
+		}, {
+			jmesPath:       "is_external_url('http://8.8.8.8')",
+			expectedResult: true,
+		}, {
+			jmesPath:       "is_external_url('http://245.48.83.160')",
+			expectedResult: true,
+		}, {
+			jmesPath:       "is_external_url('http://[fd12:3456:789a:1::1]')", //192.168.3.6
+			expectedResult: false,
+		}, {
+			jmesPath:       "is_external_url('http://[fdb7:d8fd:c4e4:5561::1]')", //182.61.200.7
+			expectedResult: false,
+		}, {
+			jmesPath:       "is_external_url('http://[::1]')", //ipv6 loopback
+			expectedResult: false,
+		}, {
+			jmesPath:       "is_external_url('http://[2001:db8::ff00:42:8329]')",
+			expectedResult: true,
+		}, {
+			jmesPath:       "is_external_url('http://www.google.com@192.168.1.3')", //url with basic auth
+			expectedResult: false,
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.jmesPath, func(t *testing.T) {
+			jp, err := newJMESPath(cfg, tc.jmesPath)
+			assert.NilError(t, err)
+			result, err := jp.Search("")
+			assert.NilError(t, err)
+			res, ok := result.(bool)
+			assert.Assert(t, ok)
+			assert.Equal(t, res, tc.expectedResult)
+		})
+	}
+}