From 2c59c66152e7e6ed53d979095ed42f44da20431b Mon Sep 17 00:00:00 2001
From: Kot C <kot@yukata.dev>
Date: Sat, 7 Dec 2024 17:46:42 -0600
Subject: [PATCH] Use httputil/NewSingleHostReverseProxy instead of our own
 implementation (#5)

---
 cmd/proxy/main.go | 42 +++---------------------------------------
 1 file changed, 3 insertions(+), 39 deletions(-)

diff --git a/cmd/proxy/main.go b/cmd/proxy/main.go
index b157956..6aad2e2 100644
--- a/cmd/proxy/main.go
+++ b/cmd/proxy/main.go
@@ -1,16 +1,13 @@
 package main
 
 import (
-	"errors"
 	"flag"
-	"fmt"
-	"io"
 	"log"
 	"log/slog"
 	"net/http"
+	"net/http/httputil"
 	"net/url"
 	"strings"
-	"time"
 
 	"tailscale.com/client/tailscale"
 	"tailscale.com/tsnet"
@@ -45,15 +42,9 @@ func main() {
 		log.Fatalf("error creating tailscale local client: %v", err)
 	}
 
+	proxy := httputil.NewSingleHostReverseProxy(endpointUrl)
 	log.Printf("proxying requests to %s", endpointUrl)
 
-	httpClient := &http.Client{
-		CheckRedirect: func(req *http.Request, via []*http.Request) error {
-			return http.ErrUseLastResponse
-		},
-		Timeout: time.Minute,
-	}
-
 	log.Fatal(http.Serve(ln,
 		http.HandlerFunc(
 			func(writer http.ResponseWriter, req *http.Request) {
@@ -72,10 +63,6 @@ func main() {
 				slog.Debug("tailscale", "whois", who)
 
 				req.Host = endpointUrl.Host
-				req.URL.Host = endpointUrl.Host
-				req.URL.Scheme = endpointUrl.Scheme
-				req.RequestURI = ""
-
 				for key, value := range req.Header {
 					if strings.HasPrefix(http.CanonicalHeaderKey("X-Remote-"), key) {
 						slog.Info("removing spoofed header", "key", key, "value", value)
@@ -89,29 +76,6 @@ func main() {
 					req.Header.Set("X-Remote-User-Id", who.UserProfile.ID.String())
 				}
 
-				res, err := httpClient.Do(req)
-				if err != nil {
-					var urlError *url.Error
-					if errors.As(err, &urlError) && urlError.Timeout() {
-						writer.WriteHeader(http.StatusGatewayTimeout)
-					} else {
-						writer.WriteHeader(http.StatusBadGateway)
-					}
-
-					_, _ = fmt.Fprint(writer, err)
-					slog.Error("proxying request", "err", err)
-					return
-				}
-
-				for key, value := range res.Header {
-					writer.Header()[key] = value
-				}
-				writer.WriteHeader(res.StatusCode)
-
-				_, err = io.Copy(writer, res.Body)
-				defer res.Body.Close()
-				if err != nil {
-					slog.Error("reading response body", "err", err)
-				}
+				proxy.ServeHTTP(writer, req)
 			})))
 }