1
0
forked from wrenn/wrenn
Files
wrenn-releases/internal/hostagent/proxy.go
pptx704 2b4c5e0176 Add pre-pause proxy connection drain and sandbox proxy caching
Introduce ConnTracker (atomic.Bool + WaitGroup) to track in-flight proxy
connections per sandbox. Before pausing a VM, the manager drains active
connections with a 2s grace period, preventing Go runtime corruption
inside the guest caused by stale TCP state surviving Firecracker
snapshot/restore.

Also add:
- AcquireProxyConn on Manager for atomic lookup + connection tracking
- Proxy cache (120s TTL) on CP SandboxProxyWrapper with single-query
  DB lookup (GetSandboxProxyTarget) to avoid two round-trips
- Reset() on ConnTracker to re-enable connections if pause fails
2026-04-01 15:09:44 +06:00

90 lines
2.3 KiB
Go

package hostagent
import (
"fmt"
"log/slog"
"net/http"
"net/http/httputil"
"strconv"
"strings"
"git.omukk.dev/wrenn/sandbox/internal/sandbox"
)
// ProxyHandler reverse-proxies HTTP requests to services running inside
// sandboxes. It handles requests of the form:
//
// /proxy/{sandbox_id}/{port}/{path...}
//
// The sandbox's HostIP (routable on this machine) is used as the upstream.
// This supports any protocol that rides on HTTP, including WebSocket upgrades.
type ProxyHandler struct {
mgr *sandbox.Manager
transport http.RoundTripper
}
// NewProxyHandler creates a new sandbox proxy handler.
func NewProxyHandler(mgr *sandbox.Manager) *ProxyHandler {
return &ProxyHandler{
mgr: mgr,
transport: http.DefaultTransport,
}
}
// ServeHTTP implements http.Handler.
func (h *ProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
// Expected path: /proxy/{sandbox_id}/{port}/...
// After trimming "/proxy/", we get "{sandbox_id}/{port}/..."
trimmed := strings.TrimPrefix(r.URL.Path, "/proxy/")
if trimmed == r.URL.Path {
http.Error(w, "invalid proxy path", http.StatusBadRequest)
return
}
parts := strings.SplitN(trimmed, "/", 3)
if len(parts) < 2 {
http.Error(w, "expected /proxy/{sandbox_id}/{port}/...", http.StatusBadRequest)
return
}
sandboxID := parts[0]
port := parts[1]
remainder := ""
if len(parts) == 3 {
remainder = parts[2]
}
// Validate port is a number in the valid range.
portNum, err := strconv.Atoi(port)
if err != nil || portNum < 1 || portNum > 65535 {
http.Error(w, "invalid port", http.StatusBadRequest)
return
}
hostIP, tracker, ok := h.mgr.AcquireProxyConn(sandboxID)
if !ok {
http.Error(w, "sandbox is not available", http.StatusServiceUnavailable)
return
}
defer tracker.Release()
targetHost := fmt.Sprintf("%s:%d", hostIP, portNum)
proxy := &httputil.ReverseProxy{
Transport: h.transport,
Director: func(req *http.Request) {
req.URL.Scheme = "http"
req.URL.Host = targetHost
req.URL.Path = "/" + remainder
req.URL.RawQuery = r.URL.RawQuery
req.Host = targetHost
},
ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) {
slog.Debug("proxy error", "sandbox_id", sandboxID, "port", port, "error", err)
http.Error(w, "proxy error: "+err.Error(), http.StatusBadGateway)
},
}
proxy.ServeHTTP(w, r)
}