refactor once again

pool-server.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"io"
 	"log"
 	"math/rand"
 	"net/http"
@@ -25,12 +26,12 @@ func NewPoolServer(pool GrainPool, pod_name string) *PoolServer {
 	}
 }
 
-func (s *PoolServer) process(id CartId, mutation interface{}) (*messages.CartState, error) {
+func (s *PoolServer) process(id CartId, mutation interface{}) (*CartGrain, error) {
 	grain, err := s.pool.Apply(id, mutation)
 	if err != nil {
 		return nil, err
 	}
-	return ToCartState(grain), nil
+	return grain, nil
 }
 
 func (s *PoolServer) HandleGet(w http.ResponseWriter, r *http.Request, id CartId) error {
@@ -39,7 +40,7 @@ func (s *PoolServer) HandleGet(w http.ResponseWriter, r *http.Request, id CartId
 		return err
 	}
 
-	return s.WriteResult(w, ToCartState(grain))
+	return s.WriteResult(w, grain)
 }
 
 func (s *PoolServer) HandleAddSku(w http.ResponseWriter, r *http.Request, id CartId) error {
@@ -62,7 +63,7 @@ func ErrorHandler(fn func(w http.ResponseWriter, r *http.Request) error) func(w
 	}
 }
 
-func (s *PoolServer) WriteResult(w http.ResponseWriter, result *messages.CartState) error {
+func (s *PoolServer) WriteResult(w http.ResponseWriter, result *CartGrain) error {
 	w.Header().Set("Content-Type", "application/json")
 	w.Header().Set("Cache-Control", "no-cache")
 	w.Header().Set("Access-Control-Allow-Origin", "*")
@@ -280,11 +281,9 @@ func NewCartId() CartId {
 
 func CookieCartIdHandler(fn func(w http.ResponseWriter, r *http.Request, cartId CartId) error) func(w http.ResponseWriter, r *http.Request) error {
 	return func(w http.ResponseWriter, r *http.Request) error {
-		// Extract / normalize cookie (preserve legacy textual IDs without rewriting).
 		var legacy CartId
 		cookies := r.CookiesNamed("cartid")
 		if len(cookies) == 0 {
-			// No cookie -> generate new canonical base62 id.
 			cid, generated, _, err := CanonicalizeOrLegacy("")
 			if err != nil {
 				return fmt.Errorf("failed to generate cart id: %w", err)
@@ -309,8 +308,6 @@ func CookieCartIdHandler(fn func(w http.ResponseWriter, r *http.Request, cartId
 				return fmt.Errorf("failed to canonicalize cart id: %w", err)
 			}
 			legacy = CartIDToLegacy(cid)
-			// Only set a new cookie if we actually generated a brand-new ID (empty input).
-			// For legacy (non-base62) ids we preserve the original text and do not overwrite.
 			if generated && wasBase62 {
 				http.SetCookie(w, &http.Cookie{
 					Name:     "cartid",
@@ -324,6 +321,12 @@ func CookieCartIdHandler(fn func(w http.ResponseWriter, r *http.Request, cartId
 				w.Header().Set("Set-Cart-Id", cid.String())
 			}
 		}
+		// Ownership proxy AFTER id extraction (cookie mode)
+		if ownershipProxyAfterExtraction != nil {
+			if handled, err := ownershipProxyAfterExtraction(legacy, w, r); handled || err != nil {
+				return err
+			}
+		}
 		return fn(w, r, legacy)
 	}
 }
@@ -351,18 +354,90 @@ func CartIdHandler(fn func(w http.ResponseWriter, r *http.Request, cartId CartId
 			return fmt.Errorf("invalid cart id: %w", err)
 		}
 		legacy := CartIDToLegacy(cid)
-		// Only emit Set-Cart-Id header if we produced a brand-new canonical id
-		// AND it is base62 (avoid rewriting legacy textual identifiers).
 		if generated && wasBase62 {
 			w.Header().Set("Set-Cart-Id", cid.String())
 		}
+		// Ownership proxy AFTER path id extraction (explicit id mode)
+		if ownershipProxyAfterExtraction != nil {
+			if handled, err := ownershipProxyAfterExtraction(legacy, w, r); handled || err != nil {
+				return err
+			}
+		}
 		return fn(w, r, legacy)
 	}
 }
 
+var ownershipProxyAfterExtraction func(cartId CartId, w http.ResponseWriter, r *http.Request) (handled bool, err error)
+
 func (s *PoolServer) Serve() *http.ServeMux {
+	// Install ownership proxy hook that runs AFTER id extraction (cookie OR path)
+	ownershipProxyAfterExtraction = func(cartId CartId, w http.ResponseWriter, r *http.Request) (bool, error) {
+		if cartId.String() == "" {
+			return false, nil
+		}
+		owner := s.pool.OwnerHost(cartId)
+		if owner == "" || owner == s.pool.Hostname() {
+			// Set / refresh cartowner cookie pointing to the local host (claim or already owned).
+			localHost := owner
+			if localHost == "" {
+				localHost = s.pool.Hostname()
+			}
+			http.SetCookie(w, &http.Cookie{
+				Name:     "cartowner",
+				Value:    localHost,
+				Path:     "/",
+				HttpOnly: true,
+				SameSite: http.SameSiteLaxMode,
+			})
+			return false, nil
+		}
+		// For remote ownership set cartowner cookie to remote host for sticky sessions.
+		http.SetCookie(w, &http.Cookie{
+			Name:     "cartowner",
+			Value:    owner,
+			Path:     "/",
+			HttpOnly: true,
+			SameSite: http.SameSiteLaxMode,
+		})
+		// Proxy logic (simplified): reuse existing request to owning host on same port.
+		target := "http://" + owner + r.URL.Path
+		if q := r.URL.RawQuery; q != "" {
+			target += "?" + q
+		}
+		req, err := http.NewRequestWithContext(r.Context(), r.Method, target, r.Body)
+		if err != nil {
+			http.Error(w, "proxy build error", http.StatusBadGateway)
+			return true, err
+		}
+		for k, v := range r.Header {
+			for _, vv := range v {
+				req.Header.Add(k, vv)
+			}
+		}
+		req.Header.Set("X-Forwarded-Host", r.Host)
+		req.Header.Set("X-Cart-Id", cartId.String())
+		req.Header.Set("X-Cart-Owner", owner)
+		resp, err := http.DefaultClient.Do(req)
+		if err != nil {
+			http.Error(w, "proxy upstream error", http.StatusBadGateway)
+			return true, err
+		}
+		defer resp.Body.Close()
+		for k, v := range resp.Header {
+			for _, vv := range v {
+				w.Header().Add(k, vv)
+			}
+		}
+		w.Header().Set("X-Cart-Owner-Routed", "true")
+		w.WriteHeader(resp.StatusCode)
+		_, copyErr := io.Copy(w, resp.Body)
+		if copyErr != nil {
+			return true, copyErr
+		}
+		return true, nil
+	}
+
 	mux := http.NewServeMux()
-	//mux.HandleFunc("/", s.RewritePath)
 	mux.HandleFunc("OPTIONS /", func(w http.ResponseWriter, r *http.Request) {
 		w.Header().Set("Access-Control-Allow-Origin", "*")
 		w.Header().Set("Access-Control-Allow-Methods", "GET, PUT, POST, DELETE")