Add graceful period to server (#499) (#500)

docs/content/configuration/_index.md

@@ -124,6 +124,7 @@ weight: 2
 |    --enabled-proxy-protocol                 | enable proxy protocol | false | PROXY_ENABLE_PROXY_PROTOCOL
 |    --max-idle-connections value             | max idle upstream / keycloak connections to keep alive, ready for reuse | 0 | PROXY_MAX_IDLE_CONNS
 |    --max-idle-connections-per-host value    | limits the number of idle connections maintained per host | 0 | PROXY_MAX_IDLE_CONNS_PER_HOST
+|    --server-grace-timeout value             | the server graceful period before shutdown | 10s | PROXY_SERVER_GRACE_TIMEOUT
 |    --server-read-timeout value              | the server read timeout on the http server | 10s | PROXY_SERVER_READ_TIMEOUT
 |    --server-write-timeout value             | the server write timeout on the http server | 10s | PROXY_SERVER_WRITE_TIMEOUT
 |    --server-idle-timeout value              | the server idle timeout on the http server | 2m0s | PROXY_SERVER_IDLE_TIMEOUT

pkg/constant/constant.go

@@ -89,6 +89,7 @@ const (
 	DefaultOpenIDProviderTimeout         = 30 * time.Second
 	DefaultOpenIDProviderRetryCount      = 3
 	DefaultSelfSignedTLSExpiration       = 3 * time.Hour
+	DefaultServerGraceTimeout            = 10 * time.Second
 	DefaultServerIdleTimeout             = 120 * time.Second
 	DefaultServerReadTimeout             = 10 * time.Second
 	DefaultServerWriteTimeout            = 10 * time.Second
@@ -100,6 +101,8 @@ const (
 	DefaultPatRetryCount                 = 5
 	DefaultPatRetryInterval              = 10 * time.Second
 	DefaultOpaTimeout                    = 10 * time.Second
+
+	ForwardingGrantTypePassword = "password"
 )
 
 var SignatureAlgs = [3]jose.SignatureAlgorithm{jose.RS256, jose.HS256, jose.HS512}

pkg/google/config/config.go

@@ -271,6 +271,8 @@ type Config struct {
 	// MaxIdleConnsPerHost limits the number of idle connections maintained per host
 	MaxIdleConnsPerHost int `env:"MAX_IDLE_CONNS_PER_HOST" json:"max-idle-connections-per-host" usage:"limits the number of idle connections maintained per host" yaml:"max-idle-connections-per-host"`
 
+	// ServerGraceTimeout is the time to wait for the server to shutdown
+	ServerGraceTimeout time.Duration `env:"SERVER_GRACE_TIMEOUT" json:"server-grace-timeout" usage:"the server wait before closing the server" yaml:"server-grace-timeout"`
 	// ServerReadTimeout is the read timeout on the http server
 	ServerReadTimeout time.Duration `env:"SERVER_READ_TIMEOUT" json:"server-read-timeout" usage:"the server read timeout on the http server" yaml:"server-read-timeout"`
 	// ServerWriteTimeout is the write timeout on the http server
@@ -350,6 +352,7 @@ func NewDefaultConfig() *Config {
 		SameSiteCookie:                constant.SameSiteLax,
 		Scopes:                        []string{"email", "profile"},
 		SecureCookie:                  true,
+		ServerGraceTimeout:            constant.DefaultServerGraceTimeout,
 		ServerIdleTimeout:             constant.DefaultServerIdleTimeout,
 		ServerReadTimeout:             constant.DefaultServerReadTimeout,
 		ServerWriteTimeout:            constant.DefaultServerWriteTimeout,

pkg/keycloak/config/config.go

@@ -287,6 +287,8 @@ type Config struct {
 	// MaxIdleConnsPerHost limits the number of idle connections maintained per host
 	MaxIdleConnsPerHost int `env:"MAX_IDLE_CONNS_PER_HOST" json:"max-idle-connections-per-host" usage:"limits the number of idle connections maintained per host" yaml:"max-idle-connections-per-host"`
 
+	// ServerGraceTimeout is the time to wait for the server to shutdown
+	ServerGraceTimeout time.Duration `env:"SERVER_GRACE_TIMEOUT" json:"server-grace-timeout" usage:"the server wait before closing the server" yaml:"server-grace-timeout"`
 	// ServerReadTimeout is the read timeout on the http server
 	ServerReadTimeout time.Duration `env:"SERVER_READ_TIMEOUT" json:"server-read-timeout" usage:"the server read timeout on the http server" yaml:"server-read-timeout"`
 	// ServerWriteTimeout is the write timeout on the http server
@@ -366,6 +368,7 @@ func NewDefaultConfig() *Config {
 		SameSiteCookie:                constant.SameSiteLax,
 		Scopes:                        []string{"email", "profile"},
 		SecureCookie:                  true,
+		ServerGraceTimeout:            constant.DefaultServerGraceTimeout,
 		ServerIdleTimeout:             constant.DefaultServerIdleTimeout,
 		ServerReadTimeout:             constant.DefaultServerReadTimeout,
 		ServerWriteTimeout:            constant.DefaultServerWriteTimeout,

pkg/keycloak/proxy/server.go

@@ -1006,6 +1006,18 @@ func (r *OauthProxy) Run() error {
 	return nil
 }
 
+// Shutdown finishes the proxy service with gracefully period
+func (r *OauthProxy) Shutdown() error {
+	ctx, cancel := context.WithTimeout(context.Background(), r.Config.ServerGraceTimeout)
+	defer cancel()
+
+	err := r.Server.Shutdown(ctx)
+	if err == nil {
+		return nil
+	}
+	return r.Server.Close()
+}
+
 // listenerConfig encapsulate listener options
 type listenerConfig struct {
 	ca                  string   // the path to a certificate authority

pkg/proxy/cli.go

@@ -90,6 +90,10 @@ func NewOauthProxyApp[T proxycore.KeycloakProvider | proxycore.GoogleProvider](p
 		signal.Notify(signalChannel, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
 		<-signalChannel
 
+		if err := proxy.Shutdown(); err != nil {
+			return utils.PrintError(err.Error())
+		}
+
 		return nil
 	}
 

pkg/proxy/core/core.go

@@ -33,6 +33,7 @@ func GetVersion() string {
 type OauthProxies interface {
 	CreateReverseProxy() error
 	Run() error
+	Shutdown() error
 }
 
 // ReverseProxy is a wrapper

pkg/testsuite/fake_proxy.go

@@ -82,6 +82,10 @@ type fakeProxy struct {
 	cookies map[string]*http.Cookie
 }
 
+func (f *fakeProxy) Shutdown() error {
+	return f.proxy.Shutdown()
+}
+
 func newFakeProxy(cfg *config.Config, authConfig *fakeAuthConfig) *fakeProxy {
 	log.SetOutput(io.Discard)
 

pkg/testsuite/fake_upstream.go

@@ -4,7 +4,9 @@ import (
 	"encoding/json"
 	"io"
 	"net/http"
+	"strconv"
 	"strings"
+	"time"
 
 	"golang.org/x/net/websocket"
 )
@@ -47,6 +49,21 @@ func (f *FakeUpstreamService) ServeHTTP(wrt http.ResponseWriter, req *http.Reque
 			wrt.WriteHeader(http.StatusInternalServerError)
 		}
 
+		var delay int
+		rawDelay := req.Header.Get("Delay")
+		if rawDelay != "" {
+			delay, err = strconv.Atoi(rawDelay)
+			if err != nil {
+				wrt.WriteHeader(http.StatusInternalServerError)
+			}
+		}
+
+		if delay > 0 {
+			// Sleep for the specified duration
+			// This is to simulate a slow upstream service
+			<-time.After(time.Duration(delay) * time.Second)
+		}
+
 		wrt.Header().Set(TestProxyAccepted, "true")
 		wrt.Header().Set("Content-Type", "application/json")
 		content, err := json.Marshal(&fakeUpstreamResponse{

pkg/testsuite/server_test.go

@@ -27,6 +27,7 @@ import (
 	"os"
 	"strconv"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
@@ -563,7 +564,7 @@ func TestSkipOpenIDProviderTLSVerifyForwardingProxy(t *testing.T) {
 	cfg.ForwardingUsername = ValidUsername
 	cfg.ForwardingPassword = ValidPassword
 	cfg.SkipOpenIDProviderTLSVerify = true
-	cfg.ForwardingGrantType = "password"
+	cfg.ForwardingGrantType = constant.ForwardingGrantTypePassword
 	s := httptest.NewServer(&FakeUpstreamService{})
 	requests := []fakeRequest{
 		{
@@ -2166,6 +2167,83 @@ func TestCustomHTTPMethod(t *testing.T) {
 	}
 }
 
+func TestGraceTimeout(t *testing.T) {
+	cfg := newFakeKeycloakConfig()
+	cfg.EnableForwarding = true
+	cfg.PatRetryCount = 5
+	cfg.PatRetryInterval = 2 * time.Second
+	cfg.OpenIDProviderTimeout = 30 * time.Second
+	cfg.ForwardingDomains = []string{}
+	cfg.ForwardingUsername = ValidUsername
+	cfg.ForwardingPassword = ValidPassword
+	cfg.SkipOpenIDProviderTLSVerify = true
+	cfg.ForwardingGrantType = constant.ForwardingGrantTypePassword
+
+	fakeServer := httptest.NewServer(&FakeUpstreamService{})
+
+	testCases := []struct {
+		Name                 string
+		ServerGraceTimeout   time.Duration
+		ResponseDelay        string
+		ExpectedCode         int
+		ExpectedRequestError string
+		ExpectedProxy        bool
+	}{
+		{
+			Name:                 "TestGraceTimeout",
+			ServerGraceTimeout:   2 * time.Second,
+			ResponseDelay:        "1",
+			ExpectedCode:         http.StatusOK,
+			ExpectedRequestError: "",
+			ExpectedProxy:        true,
+		},
+		{
+			Name:                 "TestGraceTimeoutClosedServer",
+			ServerGraceTimeout:   time.Second,
+			ResponseDelay:        "2",
+			ExpectedCode:         0,
+			ExpectedRequestError: "EOF",
+			ExpectedProxy:        false,
+		},
+	}
+
+	for _, testCase := range testCases {
+		t.Run(
+			testCase.Name,
+			func(t *testing.T) {
+				cfg.ServerGraceTimeout = testCase.ServerGraceTimeout
+
+				var waitShutdown sync.WaitGroup
+				waitShutdown.Add(1)
+
+				requests := []fakeRequest{
+					{
+						URL:                  fakeServer.URL + FakeTestURL,
+						ProxyRequest:         true,
+						ExpectedProxy:        testCase.ExpectedProxy,
+						ExpectedCode:         testCase.ExpectedCode,
+						ExpectedRequestError: testCase.ExpectedRequestError,
+						Headers:              map[string]string{"delay": testCase.ResponseDelay},
+					},
+				}
+
+				proxy := newFakeProxy(cfg, &fakeAuthConfig{EnableTLS: true})
+				<-time.After(time.Duration(100) * time.Millisecond)
+
+				go func() {
+					defer waitShutdown.Done()
+					<-time.After(time.Duration(200) * time.Millisecond)
+					if err := proxy.Shutdown(); err != nil {
+						t.Error("Failed to shutdown proxy")
+					}
+				}()
+				proxy.RunTests(t, requests)
+				waitShutdown.Wait()
+			},
+		)
+	}
+}
+
 // commented out because of see https://github.com/golang/go/issues/51416
 // func TestUpstreamProxy(t *testing.T) {
 // 	errChan := make(chan error)