Make HTTP sessions work

runtime/httpcache.go

@@ -9,11 +9,13 @@ import (
 	"fmt"
 	"math/rand"
 	"net/http"
+	"net/http/cookiejar"
 	"net/http/httputil"
 	"strconv"
 	"strings"
 	"time"
 
+	"golang.org/x/net/publicsuffix"
 	"tidbyt.dev/pixlet/runtime/modules/starlarkhttp"
 )
 
@@ -54,11 +56,22 @@ func InitHTTP(cache Cache) {
 		transport: http.DefaultTransport,
 	}
 
-	httpClient := &http.Client{
-		Transport: cc,
-		Timeout:   HTTPTimeout * 2,
+	starlarkhttp.StarlarkHTTPClient = func() *http.Client {
+		// Providing a cookie jar allows sessions and redirects to work properly. With a
+		// jar present, any cookies set in a response will automatically be added to
+		// subsequent requests. This means that we can follow redirects after logging into
+		// a session. Without a jar, any cookies will be dropped from redirects unless explicitly
+		// set in the original outgoing request.
+		// https://cs.opensource.google/go/go/+/master:src/net/http/client.go;drc=4c394b5638cc2694b1eff6418bc3e7db8132de0e;l=88
+		jar, _ := cookiejar.New(&cookiejar.Options{PublicSuffixList: publicsuffix.List}) // never returns non-nil err
+
+		httpClient := &http.Client{
+			Jar:       jar,
+			Transport: cc,
+			Timeout:   HTTPTimeout * 2,
+		}
+		return httpClient
 	}
-	starlarkhttp.StarlarkHTTPClient = httpClient
 }
 
 // RoundTrip is an approximation of what our internal HTTP proxy does. It should

runtime/httpcache_test.go

@@ -5,11 +5,14 @@ import (
 	"fmt"
 	"math/rand"
 	"net/http"
+	"net/http/httptest"
 	"os"
+	"strings"
 	"testing"
 	"time"
 
 	"github.com/stretchr/testify/assert"
+	"go.starlark.net/starlark"
 )
 
 func TestInitHTTP(t *testing.T) {
@@ -183,3 +186,54 @@ func TestDetermineTTLNoHeaders(t *testing.T) {
 	ttl := DetermineTTL(req, res)
 	assert.Equal(t, MinRequestTTL, ttl)
 }
+
+func TestSetCookieOnRedirect(t *testing.T) {
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Requests to "/login" set a cookie and redirect to /destination
+		if strings.HasSuffix(r.URL.Path, "/login") {
+			if len(r.Cookies()) != 0 {
+				t.Errorf("Cookie was already set on initial call")
+			}
+			w.Header().Set("Set-Cookie", "doodad=foobar; path=/; HttpOnly")
+			w.Header().Set("Location", "/destination")
+			w.WriteHeader(302)
+			return
+		}
+		// Requests to /destination must have cookie set
+		if strings.HasSuffix(r.URL.Path, "/destination") {
+			c, err := r.Cookie("doodad")
+			if err != nil {
+				t.Errorf("Expected cookie `doodad` not present") // Occurs if client has no cookie jar
+			}
+			if c.Value != "foobar" {
+				t.Errorf("Cookie `doodad` value mismatch. Expected foobar, got %s", c.Value)
+			}
+			if _, err := w.Write([]byte(`{"hello":"world"}`)); err != nil {
+				t.Fatal(err)
+			}
+			return
+		}
+		t.Errorf("Unexpected path requested: %s", r.URL.Path)
+	}))
+
+	starlark.Universe["test_server_url"] = starlark.String(ts.URL)
+	c := NewInMemoryCache()
+	InitHTTP(c)
+
+	b, err := os.ReadFile("testdata/httpredirect.star")
+	assert.NoError(t, err)
+
+	app, err := NewApplet("httpredirect.star", b)
+	assert.NoError(t, err)
+
+	_, err = app.Run(context.Background())
+	assert.NoError(t, err)
+
+	// Run it again and check that we're not using the same cookie jar
+	// across executions. If we were, the first request would error out.
+	app2, err := NewApplet("httpredirect.star", b)
+	assert.NoError(t, err)
+
+	_, err = app2.Run(context.Background())
+	assert.NoError(t, err)
+}

runtime/modules/starlarkhttp/starlarkhttp.go

@@ -51,9 +51,11 @@ func AsString(x starlark.Value) (string, error) {
 const ModuleName = "http.star"
 
 var (
-	// StarlarkHTTPClient is the http client used to create the http module. override with
-	// a custom client before calling LoadModule
-	StarlarkHTTPClient = http.DefaultClient
+	// StarlarkHTTPClient is a factory method for creating the http client used to create the http module.
+	// override with a custom function before calling LoadModule
+	StarlarkHTTPClient = func() *http.Client {
+		return http.DefaultClient
+	}
 	// StarlarkHTTPGuard is a global RequestGuard used in LoadModule. override with a custom
 	// implementation before calling LoadModule
 	StarlarkHTTPGuard RequestGuard
@@ -69,7 +71,7 @@ const (
 
 // LoadModule creates an http Module
 func LoadModule() (starlark.StringDict, error) {
-	var m = &Module{cli: StarlarkHTTPClient}
+	var m = &Module{cli: StarlarkHTTPClient()}
 	if StarlarkHTTPGuard != nil {
 		m.rg = StarlarkHTTPGuard
 	}

runtime/testdata/httpredirect.star

@@ -0,0 +1,10 @@
+load("assert.star", "assert")
+load("http.star", "http")
+load("render.star", "render")
+
+def main():
+    res_1 = http.post(test_server_url + "/login")
+    assert.eq(res_1.status_code, 200)
+    assert.eq(res_1.body(), '{"hello":"world"}')
+    assert.eq(res_1.json(), {"hello": "world"})
+    return render.Root(child = render.Text("pass"))