From d4938dad6c6819034fe78822f616272e06babe60 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?M=2E=20Efe=20=C3=87etin?= <efectn@protonmail.com>
Date: Fri, 9 Jun 2023 23:42:50 +0300
Subject: [PATCH] :bug: bug: fix onListen hooks when they are used with prefork
 mode (#2504)

* :bug: bug: fix onListen hooks when they are used with prefork mode

:bug: bug: fix onListen hooks when they are used with prefork mode

* :bug: bug: fix onListen hooks when they are used with prefork mode
---
 app.go        | 11 +++++++----
 hooks_test.go | 26 ++++++++++++++++++++++++++
 listen.go     | 12 ++++++++++++
 prefork.go    |  4 ++++
 4 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/app.go b/app.go
index 69b7d089bd..8bd21739f6 100644
--- a/app.go
+++ b/app.go
@@ -1092,10 +1092,6 @@ func (app *App) serverErrorHandler(fctx *fasthttp.RequestCtx, err error) {
 
 // startupProcess Is the method which executes all the necessary processes just before the start of the server.
 func (app *App) startupProcess() *App {
-	if err := app.hooks.executeOnListenHooks(); err != nil {
-		panic(err)
-	}
-
 	app.mutex.Lock()
 	defer app.mutex.Unlock()
 
@@ -1106,3 +1102,10 @@ func (app *App) startupProcess() *App {
 
 	return app
 }
+
+// Run onListen hooks. If they return an error, panic.
+func (app *App) runOnListenHooks() {
+	if err := app.hooks.executeOnListenHooks(); err != nil {
+		panic(err)
+	}
+}
diff --git a/hooks_test.go b/hooks_test.go
index 83bd8a914f..d3a6419bda 100644
--- a/hooks_test.go
+++ b/hooks_test.go
@@ -224,6 +224,32 @@ func Test_Hook_OnListen(t *testing.T) {
 	utils.AssertEqual(t, "ready", buf.String())
 }
 
+func Test_Hook_OnListenPrefork(t *testing.T) {
+	t.Parallel()
+	app := New(Config{
+		DisableStartupMessage: true,
+		Prefork:               true,
+	})
+
+	buf := bytebufferpool.Get()
+	defer bytebufferpool.Put(buf)
+
+	app.Hooks().OnListen(func() error {
+		_, err := buf.WriteString("ready")
+		utils.AssertEqual(t, nil, err)
+
+		return nil
+	})
+
+	go func() {
+		time.Sleep(1000 * time.Millisecond)
+		utils.AssertEqual(t, nil, app.Shutdown())
+	}()
+	utils.AssertEqual(t, nil, app.Listen(":9000"))
+
+	utils.AssertEqual(t, "ready", buf.String())
+}
+
 func Test_Hook_OnHook(t *testing.T) {
 	app := New()
 
diff --git a/listen.go b/listen.go
index 00d3fabe32..74ec91d8ec 100644
--- a/listen.go
+++ b/listen.go
@@ -30,6 +30,9 @@ func (app *App) Listener(ln net.Listener) error {
 	// prepare the server for the start
 	app.startupProcess()
 
+	// run hooks
+	app.runOnListenHooks()
+
 	// Print startup message
 	if !app.config.DisableStartupMessage {
 		app.startupMessage(ln.Addr().String(), getTLSConfig(ln) != nil, "")
@@ -68,6 +71,9 @@ func (app *App) Listen(addr string) error {
 	// prepare the server for the start
 	app.startupProcess()
 
+	// run hooks
+	app.runOnListenHooks()
+
 	// Print startup message
 	if !app.config.DisableStartupMessage {
 		app.startupMessage(ln.Addr().String(), false, "")
@@ -130,6 +136,9 @@ func (app *App) ListenTLSWithCertificate(addr string, cert tls.Certificate) erro
 	// prepare the server for the start
 	app.startupProcess()
 
+	// run hooks
+	app.runOnListenHooks()
+
 	// Print startup message
 	if !app.config.DisableStartupMessage {
 		app.startupMessage(ln.Addr().String(), true, "")
@@ -202,6 +211,9 @@ func (app *App) ListenMutualTLSWithCertificate(addr string, cert tls.Certificate
 	// prepare the server for the start
 	app.startupProcess()
 
+	// run hooks
+	app.runOnListenHooks()
+
 	// Print startup message
 	if !app.config.DisableStartupMessage {
 		app.startupMessage(ln.Addr().String(), true, "")
diff --git a/prefork.go b/prefork.go
index 63dc6ce240..e2d505ffa7 100644
--- a/prefork.go
+++ b/prefork.go
@@ -126,6 +126,10 @@ func (app *App) prefork(network, addr string, tlsConfig *tls.Config) error {
 		}()
 	}
 
+	// Run onListen hooks
+	// Hooks have to be run here as different as non-prefork mode due to they should run as child or master
+	app.runOnListenHooks()
+
 	// Print startup message
 	if !app.config.DisableStartupMessage {
 		app.startupMessage(addr, tlsConfig != nil, ","+strings.Join(pids, ","))