diff --git a/Makefile b/Makefile index 00dc1d0..540c44c 100644 --- a/Makefile +++ b/Makefile @@ -86,5 +86,11 @@ release: echo "tagging new version: $$new_tag"; \ git tag "$$new_tag"; +GOOS ?= $(shell go env GOOS 2>/dev/null || echo linux) +GOARCH ?= $(shell go env GOARCH 2>/dev/null || echo amd64) +wol-proxy: $(BUILD_DIR) + @echo "Building wol-proxy" + go build -o $(BUILD_DIR)/wol-proxy-$(GOOS)-$(GOARCH) cmd/wol-proxy/wol-proxy.go + # Phony targets -.PHONY: all clean ui mac linux windows simple-responder simple-responder-windows test test-all test-dev +.PHONY: all clean ui mac linux windows simple-responder simple-responder-windows test test-all test-dev wol-proxy diff --git a/cmd/wol-proxy/README.md b/cmd/wol-proxy/README.md new file mode 100644 index 0000000..65d2467 --- /dev/null +++ b/cmd/wol-proxy/README.md @@ -0,0 +1,27 @@ +# wol-proxy + +wol-proxy automatically wakes up a suspended llama-swap server using Wake-on-LAN when requests are received. + +When a request arrives and llama-swap is unavailable, wol-proxy sends a WOL packet and holds the request until the server becomes available. If the server doesn't respond within the timeout period (default: 60 seconds), the request is dropped. + +This utility helps conserve energy by allowing GPU-heavy servers to remain suspended when idle, as they can consume hundreds of watts even when not actively processing requests. + +## Usage + +```shell +# minimal +$ ./wol-proxy -mac BA:DC:0F:FE:E0:00 -upstream http://192.168.1.13:8080 + +# everything +$ ./wol-proxy -mac BA:DC:0F:FE:E0:00 -upstream http://192.168.1.13:8080 \ + # use debug log level + -log debug \ + # altenerative listening port + -listen localhost:9999 \ + # seconds to hold requests waiting for upstream to be ready + -timeout 30 +``` + +## API + +`GET /status` - that's it. Everything else is proxied to the upstream server. diff --git a/cmd/wol-proxy/wol-proxy.go b/cmd/wol-proxy/wol-proxy.go new file mode 100644 index 0000000..a002335 --- /dev/null +++ b/cmd/wol-proxy/wol-proxy.go @@ -0,0 +1,250 @@ +package main + +import ( + "context" + "errors" + "flag" + "fmt" + "io" + "log/slog" + "net" + "net/http" + "net/http/httputil" + "net/url" + "os" + "os/signal" + "sync" + "time" +) + +var ( + flagMac = flag.String("mac", "", "mac address to send WoL packet to") + flagUpstream = flag.String("upstream", "", "upstream proxy address to send requests to") + flagListen = flag.String("listen", ":8080", "listen address to listen on") + flagLog = flag.String("log", "info", "log level (debug, info, warn, error)") + flagTimeout = flag.Int("timeout", 60, "seconds requests wait for upstream response before failing") +) + +func main() { + flag.Parse() + + switch *flagLog { + case "debug": + slog.SetLogLoggerLevel(slog.LevelDebug) + case "info": + slog.SetLogLoggerLevel(slog.LevelInfo) + case "warn": + slog.SetLogLoggerLevel(slog.LevelWarn) + case "error": + slog.SetLogLoggerLevel(slog.LevelError) + default: + slog.Error("invalid log level", "logLevel", *flagLog) + return + } + + // Validate flags + if *flagListen == "" { + slog.Error("listen address is required") + return + } + + if *flagMac == "" { + slog.Error("mac address is required") + return + } + + if *flagTimeout < 1 { + slog.Error("timeout must be greater than 0") + return + } + + var upstreamURL *url.URL + var err error + // validate mac address + if _, err = net.ParseMAC(*flagMac); err != nil { + slog.Error("invalid mac address", "error", err) + return + } + + if *flagUpstream == "" { + slog.Error("upstream proxy address is required") + return + } else { + upstreamURL, err = url.ParseRequestURI(*flagUpstream) + if err != nil { + slog.Error("error parsing upstream url", "error", err) + return + } + } + + proxy := newProxy(upstreamURL) + server := &http.Server{ + Addr: *flagListen, + Handler: proxy, + } + + // start the server + go func() { + slog.Info("server starting on", "address", *flagListen) + if err := server.ListenAndServe(); err != nil { + slog.Error("error starting server", "error", err) + } + }() + + // graceful shutdown + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt) + defer stop() + <-ctx.Done() + shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + if err := server.Shutdown(shutdownCtx); err != nil { + slog.Error("server shutdown error", "error", err) + } +} + +type upstreamStatus string + +const ( + notready upstreamStatus = "not ready" + ready upstreamStatus = "ready" +) + +type proxyServer struct { + upstreamProxy *httputil.ReverseProxy + failCount int + statusMutex sync.RWMutex + status upstreamStatus +} + +func newProxy(url *url.URL) *proxyServer { + p := httputil.NewSingleHostReverseProxy(url) + proxy := &proxyServer{ + upstreamProxy: p, + status: notready, + failCount: 0, + } + + // start a goroutien to check upstream status + go func() { + checkUrl := url.Scheme + "://" + url.Host + "/wol-health" + client := &http.Client{Timeout: time.Second} + ticker := time.NewTicker(2 * time.Second) + defer ticker.Stop() + for range ticker.C { + + slog.Debug("checking upstream status at", "url", checkUrl) + resp, err := client.Get(checkUrl) + + // drain the body + if err == nil && resp != nil { + _, _ = io.Copy(io.Discard, resp.Body) + _ = resp.Body.Close() + } + + if err == nil && resp != nil && resp.StatusCode == http.StatusOK { + slog.Debug("upstream status: ready") + proxy.setStatus(ready) + proxy.statusMutex.Lock() + proxy.failCount = 0 + proxy.statusMutex.Unlock() + } else { + slog.Debug("upstream status: notready", "error", err) + proxy.setStatus(notready) + proxy.statusMutex.Lock() + proxy.failCount++ + proxy.statusMutex.Unlock() + } + + } + }() + + return proxy +} + +func (p *proxyServer) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if r.Method == "GET" && r.URL.Path == "/status" { + p.statusMutex.RLock() + status := string(p.status) + failCount := p.failCount + p.statusMutex.RUnlock() + w.Header().Set("Content-Type", "text/plain") + w.WriteHeader(200) + fmt.Fprintf(w, "status: %s\n", status) + fmt.Fprintf(w, "failures: %d\n", failCount) + return + } + + if p.getStatus() == notready { + slog.Info("upstream not ready, sending magic packet", "mac", *flagMac) + if err := sendMagicPacket(*flagMac); err != nil { + slog.Warn("failed to send magic WoL packet", "error", err) + } + ticker := time.NewTicker(250 * time.Millisecond) + timeout, cancel := context.WithTimeout(context.Background(), time.Duration(*flagTimeout)*time.Second) + defer cancel() + loop: + for { + select { + case <-timeout.Done(): + slog.Info("timeout waiting for upstream to be ready") + http.Error(w, "timeout", http.StatusRequestTimeout) + return + case <-ticker.C: + if p.getStatus() == ready { + ticker.Stop() + break loop + } + } + } + } + + p.upstreamProxy.ServeHTTP(w, r) +} + +func (p *proxyServer) getStatus() upstreamStatus { + p.statusMutex.RLock() + defer p.statusMutex.RUnlock() + return p.status +} + +func (p *proxyServer) setStatus(status upstreamStatus) { + p.statusMutex.Lock() + defer p.statusMutex.Unlock() + p.status = status +} + +func sendMagicPacket(macAddr string) error { + hwAddr, err := net.ParseMAC(macAddr) + if err != nil { + return err + } + + if len(hwAddr) != 6 { + return errors.New("invalid MAC address") + } + + // Create the magic packet. + packet := make([]byte, 102) + // Add 6 bytes of 0xFF. + for i := 0; i < 6; i++ { + packet[i] = 0xFF + } + // Repeat the MAC address 16 times. + for i := 1; i <= 16; i++ { + copy(packet[i*6:], hwAddr) + } + + // Send the packet using UDP. + addr := net.UDPAddr{ + IP: net.IPv4bcast, + Port: 9, + } + conn, err := net.DialUDP("udp", nil, &addr) + if err != nil { + return err + } + defer conn.Close() + + _, err = conn.Write(packet) + return err +} diff --git a/proxy/proxymanager.go b/proxy/proxymanager.go index 9ecef92..9888410 100644 --- a/proxy/proxymanager.go +++ b/proxy/proxymanager.go @@ -131,7 +131,15 @@ func New(config config.Config) *ProxyManager { } func (pm *ProxyManager) setupGinEngine() { + pm.ginEngine.Use(func(c *gin.Context) { + + // don't log the Wake on Lan proxy health check + if c.Request.URL.Path == "/wol-health" { + c.Next() + return + } + // Start timer start := time.Now() @@ -235,6 +243,11 @@ func (pm *ProxyManager) setupGinEngine() { c.String(http.StatusOK, "OK") }) + // see cmd/wol-proxy/wol-proxy.go, not logged + pm.ginEngine.GET("/wol-health", func(c *gin.Context) { + c.String(http.StatusOK, "OK") + }) + pm.ginEngine.GET("/favicon.ico", func(c *gin.Context) { if data, err := reactStaticFS.ReadFile("ui_dist/favicon.ico"); err == nil { c.Data(http.StatusOK, "image/x-icon", data)