From df160411d5ec7f8a19435afabc4862d3e9c33d45 Mon Sep 17 00:00:00 2001
From: dinosaursrarr <tom+github@chamberlaincurtis.com>
Date: Sat, 24 Feb 2024 16:42:09 +0000
Subject: [PATCH] Make sessions work

User simonahac reported a thorny problem on discord yesterday. See discussion beginning at https://discord.com/channels/928484660785336380/928485908842426389/1210584191142723615

They want to use a service that relies on sessions. First you POST to a login URL. That sets a cookie and redirects to a resource. The resource only works if you provide the cookie value in the GET request.

This works out of the box in python, Postman, and most other clients. But not pixlet.

Looks like the underlying cause is the default behaviour of golang's http.Client.

By default, it will ignore any cookies set in responses when making the redirect. When sending the redirected request, it would just copy over headers and cookies from the original outgoing response -- which won't have the cookie set.

If the client has a CookieJar, however, it just works as expected. Cookies from responses will be added to the jar, and all the cookies from the jar will automatically be added to all outgoing requests from that point.

https://cs.opensource.google/go/go/+/master:src/net/http/client.go;drc=b899e0b8cc1afc4534758c9ebe1e051e5220bfbd;l=88

I think this should be a safe change. There's only one existing community app (avatarsinpixels) that manually accesses the "Set-Cookie" header on a response. I've checked it works the same with and without this change.
---
 runtime/httpcache.go               | 11 ++++++++
 runtime/httpcache_test.go          | 43 ++++++++++++++++++++++++++++++
 runtime/testdata/httpredirect.star | 10 +++++++
 3 files changed, 64 insertions(+)
 create mode 100644 runtime/testdata/httpredirect.star

diff --git a/runtime/httpcache.go b/runtime/httpcache.go
index 6c77cef996..17e7ca346a 100644
--- a/runtime/httpcache.go
+++ b/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,7 +56,16 @@ func InitHTTP(cache Cache) {
 		transport: http.DefaultTransport,
 	}
 
+	// 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,
 	}
diff --git a/runtime/httpcache_test.go b/runtime/httpcache_test.go
index d4940078a4..9d33ea1d37 100644
--- a/runtime/httpcache_test.go
+++ b/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,43 @@ 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") {
+			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)
+}
diff --git a/runtime/testdata/httpredirect.star b/runtime/testdata/httpredirect.star
new file mode 100644
index 0000000000..a7236b081d
--- /dev/null
+++ b/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"))