From d412672ed487a39890ade5f52b86e30761e50373 Mon Sep 17 00:00:00 2001
From: WeebDataHoarder <weebdatahoarder@noreply.gammaspectra.live>
Date: Thu, 1 May 2025 14:16:19 +0200
Subject: [PATCH] state: explicitly free resources on Close()

---
 cmd/go-away/main.go |  6 +++++-
 lib/state.go        | 23 ++++++++++++++++++++---
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/cmd/go-away/main.go b/cmd/go-away/main.go
index e5493e9..873e8cf 100644
--- a/cmd/go-away/main.go
+++ b/cmd/go-away/main.go
@@ -238,7 +238,7 @@ func main() {
 		acmeCache = path.Join(*cachePath, "acme")
 	}
 
-	loadPolicyState := func() (http.Handler, error) {
+	loadPolicyState := func() (*lib.State, error) {
 		policyData, err := os.ReadFile(*policyFile)
 		if err != nil {
 			return nil, fmt.Errorf("failed to read policy file: %w", err)
@@ -309,6 +309,7 @@ func main() {
 			if sig != syscall.SIGHUP {
 				continue
 			}
+			oldHandler := handler
 			handler, err = loadPolicyState()
 			if err != nil {
 				slog.Error("handler configuration reload error", "err", err)
@@ -317,6 +318,9 @@ func main() {
 
 			swap(handler)
 			slog.Warn("handler configuration reloaded")
+			if oldHandler != nil {
+				_ = oldHandler.Close()
+			}
 		}
 	}()
 
diff --git a/lib/state.go b/lib/state.go
index 8c16eab..2d4f584 100644
--- a/lib/state.go
+++ b/lib/state.go
@@ -54,9 +54,8 @@ type State struct {
 	Mux *http.ServeMux
 }
 
-func NewState(p policy.Policy, opt settings.Settings, settings policy.StateSettings) (handler http.Handler, err error) {
-
-	state := new(State)
+func NewState(p policy.Policy, opt settings.Settings, settings policy.StateSettings) (state *State, err error) {
+	state = new(State)
 	state.close = make(chan struct{})
 	state.settings = settings
 	state.opt = opt
@@ -273,3 +272,21 @@ func NewState(p policy.Policy, opt settings.Settings, settings policy.StateSetti
 
 	return state, nil
 }
+
+func (state *State) Close() error {
+	select {
+	case <-state.close:
+	default:
+		close(state.close)
+		for _, c := range state.challenges {
+			if c.Object != nil {
+				err := c.Object.Close()
+				if err != nil {
+					return err
+				}
+			}
+		}
+	}
+
+	return nil
+}