From 8b2d0eab323425c724f46749a93a34e0296ad073 Mon Sep 17 00:00:00 2001
From: Ferran Borreguero <ferran.borreguero@gmail.com>
Date: Wed, 15 May 2024 12:16:51 +0100
Subject: [PATCH] Add aws lb cookie settings for session id

---
 core/vm/contracts_suave.go      | 17 +++++++++++
 core/vm/contracts_suave_test.go | 54 +++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+)

diff --git a/core/vm/contracts_suave.go b/core/vm/contracts_suave.go
index a366c8a48..92d7bda28 100644
--- a/core/vm/contracts_suave.go
+++ b/core/vm/contracts_suave.go
@@ -210,6 +210,8 @@ func (c *consoleLogPrecompile) Run(input []byte) ([]byte, error) {
 	return nil, nil
 }
 
+var contextCookieKeyPrefix = "__cookie_"
+
 func (s *suaveRuntime) doHTTPRequest(request types.HttpRequest) ([]byte, error) {
 	if request.Method != "GET" && request.Method != "POST" {
 		return nil, fmt.Errorf("only GET and POST methods are supported")
@@ -232,6 +234,13 @@ func (s *suaveRuntime) doHTTPRequest(request types.HttpRequest) ([]byte, error)
 		return nil, err
 	}
 
+	// add any cookies stored in the context
+	for key, val := range s.suaveContext.Context {
+		if strings.HasPrefix(key, contextCookieKeyPrefix) {
+			req.Header.Add("Cookie", string(val))
+		}
+	}
+
 	for _, header := range request.Headers {
 		indx := strings.Index(header, ":")
 		if indx == -1 {
@@ -269,6 +278,14 @@ func (s *suaveRuntime) doHTTPRequest(request types.HttpRequest) ([]byte, error)
 	if resp.StatusCode > 299 {
 		return nil, fmt.Errorf("http error: %s: %v", resp.Status, data)
 	}
+
+	// parse the LB cookies (AWSALB, AWSALBCORS) and set them in the context
+	for _, cookie := range resp.Cookies() {
+		if cookie.Name == "AWSALB" || cookie.Name == "AWSALBCORS" {
+			s.suaveContext.Context[contextCookieKeyPrefix+cookie.Name] = []byte(cookie.String())
+		}
+	}
+
 	return data, nil
 }
 
diff --git a/core/vm/contracts_suave_test.go b/core/vm/contracts_suave_test.go
index cc0c0d3fb..4e5cc0f45 100644
--- a/core/vm/contracts_suave_test.go
+++ b/core/vm/contracts_suave_test.go
@@ -321,3 +321,57 @@ func TestSuave_HttpRequest_FlashbotsSignatue(t *testing.T) {
 	_, err := s.doHTTPRequest(req)
 	require.NoError(t, err)
 }
+
+func TestSuave_HttpRequest_Cookies(t *testing.T) {
+	cookies := map[string]http.Cookie{
+		"AWSALB":     {Name: "AWSALB", Value: "value1"},
+		"AWSALBCORS": {Name: "AWSALBCORS", Value: "value2"},
+		"OTHER":      {Name: "OTHER", Value: "value3"},
+	}
+
+	firstCall := true
+	srv := httptest.NewServer(&httpTestHandler{
+		fn: func(w http.ResponseWriter, r *http.Request) {
+			if firstCall {
+				firstCall = false
+				for _, c := range cookies {
+					http.SetCookie(w, &c)
+				}
+			} else {
+				// check the cookies in the second call
+				for _, c := range r.Cookies() {
+					if val, found := cookies[c.Name]; found {
+						require.Equal(t, val.Value, c.Value)
+					}
+				}
+			}
+			w.Write([]byte("ok"))
+		},
+	})
+
+	s := &suaveRuntime{
+		suaveContext: &SuaveContext{
+			Context: map[string][]byte{},
+			Backend: &SuaveExecutionBackend{
+				ExternalWhitelist:    []string{"127.0.0.1"},
+				ServiceAliasRegistry: map[string]string{"goerli": srv.URL},
+			},
+		},
+	}
+
+	defer srv.Close()
+
+	req := types.HttpRequest{Url: srv.URL, Method: "GET"}
+	_, err := s.doHTTPRequest(req)
+	require.NoError(t, err)
+
+	// validate the only the AWS cookies are stored
+	require.Len(t, s.suaveContext.Context, 2)
+	for key, val := range s.suaveContext.Context {
+		require.True(t, strings.HasPrefix(key, contextCookieKeyPrefix))
+		require.Contains(t, []string{"AWSALB", "AWSALBCORS"}, strings.Split(string(val), "=")[0])
+	}
+
+	_, err = s.doHTTPRequest(req)
+	require.NoError(t, err)
+}