From c7d4838f1cf24bf176f872bed9d352902c82d4c0 Mon Sep 17 00:00:00 2001
From: Siguza <siguza@siguza.net>
Date: Fri, 18 Oct 2024 09:13:57 +0200
Subject: [PATCH] Unix socket support

---
 README.md |  5 ++++-
 main.go   | 21 +++++++++++++++------
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index 4020295..4d92665 100644
--- a/README.md
+++ b/README.md
@@ -40,6 +40,7 @@ Command-line flags:
 ------------------
 
 	-listen=":8080": The [IP]:port to listen for incoming connections on.
+	-socket="/run/my.sock": Unix socket to listen to (overrides TCP address).
 	-max_connections=4096: The maximum number of incoming connections allowed.
 	-period=16s: Time between each byte sent on a connection.
 	-response_len=10485760: The number of bytes to send total per connection.
@@ -47,7 +48,9 @@ Command-line flags:
 	-workers=4: The number of worker threads to execute.
 
 It defaults to dual-stack IPv4/IPv6.  If you want IPv4-only, specify an IPv4
-listen address, like -listen="0.0.0.0:8080".
+listen address, like `-listen="0.0.0.0:8080"`.  
+Unix socket defaults to empty string (i.e. unused), but will override the
+`-listen` address for binding purposes, if provided.
 
 It will try to raise "ulimit -n" to the max_connections that you specify. 
 It defaults to raising the limit as much as it can; if you want it higher
diff --git a/main.go b/main.go
index 3f1cdca..a3ce600 100644
--- a/main.go
+++ b/main.go
@@ -13,6 +13,7 @@ import (
 
 var (
 	listenAddr     = flag.String("listen", ":8080", "The [IP]:port to listen for incoming connections on.")
+	unixSocket     = flag.String("socket", "", "Unix socket to listen to (overrides TCP address).")
 	workers        = flag.Int("workers", runtime.NumCPU(), "The number of worker threads to execute.")
 	period         = flag.Duration("period", 16*time.Second, "Time between each byte sent on a connection.")
 	timeslice      = flag.Duration("timeslice", 50*time.Millisecond, "How often each thread should wake up to send.")
@@ -38,14 +39,22 @@ func main() {
 	http.HandleFunc("/", tp.Handler)
 	http.HandleFunc("/robots.txt", robotsDisallowHandler)
 
-	log.Fatal(listenAndServe(*listenAddr))
+	log.Fatal(listenAndServe(*listenAddr, *unixSocket))
 }
 
-func listenAndServe(addr string) error {
+func listenAndServe(addr string, socket string) error {
 	srv := &http.Server{Addr: addr}
-	ln, err := net.Listen("tcp", addr)
-	if err != nil {
-		return err
+	if socket == "" {
+		ln, err := net.Listen("tcp", addr)
+		if err != nil {
+			return err
+		}
+		return srv.Serve(NewBufSizeListener(*rcvBuf, *sndBuf, ln.(*net.TCPListener)))
+	} else {
+		ln, err := net.Listen("unix", socket)
+		if err != nil {
+			return err
+		}
+		return srv.Serve(ln)
 	}
-	return srv.Serve(NewBufSizeListener(*rcvBuf, *sndBuf, ln.(*net.TCPListener)))
 }