Skip to content
Merged
Show file tree
Hide file tree
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 7 additions & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -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
27 changes: 27 additions & 0 deletions cmd/wol-proxy/README.md
Original file line number Diff line number Diff line change
@@ -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.
249 changes: 249 additions & 0 deletions cmd/wol-proxy/wol-proxy.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,249 @@
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)
}

// 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" {
w.Header().Set("Content-Type", "text/plain")
w.WriteHeader(200)
fmt.Fprintf(w, "status: %s\n", string(p.status))
fmt.Fprintf(w, "failures: %d\n", p.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)
}
timeout, cancel := context.WithTimeout(context.Background(), time.Duration(*flagTimeout)*time.Second)
defer cancel()
for {
select {
case <-timeout.Done():
slog.Info("timeout waiting for upstream to be ready")
http.Error(w, "timeout", http.StatusRequestTimeout)
return
default:
if p.getStatus() == ready {
break
}
// prevent busy waiting
<-time.After(100 * time.Millisecond)
}

// break the loop
if p.getStatus() == ready {
break
}
}
}

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
}
13 changes: 13 additions & 0 deletions proxy/proxymanager.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()

Expand Down Expand Up @@ -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)
Expand Down