Merge pull request #166 from SenseUnit/proper_shutdown

Snawoot · web-flow · commit fb4cc9f8322e · 2025-11-15T22:30:53.000+02:00
Tidy shutdown
diff --git a/auth/auth.go b/auth/auth.go
@@ -3,6 +3,7 @@ package auth
 import (
 	"context"
 	"errors"
+	"io"
 	"net/http"
 	"net/url"
 	"strings"
@@ -12,7 +13,7 @@ import (
 
 type Auth interface {
 	Validate(ctx context.Context, wr http.ResponseWriter, req *http.Request) (string, bool)
-	Stop()
+	io.Closer
 }
 
 func NewAuth(paramstr string, logger *clog.CondLogger) (Auth, error) {
diff --git a/auth/basic.go b/auth/basic.go
@@ -173,11 +173,13 @@ func (auth *BasicAuth) Validate(ctx context.Context, wr http.ResponseWriter, req
 	return requireBasicAuth(ctx, wr, req, auth.hiddenDomain, auth.next)
 }
 
-func (auth *BasicAuth) Stop() {
+func (auth *BasicAuth) Close() error {
+	var err error
 	auth.stopOnce.Do(func() {
 		if auth.next != nil {
-			auth.next.Stop()
+			err = auth.next.Close()
 		}
 		close(auth.stopChan)
 	})
+	return err
 }
diff --git a/auth/cert.go b/auth/cert.go
@@ -16,6 +16,7 @@ import (
 	"time"
 
 	clog "github.com/SenseUnit/dumbproxy/log"
+	"github.com/hashicorp/go-multierror"
 
 	us "github.com/Snawoot/uniqueslice"
 )
@@ -108,16 +109,22 @@ func (auth *CertAuth) Validate(ctx context.Context, wr http.ResponseWriter, req
 	), true
 }
 
-func (auth *CertAuth) Stop() {
+func (auth *CertAuth) Close() error {
+	var err error
 	auth.stopOnce.Do(func() {
 		if auth.next != nil {
-			auth.next.Stop()
+			if closeErr := auth.next.Close(); closeErr != nil {
+				err = multierror.Append(err, closeErr)
+			}
 		}
 		if auth.reject != nil {
-			auth.reject.Stop()
+			if closeErr := auth.reject.Close(); closeErr != nil {
+				err = multierror.Append(err, closeErr)
+			}
 		}
 		close(auth.stopChan)
 	})
+	return err
 }
 
 func (auth *CertAuth) reload() error {
diff --git a/auth/hmac.go b/auth/hmac.go
@@ -143,12 +143,14 @@ func (auth *HMACAuth) Validate(ctx context.Context, wr http.ResponseWriter, req
 	return requireBasicAuth(ctx, wr, req, auth.hiddenDomain, auth.next)
 }
 
-func (auth *HMACAuth) Stop() {
+func (auth *HMACAuth) Close() error {
+	var err error
 	auth.stopOnce.Do(func() {
 		if auth.next != nil {
-			auth.next.Stop()
+			err = auth.next.Close()
 		}
 	})
+	return err
 }
 
 func CalculateHMACSignature(secret []byte, username string, expire int64) []byte {
diff --git a/auth/noauth.go b/auth/noauth.go
@@ -11,4 +11,6 @@ func (_ NoAuth) Validate(_ context.Context, _ http.ResponseWriter, _ *http.Reque
 	return "", true
 }
 
-func (_ NoAuth) Stop() {}
+func (_ NoAuth) Close() error {
+	return nil
+}
diff --git a/auth/redis.go b/auth/redis.go
@@ -13,6 +13,7 @@ import (
 	"time"
 
 	clog "github.com/SenseUnit/dumbproxy/log"
+	"github.com/hashicorp/go-multierror"
 
 	"github.com/redis/go-redis/v9"
 )
@@ -134,11 +135,17 @@ func (auth *RedisAuth) Validate(ctx context.Context, wr http.ResponseWriter, req
 	return requireBasicAuth(ctx, wr, req, auth.hiddenDomain, auth.next)
 }
 
-func (auth *RedisAuth) Stop() {
+func (auth *RedisAuth) Close() error {
+	var err error
 	auth.stopOnce.Do(func() {
 		if auth.next != nil {
-			auth.next.Stop()
+			if closeErr := auth.next.Close(); closeErr != nil {
+				err = multierror.Append(err, closeErr)
+			}
+		}
+		if closeErr := auth.r.Close(); closeErr != nil {
+			err = multierror.Append(err, closeErr)
 		}
-		auth.r.Close()
 	})
+	return err
 }
diff --git a/auth/rejecthttp.go b/auth/rejecthttp.go
@@ -56,4 +56,6 @@ func (_ *RejectHTTPAuth) Valid(_, _, _ string) bool {
 	return false
 }
 
-func (_ *RejectHTTPAuth) Stop() {}
+func (_ *RejectHTTPAuth) Close() error {
+	return nil
+}
diff --git a/auth/rejectstatic.go b/auth/rejectstatic.go
@@ -105,4 +105,6 @@ func (_ *StaticRejectAuth) Valid(_, _, _ string) bool {
 	return false
 }
 
-func (_ *StaticRejectAuth) Stop() {}
+func (_ *StaticRejectAuth) Close() error {
+	return nil
+}
diff --git a/certcache/cryptobox.go b/certcache/cryptobox.go
@@ -5,6 +5,7 @@ import (
 	"crypto/cipher"
 	cryptorand "crypto/rand"
 	"errors"
+	"io"
 
 	"golang.org/x/crypto/acme/autocert"
 	"golang.org/x/crypto/chacha20poly1305"
@@ -64,3 +65,10 @@ func (c *EncryptedCache) Put(ctx context.Context, key string, data []byte) error
 func (c *EncryptedCache) Delete(ctx context.Context, key string) error {
 	return c.next.Delete(ctx, key)
 }
+
+func (c *EncryptedCache) Close() error {
+	if cacheCloser, ok := c.next.(io.Closer); ok {
+		return cacheCloser.Close()
+	}
+	return nil
+}
diff --git a/certcache/local.go b/certcache/local.go
@@ -2,6 +2,7 @@ package certcache
 
 import (
 	"context"
+	"io"
 	"sync"
 	"time"
 
@@ -16,10 +17,9 @@ type certCacheValue struct {
 }
 
 type LocalCertCache struct {
-	cache     *ttlcache.Cache[certCacheKey, certCacheValue]
-	next      autocert.Cache
-	startOnce sync.Once
-	stopOnce  sync.Once
+	cache    *ttlcache.Cache[certCacheKey, certCacheValue]
+	next     autocert.Cache
+	stopOnce sync.Once
 }
 
 func NewLocalCertCache(next autocert.Cache, ttl, timeout time.Duration) *LocalCertCache {
@@ -41,6 +41,7 @@ func NewLocalCertCache(next autocert.Cache, ttl, timeout time.Duration) *LocalCe
 				nil),
 		),
 	)
+	go cache.Start()
 	return &LocalCertCache{
 		cache: cache,
 		next:  next,
@@ -62,14 +63,15 @@ func (cc *LocalCertCache) Delete(ctx context.Context, key string) error {
 	return cc.next.Delete(ctx, key)
 }
 
-func (cc *LocalCertCache) Start() {
-	cc.startOnce.Do(func() {
-		go cc.cache.Start()
+func (cc *LocalCertCache) Close() error {
+	var err error
+	cc.stopOnce.Do(func() {
+		cc.cache.Stop()
+		if cacheCloser, ok := cc.next.(io.Closer); ok {
+			err = cacheCloser.Close()
+		}
 	})
-}
-
-func (cc *LocalCertCache) Stop() {
-	cc.stopOnce.Do(cc.cache.Stop)
+	return err
 }
 
 var _ autocert.Cache = new(LocalCertCache)
diff --git a/certcache/redis.go b/certcache/redis.go
@@ -2,17 +2,25 @@ package certcache
 
 import (
 	"context"
+	"io"
+	"sync"
 
 	"github.com/redis/go-redis/v9"
 	"golang.org/x/crypto/acme/autocert"
 )
 
+type CmdableCloser interface {
+	redis.Cmdable
+	io.Closer
+}
+
 type RedisCache struct {
-	r   redis.Cmdable
-	pfx string
+	r        CmdableCloser
+	pfx      string
+	stopOnce sync.Once
 }
 
-func NewRedisCache(r redis.Cmdable, prefix string) *RedisCache {
+func NewRedisCache(r CmdableCloser, prefix string) *RedisCache {
 	return &RedisCache{
 		r:   r,
 		pfx: prefix,
@@ -38,6 +46,14 @@ func (r *RedisCache) Delete(ctx context.Context, key string) error {
 	return r.r.Del(ctx, r.pfx+key).Err()
 }
 
+func (r *RedisCache) Close() error {
+	var err error
+	r.stopOnce.Do(func() {
+		err = r.r.Close()
+	})
+	return err
+}
+
 func RedisCacheFromURL(url string, prefix string) (*RedisCache, error) {
 	opts, err := redis.ParseURL(url)
 	if err != nil {
diff --git a/main.go b/main.go
@@ -561,7 +561,7 @@ func run() int {
 		mainLogger.Critical("Failed to instantiate auth provider: %v", err)
 		return 3
 	}
-	defer authProvider.Stop()
+	defer authProvider.Close()
 
 	// setup access filters
 	var filterRoot access.Filter = access.AlwaysAllow{}
@@ -744,6 +744,9 @@ func run() int {
 				mainLogger.Critical("redis cluster cache construction failed: %v", err)
 				return 3
 			}
+		default:
+			mainLogger.Critical("unknown cert cache type %#v", args.autocertCache.kind)
+			return 3
 		}
 		if len(args.autocertCacheEncKey.Value()) > 0 {
 			certCache, err = certcache.NewEncryptedCache(args.autocertCacheEncKey.Value(), certCache)
@@ -758,10 +761,11 @@ func run() int {
 				args.autocertLocalCacheTTL,
 				args.autocertLocalCacheTimeout,
 			)
-			lcc.Start()
-			defer lcc.Stop()
 			certCache = lcc
 		}
+		if cacheCloser, ok := certCache.(io.Closer); ok {
+			defer cacheCloser.Close()
+		}
 
 		m := &autocert.Manager{
 			Cache:  certCache,

Original file line number	Diff line number	Diff line change
`@@ -173,11 +173,13 @@ func (auth *BasicAuth) Validate(ctx context.Context, wr http.ResponseWriter, req`
`173`	`173`	`return requireBasicAuth(ctx, wr, req, auth.hiddenDomain, auth.next)`
`174`	`174`	`}`
`175`	`175`
`176`		`-func (auth *BasicAuth) Stop() {`
	`176`	`+func (auth *BasicAuth) Close() error {`
	`177`	`+ var err error`
`177`	`178`	`auth.stopOnce.Do(func() {`
`178`	`179`	`if auth.next != nil {`
`179`		`- auth.next.Stop()`
	`180`	`+ err = auth.next.Close()`
`180`	`181`	`}`
`181`	`182`	`close(auth.stopChan)`
`182`	`183`	`})`
	`184`	`+ return err`
`183`	`185`	`}`
Original file line number	Diff line number	Diff line change
`@@ -16,6 +16,7 @@ import (`
`16`	`16`	`"time"`
`17`	`17`
`18`	`18`	`clog "github.com/SenseUnit/dumbproxy/log"`
	`19`	`+ "github.com/hashicorp/go-multierror"`
`19`	`20`
`20`	`21`	`us "github.com/Snawoot/uniqueslice"`
`21`	`22`	`)`
`@@ -108,16 +109,22 @@ func (auth *CertAuth) Validate(ctx context.Context, wr http.ResponseWriter, req`
`108`	`109`	`), true`
`109`	`110`	`}`
`110`	`111`
`111`		`-func (auth *CertAuth) Stop() {`
	`112`	`+func (auth *CertAuth) Close() error {`
	`113`	`+ var err error`
`112`	`114`	`auth.stopOnce.Do(func() {`
`113`	`115`	`if auth.next != nil {`
`114`		`- auth.next.Stop()`
	`116`	`+ if closeErr := auth.next.Close(); closeErr != nil {`
	`117`	`+ err = multierror.Append(err, closeErr)`
	`118`	`+ }`
`115`	`119`	`}`
`116`	`120`	`if auth.reject != nil {`
`117`		`- auth.reject.Stop()`
	`121`	`+ if closeErr := auth.reject.Close(); closeErr != nil {`
	`122`	`+ err = multierror.Append(err, closeErr)`
	`123`	`+ }`
`118`	`124`	`}`
`119`	`125`	`close(auth.stopChan)`
`120`	`126`	`})`
	`127`	`+ return err`
`121`	`128`	`}`
`122`	`129`
`123`	`130`	`func (auth *CertAuth) reload() error {`
Original file line number	Diff line number	Diff line change
`@@ -143,12 +143,14 @@ func (auth *HMACAuth) Validate(ctx context.Context, wr http.ResponseWriter, req`
`143`	`143`	`return requireBasicAuth(ctx, wr, req, auth.hiddenDomain, auth.next)`
`144`	`144`	`}`
`145`	`145`
`146`		`-func (auth *HMACAuth) Stop() {`
	`146`	`+func (auth *HMACAuth) Close() error {`
	`147`	`+ var err error`
`147`	`148`	`auth.stopOnce.Do(func() {`
`148`	`149`	`if auth.next != nil {`
`149`		`- auth.next.Stop()`
	`150`	`+ err = auth.next.Close()`
`150`	`151`	`}`
`151`	`152`	`})`
	`153`	`+ return err`
`152`	`154`	`}`
`153`	`155`
`154`	`156`	`func CalculateHMACSignature(secret []byte, username string, expire int64) []byte {`
Original file line number	Diff line number	Diff line change
`@@ -11,4 +11,6 @@ func (_ NoAuth) Validate(_ context.Context, _ http.ResponseWriter, _ *http.Reque`
`11`	`11`	`return "", true`
`12`	`12`	`}`
`13`	`13`
`14`		`-func (_ NoAuth) Stop() {}`
	`14`	`+func (_ NoAuth) Close() error {`
	`15`	`+ return nil`
	`16`	`+}`
Original file line number	Diff line number	Diff line change
`@@ -56,4 +56,6 @@ func (_ *RejectHTTPAuth) Valid(_, _, _ string) bool {`
`56`	`56`	`return false`
`57`	`57`	`}`
`58`	`58`
`59`		`-func (_ *RejectHTTPAuth) Stop() {}`
	`59`	`+func (_ *RejectHTTPAuth) Close() error {`
	`60`	`+ return nil`
	`61`	`+}`
Original file line number	Diff line number	Diff line change
`@@ -105,4 +105,6 @@ func (_ *StaticRejectAuth) Valid(_, _, _ string) bool {`
`105`	`105`	`return false`
`106`	`106`	`}`
`107`	`107`
`108`		`-func (_ *StaticRejectAuth) Stop() {}`
	`108`	`+func (_ *StaticRejectAuth) Close() error {`
	`109`	`+ return nil`
	`110`	`+}`