caddy: run OnCancel callbacks and restore std logger on reload

Callbacks registered with OnCancel were appended to a copy of the
Context and never ran on config unload, so e.g. the std log redirect to
a removed sink was never undone. Share callbacks across copies of a
Context so they run exactly once on cancellation; callbacks registered
after cancellation run immediately, and OnCancel on a zero-value
Context remains safe.

Since overlapping configs mean an old config's cleanup runs after the
new one is set up, track ownership of the default logger and std log
redirect so cleanup restores the most recent live config's state.

Addresses scan report 88637e95bf243f5a.

Code written by GitHub Copilot (GPT-6.1 Sol, Claude Opus 5.5); reviewed by mholt.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

context.go

@@ -49,7 +49,7 @@ type Context struct {
 	moduleInstances map[string][]Module
 	cfg             *Config
 	ancestry        []Module
-	cleanupFuncs    []func()                // invoked at every config unload
+	cleanupFuncs    *contextCleanup         // shared by value copies; invoked at every config unload
 	exitFuncs       []func(context.Context) // invoked at config unload ONLY IF the process is exiting (EXPERIMENTAL)
 	metricsRegistry *prometheus.Registry
 }
@@ -70,13 +70,18 @@ func NewContext(ctx Context) (Context, context.CancelFunc) {
 // NewContextWithCause is like NewContext but returns a context.CancelCauseFunc.
 // EXPERIMENTAL: This API is subject to change.
 func NewContextWithCause(ctx Context) (Context, context.CancelCauseFunc) {
-	newCtx := Context{moduleInstances: make(map[string][]Module), cfg: ctx.cfg, metricsRegistry: prometheus.NewPedanticRegistry()}
+	newCtx := Context{
+		moduleInstances: make(map[string][]Module),
+		cfg:             ctx.cfg,
+		metricsRegistry: prometheus.NewPedanticRegistry(),
+		cleanupFuncs:    new(contextCleanup),
+	}
 	c, cancel := context.WithCancelCause(ctx.Context)
 	wrappedCancel := func(cause error) {
 		cancel(cause)
 
-		for _, f := range ctx.cleanupFuncs {
-			f()
+		if !newCtx.cleanupFuncs.run() {
+			return
 		}
 
 		for modName, modInstances := range newCtx.moduleInstances {
@@ -95,9 +100,49 @@ func NewContextWithCause(ctx Context) (Context, context.CancelCauseFunc) {
 	return newCtx, wrappedCancel
 }
 
-// OnCancel executes f when ctx is canceled.
+// OnCancel executes f once when ctx is canceled. Registrations on value copies
+// and WithValue contexts share the same lifetime. Registrations on a context
+// created by NewContext belong to that new lifetime, not its parent.
+// If cancellation has already begun, f is executed immediately.
 func (ctx *Context) OnCancel(f func()) {
-	ctx.cleanupFuncs = append(ctx.cleanupFuncs, f)
+	if ctx.cleanupFuncs == nil {
+		// a Context not created by NewContext has no cancel func to run f
+		ctx.cleanupFuncs = new(contextCleanup)
+	}
+	ctx.cleanupFuncs.add(f)
+}
+
+type contextCleanup struct {
+	mu       sync.Mutex
+	canceled bool
+	funcs    []func()
+}
+
+func (c *contextCleanup) add(f func()) {
+	c.mu.Lock()
+	if c.canceled {
+		c.mu.Unlock()
+		f()
+		return
+	}
+	c.funcs = append(c.funcs, f)
+	c.mu.Unlock()
+}
+
+func (c *contextCleanup) run() bool {
+	c.mu.Lock()
+	if c.canceled {
+		c.mu.Unlock()
+		return false
+	}
+	c.canceled = true
+	funcs := c.funcs
+	c.funcs = nil
+	c.mu.Unlock()
+	for _, f := range funcs {
+		f()
+	}
+	return true
 }
 
 // FileSystems returns a ref to the FilesystemMap.

context_test.go

@@ -17,10 +17,85 @@ package caddy
 import (
 	"context"
 	"encoding/json"
+	"errors"
 	"io"
+	"sync"
+	"sync/atomic"
 	"testing"
 )
 
+func TestContextCancellationCallbacks(t *testing.T) {
+	ctx, cancel := NewContextWithCause(Context{Context: context.Background()})
+	copy := ctx
+	derived := ctx.WithValue("key", "value")
+	var count atomic.Int32
+	for _, c := range []*Context{&ctx, &copy, &derived} {
+		c.OnCancel(func() { count.Add(1) })
+	}
+	cause := errors.New("test cancellation")
+	cancel(cause)
+	cancel(errors.New("second cancellation"))
+	copy.OnCancel(func() { count.Add(1) })
+	if got := count.Load(); got != 4 {
+		t.Fatalf("callbacks ran %d times, want 4", got)
+	}
+	if context.Cause(ctx) != cause {
+		t.Fatalf("cancellation cause = %v, want %v", context.Cause(ctx), cause)
+	}
+}
+
+func TestContextCancellationLifetimes(t *testing.T) {
+	parent, cancelParent := NewContext(Context{Context: context.Background()})
+	child, cancelChild := NewContext(parent)
+	var parents, children int
+	parent.OnCancel(func() { parents++ })
+	child.OnCancel(func() { children++ })
+	cancelChild()
+	if parents != 0 || children != 1 {
+		t.Fatalf("after child cancel: parent=%d child=%d", parents, children)
+	}
+	cancelParent()
+	cancelChild()
+	if parents != 1 || children != 1 {
+		t.Fatalf("cleanup counts: parent=%d child=%d", parents, children)
+	}
+}
+
+func TestContextOnCancelZeroValue(t *testing.T) {
+	var ctx Context
+	ctx.OnCancel(func() { t.Fatal("callback on context without cancel func must not run") })
+}
+
+func TestContextCancellationConcurrent(t *testing.T) {
+	ctx, cancel := NewContext(Context{Context: context.Background()})
+	cleaner := new(contextTestCleaner)
+	ctx.moduleInstances["test"] = []Module{cleaner}
+	var count atomic.Int32
+	var wg sync.WaitGroup
+	for range 50 {
+		wg.Go(func() {
+			copy := ctx
+			copy.OnCancel(func() { count.Add(1) })
+			cancel()
+		})
+	}
+	wg.Wait()
+	if count.Load() != 50 {
+		t.Fatalf("callbacks ran %d times, want 50", count.Load())
+	}
+	if cleaner.count.Load() != 1 {
+		t.Fatalf("module cleanup ran %d times, want 1", cleaner.count.Load())
+	}
+}
+
+type contextTestCleaner struct{ count atomic.Int32 }
+
+func (*contextTestCleaner) CaddyModule() ModuleInfo { return ModuleInfo{ID: "test"} }
+func (c *contextTestCleaner) Cleanup() error {
+	c.count.Add(1)
+	return nil
+}
+
 func TestContextWithValuePreservesMetricsRegistry(t *testing.T) {
 	ctx, cancel := NewContext(Context{Context: context.Background()})
 	t.Cleanup(cancel)

logging.go

@@ -78,6 +78,8 @@ type Logging struct {
 	// must have their keys added to this list so they
 	// can be closed when cleaning up
 	writerKeys []string
+
+	owner *loggingOwner
 }
 
 // openLogs sets up the config and opens all the configured writers.
@@ -99,6 +101,11 @@ func (logging *Logging) openLogs(ctx Context) error {
 			return fmt.Errorf("setting up sink log: %v", err)
 		}
 	}
+	var sink *zap.Logger
+	if logging.Sink != nil {
+		sink = logging.Sink.logger
+	}
+	logging.owner = claimLogging(sink)
 
 	// as a special case, set up the default structured Caddy log next
 	if err := logging.setupNewDefault(ctx); err != nil {
@@ -172,6 +179,7 @@ func (logging *Logging) setupNewDefault(ctx Context) error {
 	defaultLoggerMu.Lock()
 	oldDefault := defaultLogger
 	defaultLogger = newDefault
+	logging.owner.defaultLog = newDefault
 	defaultLoggerMu.Unlock()
 
 	// if the new writer is different, indicate it in the logs for convenience
@@ -206,6 +214,9 @@ func (logging *Logging) setupNewDefault(ctx Context) error {
 // A successful call to openLogs calls this automatically
 // when the context is canceled.
 func (logging *Logging) closeLogs() error {
+	if logging.owner != nil {
+		logging.owner.release()
+	}
 	for _, key := range logging.writerKeys {
 		_, err := writers.Delete(key)
 		if err != nil {
@@ -215,6 +226,69 @@ func (logging *Logging) closeLogs() error {
 	return nil
 }
 
+// loggingOwner tracks overlapping configurations, including ones still being
+// provisioned. Old configuration cleanup must not undo newer redirects, and a
+// failed new configuration must restore the most recent still-live owner.
+type loggingOwner struct {
+	defaultLog *defaultCustomLog
+	stdLog     stdLogState
+}
+
+type stdLogState struct {
+	output io.Writer
+	flags  int
+	prefix string
+}
+
+func currentStdLog() stdLogState {
+	return stdLogState{log.Writer(), log.Flags(), log.Prefix()}
+}
+
+func (s stdLogState) restore() {
+	log.SetOutput(s.output)
+	log.SetFlags(s.flags)
+	log.SetPrefix(s.prefix)
+}
+
+func claimLogging(sink *zap.Logger) *loggingOwner {
+	defaultLoggerMu.Lock()
+	defer defaultLoggerMu.Unlock()
+	if len(loggingOwners) == 0 {
+		loggingBaseline = loggingOwner{defaultLogger, currentStdLog()}
+	}
+	owner := &loggingOwner{defaultLog: defaultLogger, stdLog: loggingBaseline.stdLog}
+	if sink != nil {
+		// zap's undo function always restores stderr, not the previous writer.
+		// Save the full state ourselves so failed provisioning can roll back.
+		zap.RedirectStdLog(sink)
+		owner.stdLog = currentStdLog()
+	} else {
+		owner.stdLog.restore()
+	}
+	loggingOwners = append(loggingOwners, owner)
+	return owner
+}
+
+func (o *loggingOwner) release() {
+	defaultLoggerMu.Lock()
+	defer defaultLoggerMu.Unlock()
+	for i, owner := range loggingOwners {
+		if owner != o {
+			continue
+		}
+		loggingOwners = slices.Delete(loggingOwners, i, i+1)
+		if i == len(loggingOwners) {
+			restore := &loggingBaseline
+			if i > 0 {
+				restore = loggingOwners[i-1]
+			}
+			defaultLogger = restore.defaultLog
+			restore.stdLog.restore()
+		}
+		return
+	}
+}
+
 // Logger returns a logger that is ready for the module to use.
 func (logging *Logging) Logger(mod Module) *zap.Logger {
 	modID := string(mod.CaddyModule().ID)
@@ -444,9 +518,12 @@ func (cl *BaseLog) buildOptions() ([]zap.Option, error) {
 // global logger in the log package. This is necessary because
 // module dependencies which are not built specifically for
 // Caddy will use the standard logger. This is also known as
-// the "sink" logger.
+// the "sink" logger. Removing it on reload restores the standard logger's
+// original output, prefix, and flags.
 type SinkLog struct {
 	BaseLog
+
+	logger *zap.Logger
 }
 
 func (sll *SinkLog) provision(ctx Context, logging *Logging) error {
@@ -459,8 +536,7 @@ func (sll *SinkLog) provision(ctx Context, logging *Logging) error {
 		return err
 	}
 
-	logger := zap.New(sll.core, options...)
-	ctx.cleanupFuncs = append(ctx.cleanupFuncs, zap.RedirectStdLog(logger))
+	sll.logger = zap.New(sll.core, options...)
 	return nil
 }
 
@@ -819,6 +895,8 @@ var (
 	coloringEnabled  = os.Getenv("NO_COLOR") == "" && os.Getenv("TERM") != "xterm-mono"
 	defaultLogger, _ = newDefaultProductionLog()
 	defaultLoggerMu  sync.RWMutex
+	loggingOwners    []*loggingOwner // protected by defaultLoggerMu
+	loggingBaseline  loggingOwner
 )
 
 var writers = NewUsagePool()

logging_lifecycle_test.go

@@ -0,0 +1,234 @@
+// Copyright 2015 Matthew Holt and The Caddy Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package caddy
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+	"log"
+	"strings"
+	"sync"
+	"testing"
+)
+
+func TestLoggingSinkReloadLifecycle(t *testing.T) {
+	baseline, originalDefault := prepareLoggingLifecycle(t)
+	a, sinkA, _ := lifecycleLogging(t, true)
+	cancelA := provisionLifecycleLogging(t, a)
+	log.Print("first sink")
+	if !strings.Contains(sinkA.contents(), "first sink") {
+		t.Fatal("standard logger did not reach initial sink")
+	}
+
+	b, sinkB, _ := lifecycleLogging(t, true)
+	cancelB := provisionLifecycleLogging(t, b)
+	cancelA()
+	log.Print("replacement sink")
+	if !strings.Contains(sinkB.contents(), "replacement sink") || !sinkA.isClosed() {
+		t.Fatal("old config cleanup undid the replacement sink or leaked its writer")
+	}
+	if strings.Contains(sinkA.contents(), "replacement sink") {
+		t.Fatal("replacement log reached old sink")
+	}
+
+	c, _, _ := lifecycleLogging(t, false)
+	cancelC := provisionLifecycleLogging(t, c)
+	cancelB()
+	log.Print("sink removed")
+	if !strings.Contains(baseline.contents(), "baseline: sink removed") {
+		t.Fatalf("removed sink did not restore original logger: %s", baseline.contents())
+	}
+	if strings.Contains(sinkB.contents(), "sink removed") || !sinkB.isClosed() {
+		t.Fatal("removed sink received another log or leaked its writer")
+	}
+
+	d, sinkD, _ := lifecycleLogging(t, true)
+	cancelD := provisionLifecycleLogging(t, d)
+	cancelC()
+	log.Print("sink reintroduced")
+	if !strings.Contains(sinkD.contents(), "sink reintroduced") {
+		t.Fatal("sinkless config cleanup undid newer sink")
+	}
+	cancelD()
+	if log.Writer() != baseline || log.Prefix() != "baseline: " || log.Flags() != log.Lmsgprefix {
+		t.Fatal("final cleanup did not restore original standard logger state")
+	}
+	if defaultLogger != originalDefault {
+		t.Fatal("final cleanup did not restore original Caddy logger")
+	}
+}
+
+func TestLoggingFailedProvisioningRestoresOwner(t *testing.T) {
+	for _, sink := range []bool{false, true} {
+		t.Run(fmt.Sprintf("active_sink=%t", sink), func(t *testing.T) {
+			testLoggingFailedProvisioningRestoresOwner(t, sink)
+		})
+	}
+}
+
+func testLoggingFailedProvisioningRestoresOwner(t *testing.T, sink bool) {
+	baseline, _ := prepareLoggingLifecycle(t)
+	active, activeSink, activeDefault := lifecycleLogging(t, sink)
+	activeWriter := activeSink
+	if activeWriter == nil {
+		activeWriter = baseline
+	}
+	provisionLifecycleLogging(t, active)
+	originalDefault := defaultLogger
+	for _, stage := range []string{"sink", "default", "custom", "storage"} {
+		for _, sink := range []bool{false, true} {
+			if stage == "sink" && !sink {
+				continue
+			}
+			t.Run(stage+map[bool]string{false: "/without_sink", true: "/with_sink"}[sink], func(t *testing.T) {
+				failed, failedSink, failedDefault := lifecycleLogging(t, sink)
+				switch stage {
+				case "sink":
+					failed.Sink.Level = "invalid"
+				case "default":
+					failed.Logs[DefaultLoggerName].Level = "invalid"
+				case "custom":
+					failed.Logs["broken"] = &CustomLog{BaseLog: BaseLog{Level: "invalid"}}
+				}
+				var cancel context.CancelFunc
+				if stage == "storage" {
+					ctx, err := provisionContext(&Config{
+						Logging:    failed,
+						StorageRaw: json.RawMessage(`{"module":"nonexistent"}`),
+					}, false)
+					if err == nil {
+						t.Fatal("expected storage provisioning failure")
+					}
+					cancel = func() { ctx.cfg.cancelFunc(nil) }
+				} else {
+					ctx, cancelContext := NewContext(Context{Context: context.Background()})
+					cancel = cancelContext
+					if err := failed.openLogs(ctx); err == nil {
+						t.Fatal("expected logging provisioning failure")
+					}
+				}
+				t.Cleanup(cancel)
+				cancel()
+				cancel()
+				log.Print("after failed provisioning")
+				Log().Info("Caddy logger after failed provisioning")
+				if !strings.Contains(activeWriter.contents(), "after failed provisioning") {
+					t.Fatal("failure did not restore active sink")
+				}
+				if defaultLogger != originalDefault || activeWriter.isClosed() || activeDefault.isClosed() {
+					t.Fatal("failure replaced or closed active loggers")
+				}
+				if stage != "sink" && !failedDefault.isClosed() {
+					t.Fatal("failed configuration's default writer leaked")
+				}
+				if sink && !failedSink.isClosed() {
+					t.Fatal("failed configuration's sink writer leaked")
+				}
+			})
+		}
+	}
+}
+
+func prepareLoggingLifecycle(t *testing.T) (*lifecycleLogWriter, *defaultCustomLog) {
+	t.Helper()
+	originalWriter, originalFlags, originalPrefix := log.Writer(), log.Flags(), log.Prefix()
+	defaultLoggerMu.Lock()
+	originalDefault := defaultLogger
+	originalOwners, originalBaseline := loggingOwners, loggingBaseline
+	loggingOwners = nil
+	defaultLoggerMu.Unlock()
+	baseline := new(lifecycleLogWriter)
+	log.SetOutput(baseline)
+	log.SetPrefix("baseline: ")
+	log.SetFlags(log.Lmsgprefix)
+	t.Cleanup(func() {
+		defaultLoggerMu.Lock()
+		if len(loggingOwners) != 0 {
+			t.Error("configuration logger ownership leaked")
+		}
+		loggingOwners, loggingBaseline = originalOwners, originalBaseline
+		defaultLogger = originalDefault
+		defaultLoggerMu.Unlock()
+		log.SetOutput(originalWriter)
+		log.SetFlags(originalFlags)
+		log.SetPrefix(originalPrefix)
+	})
+	return baseline, originalDefault
+}
+
+func lifecycleLogging(t *testing.T, sink bool) (*Logging, *lifecycleLogWriter, *lifecycleLogWriter) {
+	t.Helper()
+	defaultWriter := new(lifecycleLogWriter)
+	defaultWriter.key = fmt.Sprintf("%s/default/%p", t.Name(), defaultWriter)
+	logging := &Logging{Logs: map[string]*CustomLog{
+		DefaultLoggerName: {BaseLog: BaseLog{writerOpener: defaultWriter}},
+	}}
+	var sinkWriter *lifecycleLogWriter
+	if sink {
+		sinkWriter = new(lifecycleLogWriter)
+		sinkWriter.key = fmt.Sprintf("%s/sink/%p", t.Name(), sinkWriter)
+		logging.Sink = &SinkLog{BaseLog: BaseLog{writerOpener: sinkWriter}}
+	}
+	return logging, sinkWriter, defaultWriter
+}
+
+func provisionLifecycleLogging(t *testing.T, logging *Logging) context.CancelFunc {
+	t.Helper()
+	ctx, cancel := NewContext(Context{Context: context.Background()})
+	t.Cleanup(cancel)
+	if err := logging.openLogs(ctx); err != nil {
+		t.Fatal(err)
+	}
+	return cancel
+}
+
+type lifecycleLogWriter struct {
+	mu     sync.Mutex
+	buffer bytes.Buffer
+	key    string
+	closed bool
+}
+
+func (w *lifecycleLogWriter) WriterKey() string                   { return w.key }
+func (w *lifecycleLogWriter) String() string                      { return w.key }
+func (w *lifecycleLogWriter) OpenWriter() (io.WriteCloser, error) { return w, nil }
+func (w *lifecycleLogWriter) Write(p []byte) (int, error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	if w.closed {
+		return 0, errors.New("write to closed lifecycle log")
+	}
+	return w.buffer.Write(p)
+}
+func (w *lifecycleLogWriter) Close() error {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	w.closed = true
+	return nil
+}
+func (w *lifecycleLogWriter) contents() string {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.buffer.String()
+}
+func (w *lifecycleLogWriter) isClosed() bool {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.closed
+}