Merge branch 'main' of github.com:e2b-dev/infra into lev-block-cache-bitmap

Author: levb

packages/db/pkg/pool/main.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 
 	"github.com/exaring/otelpgx"
+	"github.com/jackc/pgx/v5"
 	"github.com/jackc/pgx/v5/pgxpool"
 	"go.opentelemetry.io/otel/attribute"
 
@@ -28,6 +29,9 @@ func New(ctx context.Context, databaseURL string, poolName string, options ...Op
 	// expose otel traces
 	config.ConnConfig.Tracer = otelpgx.NewTracer()
 
+	// Disable statement caching to avoid issues with prepared statements in transactions
+	config.ConnConfig.DefaultQueryExecMode = pgx.QueryExecModeExec
+
 	// Create the connection pool
 	pool, err := pgxpool.NewWithConfig(ctx, config)
 	if err != nil {
@@ -41,9 +45,5 @@ func New(ctx context.Context, databaseURL string, poolName string, options ...Op
 		return nil, nil, fmt.Errorf("failed to record stats: %w", err)
 	}
 
-	// TODO [ENG-3437]: Uncomment
-	// Disable statement caching to avoid issues with prepared statements in transactions
-	// config.ConnConfig.DefaultQueryExecMode = pgx.QueryExecModeExec
-
 	return retry.Wrap(pool, retryConfig), pool, nil
 }

packages/db/pkg/types/types.go

@@ -1,7 +1,20 @@
 package types
 
-import "encoding/json"
+import (
+	"database/sql/driver"
+	"encoding/json"
+)
+
+func jsonbValue(v any) (driver.Value, error) {
+	buf, err := json.Marshal(v)
+	if err != nil {
+		return nil, err
+	}
 
+	return string(buf), nil
+}
+
+//nolint:recvcheck // JSONBStringMap needs pointer receiver for unmarshal allocation and value receivers for marshal/driver.Valuer.
 type JSONBStringMap map[string]string
 
 // MarshalJSON ensures a nil map serializes as "{}" instead of "null",
@@ -14,6 +27,10 @@ func (m JSONBStringMap) MarshalJSON() ([]byte, error) {
 	return json.Marshal(map[string]string(m))
 }
 
+func (m JSONBStringMap) Value() (driver.Value, error) {
+	return jsonbValue(m)
+}
+
 // UnmarshalJSON ensures JSON null deserializes as an empty map instead of nil.
 func (m *JSONBStringMap) UnmarshalJSON(data []byte) error {
 	if string(data) == "null" {
@@ -40,6 +57,10 @@ type BuildReason struct {
 	Step *string `json:"step,omitempty"`
 }
 
+func (r BuildReason) Value() (driver.Value, error) {
+	return jsonbValue(r)
+}
+
 const PausedSandboxConfigVersion = "v1"
 
 type SandboxNetworkEgressConfig struct {
@@ -85,6 +106,10 @@ type PausedSandboxConfig struct {
 	VolumeMounts []*SandboxVolumeMountConfig `json:"volumeMounts,omitempty"`
 }
 
+func (c PausedSandboxConfig) Value() (driver.Value, error) {
+	return jsonbValue(c)
+}
+
 // BuildStatus represents the raw status value written to the env_builds table.
 // Use BuildStatusGroup for read-side comparisons.
 type BuildStatus string

packages/db/scripts/migrator.go

@@ -45,6 +45,7 @@ func main() {
 
 		return err
 	}
+	poolConfig.ConnConfig.DefaultQueryExecMode = pgx.QueryExecModeExec
 
 	pool, err := pgxpool.NewWithConfig(ctx, poolConfig)
 	if err != nil {

packages/orchestrator/pkg/hyperloopserver/handlers/logs.go

@@ -3,6 +3,7 @@ package handlers
 import (
 	"bytes"
 	"encoding/json"
+	"fmt"
 	"net"
 	"net/http"
 
@@ -33,6 +34,14 @@ func (h *APIStore) Logs(c *gin.Context) {
 		return
 	}
 
+	err = h.validatePayloadSandboxID(payload, sbxID)
+	if err != nil {
+		h.sendAPIStoreError(c, http.StatusBadRequest, "Invalid sandboxID in logs payload")
+		h.logger.Error(ctx, "error when parsing sandbox logs request", zap.Error(err), logger.WithSandboxID(sbxID))
+
+		return
+	}
+
 	// Overwrite instanceID and teamID to avoid spoofing
 	payload["instanceID"] = sbxID
 	payload["teamID"] = sbx.Runtime.TeamID
@@ -65,3 +74,21 @@ func (h *APIStore) Logs(c *gin.Context) {
 
 	c.Status(http.StatusOK)
 }
+
+// validatePayloadSandboxID checks if the payload contains correct instanceID to prevent slow requests to contaminating the logs of other sandboxes.
+func (h *APIStore) validatePayloadSandboxID(payload map[string]any, sbxID string) error {
+	if payload["instanceID"] == nil {
+		return fmt.Errorf("missing sandboxID in logs payload")
+	}
+
+	payloadSandboxID, ok := payload["instanceID"].(string)
+	if !ok {
+		return fmt.Errorf("instanceID in logs payload is not a string: %v", payload["instanceID"])
+	}
+
+	if payloadSandboxID != sbxID {
+		return fmt.Errorf("sandboxID in logs payload does not match the sandboxID of the source sandbox (%s != %s)", payloadSandboxID, sbxID)
+	}
+
+	return nil
+}