diff --git a/runtime/httpcache.go b/runtime/httpcache.go index 6c77cef996..55710cc207 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,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 diff --git a/runtime/httpcache_test.go b/runtime/httpcache_test.go index d4940078a4..79492da4fe 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,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) +} diff --git a/runtime/modules/starlarkhttp/starlarkhttp.go b/runtime/modules/starlarkhttp/starlarkhttp.go index 5fb8764e21..03c3132d03 100644 --- a/runtime/modules/starlarkhttp/starlarkhttp.go +++ b/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 } 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"))