[fix] custom http module loader (#88)

* better analysis

* update comments

* for http tests

* test func

* clean it

* another test

* not allowed

* fix for windows"
g s
g "

* better refacotr

Author: hyorigo

.github/workflows/build.yml

@@ -68,17 +68,25 @@ jobs:
           CODACY_PROJECT_TOKEN: ${{ secrets.CODACY_PROJECT_TOKEN }}
         run:
           bash <(curl -Ls https://coverage.codacy.com/get.sh) report --force-coverage-parser go -r coverage.txt
-      - name: Analyze
+      - name: Analyze on Linux
         if: ${{ runner.os == 'Linux' }}
         run: |
-          # tokei -- count
+          # Setup
           wget -cqL https://github.com/XAMPPRocky/tokei/releases/download/v12.1.2/tokei-i686-unknown-linux-musl.tar.gz -O tokei.tgz
           tar zxf tokei.tgz tokei && chmod +x tokei && $SUDO mv tokei /usr/local/bin && rm tokei.tgz
-          echo "=== Tokei Result ==="
-          tokei
-          # revive -- lint
-          wget -cqL https://github.com/mgechev/revive/releases/download/v1.2.4/revive_1.2.4_Linux_x86_64.tar.gz -O revive.tgz
+          wget -cqL https://github.com/mgechev/revive/releases/download/v1.3.7/revive_linux_amd64.tar.gz -O revive.tgz
           tar zxf revive.tgz revive && chmod +x revive && $SUDO mv revive /usr/local/bin && rm revive.tgz
-          wget -cqL https://bitbucket.org/ai69/common/raw/master/revive.toml -O revive.toml
-          echo "=== Revive Result ==="
-          revive -config revive.toml -formatter friendly ./...
+          wget -cqL https://raw.githubusercontent.com/1set/meta/master/revive.toml -O revive.toml
+          # Analyze
+          echo "# Analysis on Linux" > $GITHUB_STEP_SUMMARY
+          uname -a >> $GITHUB_STEP_SUMMARY
+          # --- count lines of code
+          echo "## Tokei Result" >> $GITHUB_STEP_SUMMARY
+          printf '\n```\n' >> $GITHUB_STEP_SUMMARY
+          tokei >> $GITHUB_STEP_SUMMARY
+          printf '```\n\n' >> $GITHUB_STEP_SUMMARY
+          # --- lint
+          echo "## Revive Result" >> $GITHUB_STEP_SUMMARY
+          printf '\n```\n' >> $GITHUB_STEP_SUMMARY
+          revive -config revive.toml -formatter friendly ./... >> $GITHUB_STEP_SUMMARY
+          printf '```\n\n' >> $GITHUB_STEP_SUMMARY

lib/http/http.go

@@ -50,24 +50,42 @@ type RequestGuard interface {
 
 // LoadModule creates an http Module
 func LoadModule() (starlark.StringDict, error) {
-	var m = &Module{}
+	return NewModule().LoadModule()
+}
+
+// Module defines the actual HTTP module with methods for making requests.
+type Module struct {
+	cli *http.Client
+	rg  RequestGuard
+}
+
+// NewModule creates a new http module with default settings.
+func NewModule() *Module {
+	m := &Module{}
 	if Client != nil {
 		m.cli = Client
 	}
 	if Guard != nil {
 		m.rg = Guard
 	}
-	ns := starlark.StringDict{
-		ModuleName: m.Struct(),
-	}
-	return ns, nil
+	return m
 }
 
-// Module joins http tools to a dataset, allowing dataset
-// to follow along with http requests
-type Module struct {
-	cli *http.Client
-	rg  RequestGuard
+// SetClient sets the http client for this module, useful for setting custom clients for testing or multiple loadings.
+func (m *Module) SetClient(c *http.Client) {
+	m.cli = c
+}
+
+// SetGuard sets the request guard for this module, useful for setting custom guards for testing or multiple loadings.
+func (m *Module) SetGuard(g RequestGuard) {
+	m.rg = g
+}
+
+// LoadModule creates an http Module.
+func (m *Module) LoadModule() (starlark.StringDict, error) {
+	return starlark.StringDict{
+		ModuleName: m.Struct(),
+	}, nil
 }
 
 // Struct returns this module's methods as a starlark Struct

lib/http/http_test.go

@@ -1,16 +1,17 @@
-package http
+package http_test
 
 import (
 	"encoding/json"
-	"io/ioutil"
+	"errors"
 	"net/http"
 	"net/http/httptest"
 	"net/http/httputil"
-	"strings"
+	"net/url"
 	"testing"
 	"time"
 
 	itn "github.com/1set/starlet/internal"
+	lh "github.com/1set/starlet/lib/http"
 	"github.com/1set/starlight/convert"
 	"go.starlark.net/starlark"
 	"go.starlark.net/starlarktest"
@@ -27,7 +28,7 @@ func TestAsString(t *testing.T) {
 	}
 
 	for i, c := range cases {
-		got, err := AsString(c.in)
+		got, err := lh.AsString(c.in)
 		if !(err == nil && c.err == "" || err != nil && err.Error() == c.err) {
 			t.Errorf("case %d error mismatch. expected: '%s', got: '%s'", i, c.err, err)
 			continue
@@ -49,7 +50,7 @@ func TestLoadModule_HTTP_One(t *testing.T) {
 	defer ts.Close()
 	starlark.Universe["test_server_url"] = starlark.String(ts.URL)
 
-	thread := &starlark.Thread{Load: itn.NewAssertLoader(ModuleName, LoadModule)}
+	thread := &starlark.Thread{Load: itn.NewAssertLoader(lh.ModuleName, lh.LoadModule)}
 	starlarktest.SetReporter(thread, t)
 
 	code := itn.HereDoc(`
@@ -296,7 +297,7 @@ func TestLoadModule_HTTP(t *testing.T) {
 		{
 			name: `POST with UA Set`,
 			preset: func() {
-				UserAgent = "GqQdYX3eIJw2DTt"
+				lh.UserAgent = "GqQdYX3eIJw2DTt"
 			},
 			script: itn.HereDoc(`
 				load('http', 'post')
@@ -575,8 +576,8 @@ func TestLoadModule_HTTP(t *testing.T) {
 			if tt.preset != nil {
 				tt.preset()
 			}
-			TimeoutSecond = 30.0
-			res, err := itn.ExecModuleWithErrorTest(t, ModuleName, LoadModule, tt.script, tt.wantErr, nil)
+			lh.TimeoutSecond = 30.0
+			res, err := itn.ExecModuleWithErrorTest(t, lh.ModuleName, lh.LoadModule, tt.script, tt.wantErr, nil)
 			if (err != nil) != (tt.wantErr != "") {
 				t.Errorf("http(%q) expects error = '%v', actual error = '%v', result = %v", tt.name, tt.wantErr, err, res)
 				return
@@ -585,65 +586,91 @@ func TestLoadModule_HTTP(t *testing.T) {
 	}
 }
 
-// we're ok with testing private functions if it simplifies the test :)
-func TestSetBody(t *testing.T) {
-	fd := map[string]string{
-		"foo": "bar baz",
+// DomainWhitelistGuard allows requests only to domains in its whitelist.
+type DomainWhitelistGuard struct {
+	whitelist map[string]struct{} // Set of allowed domains
+}
+
+// NewDomainWhitelistGuard creates a new DomainWhitelistGuard with the specified domains.
+func NewDomainWhitelistGuard(domains []string) *DomainWhitelistGuard {
+	whitelist := make(map[string]struct{})
+	for _, domain := range domains {
+		whitelist[domain] = struct{}{}
 	}
+	return &DomainWhitelistGuard{whitelist: whitelist}
+}
 
-	cases := []struct {
-		rawBody      starlark.String
-		formData     map[string]string
-		formEncoding starlark.String
-		jsonData     starlark.Value
-		body         string
-		err          string
-	}{
-		{starlark.String("hallo"), nil, starlark.String(""), nil, "hallo", ""},
-		{starlark.String(""), fd, starlark.String(""), nil, "foo=bar+baz", ""},
-		// TODO - this should check multipart form data is being set
-		{starlark.String(""), fd, starlark.String("multipart/form-data"), nil, "", ""},
-		{starlark.String(""), nil, starlark.String(""), starlark.Tuple{starlark.Bool(true), starlark.MakeInt(1), starlark.String("der")}, "[true,1,\"der\"]", ""},
+// Allowed checks if the request's domain is in the whitelist.
+func (g *DomainWhitelistGuard) Allowed(thread *starlark.Thread, req *http.Request) (*http.Request, error) {
+	if _, ok := g.whitelist[req.URL.Host]; ok {
+		// Domain is in the whitelist, allow the request
+		return req, nil
 	}
+	// Domain is not in the whitelist, deny the request
+	return nil, errors.New("request to this domain is not allowed")
+}
 
-	for i, c := range cases {
-		var formData *starlark.Dict
-		if c.formData != nil {
-			formData = starlark.NewDict(len(c.formData))
-			for k, v := range c.formData {
-				if err := formData.SetKey(starlark.String(k), starlark.String(v)); err != nil {
-					t.Fatal(err)
-				}
-			}
-		}
+func TestLoadModule_CustomLoad(t *testing.T) {
+	md := lh.NewModule()
+	proxyURL, _ := url.Parse("http://127.0.0.1:9999")
+	client := &http.Client{
+		Transport: &http.Transport{
+			Proxy: http.ProxyURL(proxyURL),
+		},
+	}
+	md.SetClient(client)
+	guard := NewDomainWhitelistGuard([]string{"allowed.com"})
+	md.SetGuard(guard)
 
-		req := httptest.NewRequest("get", "https://example.com", nil)
-		err := setBody(req, c.rawBody, formData, c.formEncoding, c.jsonData)
-		if !(err == nil && c.err == "" || (err != nil && err.Error() == c.err)) {
-			t.Errorf("case %d error mismatch. expected: %s, got: %s", i, c.err, err)
-			continue
+	httpHand := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		b, err := httputil.DumpRequest(r, true)
+		if err != nil {
+			t.Errorf("Error dumping request: %v", err)
 		}
+		t.Logf("Web server received request: [[%s]]", b)
+		time.Sleep(10 * time.Millisecond)
+		w.Write(b)
+	})
+	ts := httptest.NewServer(httpHand)
+	defer ts.Close()
 
-		if strings.HasPrefix(req.Header.Get("Content-Type"), "multipart/form-data;") {
-			if err := req.ParseMultipartForm(0); err != nil {
-				t.Fatal(err)
-			}
-
-			for k, v := range c.formData {
-				fv := req.FormValue(k)
-				if fv != v {
-					t.Errorf("case %d error mismatch. expected %s=%s, got: %s", i, k, v, fv)
-				}
-			}
-		} else {
-			body, err := ioutil.ReadAll(req.Body)
-			if err != nil {
-				t.Fatal(err)
+	tests := []struct {
+		name    string
+		preset  func()
+		script  string
+		wantErr string
+	}{
+		{
+			name: `Simple GET`,
+			script: itn.HereDoc(`
+				load('http', 'get')
+				res = get("http://allowed.com/hello")
+				assert.eq(res.status_code, 200)
+			`),
+			wantErr: `proxyconnect tcp: dial tcp 127.0.0.1:9999: connect`,
+		},
+		{
+			name: `Not Allowed`,
+			script: itn.HereDoc(`
+				load('http', 'get')
+				res = get("http://topsecret.com/text")
+				assert.eq(res.status_code, 200)
+			`),
+			wantErr: `request to this domain is not allowed`,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if tt.preset != nil {
+				tt.preset()
 			}
-
-			if string(body) != c.body {
-				t.Errorf("case %d body mismatch. expected: %s, got: %s", i, c.body, string(body))
+			res, err := itn.ExecModuleWithErrorTest(t, lh.ModuleName, md.LoadModule, tt.script, tt.wantErr, starlark.StringDict{
+				"test_server_url": starlark.String(ts.URL),
+			})
+			if (err != nil) != (tt.wantErr != "") {
+				t.Errorf("http(%q) expects error = '%v', actual error = '%v', result = %v", tt.name, tt.wantErr, err, res)
+				return
 			}
-		}
+		})
 	}
 }

lib/http/internal_test.go

@@ -0,0 +1,73 @@
+package http
+
+import (
+	"io/ioutil"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"go.starlark.net/starlark"
+)
+
+// we're ok with testing private functions if it simplifies the test :)
+func TestSetBody(t *testing.T) {
+	fd := map[string]string{
+		"foo": "bar baz",
+	}
+
+	cases := []struct {
+		rawBody      starlark.String
+		formData     map[string]string
+		formEncoding starlark.String
+		jsonData     starlark.Value
+		body         string
+		err          string
+	}{
+		{starlark.String("hallo"), nil, starlark.String(""), nil, "hallo", ""},
+		{starlark.String(""), fd, starlark.String(""), nil, "foo=bar+baz", ""},
+		// TODO - this should check multipart form data is being set
+		{starlark.String(""), fd, starlark.String("multipart/form-data"), nil, "", ""},
+		{starlark.String(""), nil, starlark.String(""), starlark.Tuple{starlark.Bool(true), starlark.MakeInt(1), starlark.String("der")}, "[true,1,\"der\"]", ""},
+	}
+
+	for i, c := range cases {
+		var formData *starlark.Dict
+		if c.formData != nil {
+			formData = starlark.NewDict(len(c.formData))
+			for k, v := range c.formData {
+				if err := formData.SetKey(starlark.String(k), starlark.String(v)); err != nil {
+					t.Fatal(err)
+				}
+			}
+		}
+
+		req := httptest.NewRequest("get", "https://example.com", nil)
+		err := setBody(req, c.rawBody, formData, c.formEncoding, c.jsonData)
+		if !(err == nil && c.err == "" || (err != nil && err.Error() == c.err)) {
+			t.Errorf("case %d error mismatch. expected: %s, got: %s", i, c.err, err)
+			continue
+		}
+
+		if strings.HasPrefix(req.Header.Get("Content-Type"), "multipart/form-data;") {
+			if err := req.ParseMultipartForm(0); err != nil {
+				t.Fatal(err)
+			}
+
+			for k, v := range c.formData {
+				fv := req.FormValue(k)
+				if fv != v {
+					t.Errorf("case %d error mismatch. expected %s=%s, got: %s", i, k, v, fv)
+				}
+			}
+		} else {
+			body, err := ioutil.ReadAll(req.Body)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			if string(body) != c.body {
+				t.Errorf("case %d body mismatch. expected: %s, got: %s", i, c.body, string(body))
+			}
+		}
+	}
+}