Repository: gofr-dev/gofr

Issue:
Cannot Close websocket conneciton in case of unauthorised request
**Describe the bug**
I have validations set for my server to authorise users based on Bearer tokens, when using the websocket, The server is showing connection successfull and I get error logs for connection highjack because I am writing 401 errors to the response writter. 

Even If I ignore the un-authorized request and not move forward, the connection remains open.

**To Reproduce**
Add a middleware which returns and writes to the http.ResponseWriter and hijack can be reproduced.
Check for authorization and return without calling inner.ServeHTTP(w,r), the connection to websocket remains open.

1. The code is

   ```go
   
   ```

2. The error is

   ```
   
   ```

**Expected behavior**
I should be able to get the connection in middleware for the particular websocket request and close connection in case of unauthorised request

**Screenshots**
If applicable, add screenshots to help explain your problem.

**Environments (please complete the following information):**
- OS: [e.g. Linux]
- gofr version [e.g. v1.5.0]
- go version [e.g. 1.21]

**More description**
Add any other context about the problem here.


Discussion:

Solution patch:
diff --git a/pkg/gofr/container/mock_container.go b/pkg/gofr/container/mock_container.go
index 973e9d7a..2f492edb 100644
--- a/pkg/gofr/container/mock_container.go
+++ b/pkg/gofr/container/mock_container.go
@@ -2,8 +2,6 @@ package container
 
 import (
 	"context"
-	"fmt"
-	"net/http"
 	"testing"
 
 	"go.uber.org/mock/gomock"
@@ -33,7 +31,7 @@ type Mocks struct {
 
 type options func(c *Container, ctrl *gomock.Controller) any
 
-//nolint:revive //Because user should not access the options, and we might change it to an interface in the future.
+//nolint:revive // Because user should not access the options, and we might change it to an interface in the future.
 func WithMockHTTPService(httpServiceNames ...string) options {
 	return func(c *Container, ctrl *gomock.Controller) any {
 		mockservice := service.NewMockHTTP(ctrl)
@@ -124,8 +122,9 @@ func NewMockContainer(t *testing.T, options ...options) (*Container, *Mocks) {
 		Metrics:     mockMetrics,
 	}
 
-	mockMetrics.EXPECT().RecordHistogram(gomock.Any(), "app_http_service_response", gomock.Any(), "path", gomock.Any(),
-		"method", gomock.Any(), "status", fmt.Sprintf("%v", http.StatusInternalServerError)).AnyTimes()
+	// TODO: Remove this expectation from mock container (previous generalisation) to the actual tests where their expectations are being set.
+	mockMetrics.EXPECT().RecordHistogram(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(),
+		gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes()
 
 	return container, &mocks
 }
diff --git a/pkg/gofr/factory.go b/pkg/gofr/factory.go
index d5669aa7..0c12684a 100644
--- a/pkg/gofr/factory.go
+++ b/pkg/gofr/factory.go
@@ -8,7 +8,6 @@ import (
 
 	"gofr.dev/pkg/gofr/cmd/terminal"
 	"gofr.dev/pkg/gofr/container"
-	"gofr.dev/pkg/gofr/http/middleware"
 	"gofr.dev/pkg/gofr/logging"
 )
 
@@ -38,7 +37,7 @@ func New() *App {
 		port = defaultHTTPPort
 	}
 
-	app.httpServer = newHTTPServer(app.container, port, middleware.GetConfigs(app.Config))
+	app.httpServer = newHTTPServer(port)
 	app.httpServer.certFile = app.Config.GetOrDefault("CERT_FILE", "")
 	app.httpServer.keyFile = app.Config.GetOrDefault("KEY_FILE", "")
 
diff --git a/pkg/gofr/gofr.go b/pkg/gofr/gofr.go
index 14a47644..cbb272ac 100644
--- a/pkg/gofr/gofr.go
+++ b/pkg/gofr/gofr.go
@@ -23,6 +23,10 @@ import (
 	"gofr.dev/pkg/gofr/service"
 )
 
+const (
+	configLocation = "./configs"
+)
+
 // App is the main application in the GoFr framework.
 type App struct {
 	// Config can be used by applications to fetch custom configurations from environment or file.
@@ -136,18 +140,19 @@ func (a *App) startSubscriptions(ctx context.Context) error {
 
 // readConfig reads the configuration from the default location.
 func (a *App) readConfig(isAppCMD bool) {
-	var configLocation string
-	if _, err := os.Stat("./configs"); err == nil {
-		configLocation = "./configs"
+	var location string
+
+	if _, err := os.Stat(configLocation); err == nil {
+		location = configLocation
 	}
 
 	if isAppCMD {
-		a.Config = config.NewEnvFile(configLocation, logging.NewFileLogger(""))
+		a.Config = config.NewEnvFile(location, logging.NewFileLogger(""))
 
 		return
 	}
 
-	a.Config = config.NewEnvFile(configLocation, logging.NewLogger(logging.INFO))
+	a.Config = config.NewEnvFile(location, logging.NewLogger(logging.INFO))
 }
 
 // AddHTTPService registers HTTP service in container.
diff --git a/pkg/gofr/http_server.go b/pkg/gofr/http_server.go
index dab38f8c..ba340f73 100644
--- a/pkg/gofr/http_server.go
+++ b/pkg/gofr/http_server.go
@@ -29,18 +29,10 @@ var (
 	errInvalidKeyFile         = errors.New("invalid key file")
 )
 
-func newHTTPServer(c *container.Container, port int, middlewareConfigs map[string]string) *httpServer {
+func newHTTPServer(port int) *httpServer {
 	r := gofrHTTP.NewRouter()
 	wsManager := websocket.New()
 
-	r.Use(
-		middleware.WSHandlerUpgrade(c, wsManager),
-		middleware.Tracer,
-		middleware.Logging(c.Logger),
-		middleware.CORS(middlewareConfigs, r.RegisteredRoutes),
-		middleware.Metrics(c.Metrics()),
-	)
-
 	return &httpServer{
 		router: r,
 		port:   port,
@@ -70,7 +62,21 @@ func (s *httpServer) RegisterProfilingRoutes() {
 	s.router.NewRoute().Methods(http.MethodGet).PathPrefix("/debug/pprof/").HandlerFunc(pprof.Index)
 }
 
-func (s *httpServer) Run(c *container.Container) {
+func (s *httpServer) run(c *container.Container, middlewareConfigs map[string]string) {
+	// Developer Note:
+	//	WebSocket connections do not inherently support authentication mechanisms.
+	//	It is recommended to authenticate users before upgrading to a WebSocket connection.
+	//	Hence, we are registering middlewares here, to ensure that authentication or other
+	//	middleware logic is executed during the initial HTTP handshake request, prior to upgrading
+	//	the connection to WebSocket, if any.
+	s.router.Use(
+		middleware.WSHandlerUpgrade(c, s.ws),
+		middleware.Tracer,
+		middleware.CORS(middlewareConfigs, s.router.RegisteredRoutes),
+		middleware.Logging(c.Logger),
+		middleware.Metrics(c.Metrics()),
+	)
+
 	if s.srv != nil {
 		c.Logf("Server already running on port: %d", s.port)
 		return
diff --git a/pkg/gofr/run.go b/pkg/gofr/run.go
index ae0067a8..8b6045fa 100644
--- a/pkg/gofr/run.go
+++ b/pkg/gofr/run.go
@@ -7,6 +7,8 @@ import (
 	"os/signal"
 	"sync"
 	"syscall"
+
+	"gofr.dev/pkg/gofr/http/middleware"
 )
 
 // Run starts the application. If it is an HTTP server, it will start the server.
@@ -68,7 +70,7 @@ func (a *App) Run() {
 
 		go func(s *httpServer) {
 			defer wg.Done()
-			s.Run(a.container)
+			s.run(a.container, middleware.GetConfigs(a.Config))
 		}(a.httpServer)
 	}
 



Test patch:
diff --git a/pkg/gofr/context_test.go b/pkg/gofr/context_test.go
index 0c229f22..23ab4420 100644
--- a/pkg/gofr/context_test.go
+++ b/pkg/gofr/context_test.go
@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 
 	"github.com/golang-jwt/jwt/v5"
 	"github.com/gorilla/websocket"
@@ -78,42 +79,37 @@ func TestContext_AddTrace(t *testing.T) {
 
 func TestContext_WriteMessageToSocket(t *testing.T) {
 	port := testutil.GetFreePort(t)
-
 	t.Setenv("HTTP_PORT", fmt.Sprint(port))
 
 	app := New()
 
-	server := httptest.NewServer(app.httpServer.router)
-	defer server.Close()
-
 	app.WebSocket("/ws", func(ctx *Context) (any, error) {
 		socketErr := ctx.WriteMessageToSocket("Hello! GoFr")
 		if socketErr != nil {
 			return nil, socketErr
 		}
 
-		// TODO: returning error here to close the connection to the websocket
-		// as the websocket close error is not caught because we are using no bind function here.
-		// this must not be necessary. We should put an actual check in handleWebSocketConnection method instead.
-		return nil, &websocket.CloseError{Code: websocket.CloseNormalClosure, Text: "Error closing"}
+		conn := ctx.GetConnectionFromContext(ctx)
+		defer conn.Close()
+
+		return "", socketErr
 	})
 
 	go app.Run()
+	time.Sleep(100 * time.Millisecond) // Wait for the server to boot
 
-	wsURL := "ws" + server.URL[len("http"):] + "/ws"
+	wsURL := fmt.Sprintf("ws://localhost:%d/ws", port)
 
-	// Create a WebSocket client
 	ws, resp, err := websocket.DefaultDialer.Dial(wsURL, nil)
-	require.NoError(t, err)
+	require.NoError(t, err, "WebSocket handshake failed")
 
 	defer resp.Body.Close()
 	defer ws.Close()
 
 	_, message, err := ws.ReadMessage()
-	require.NoError(t, err)
+	require.NoError(t, err, "Failed to read WebSocket message")
 
-	expectedResponse := "Hello! GoFr"
-	assert.Equal(t, expectedResponse, string(message))
+	assert.Equal(t, "Hello! GoFr", string(message))
 }
 
 func TestGetAuthInfo_BasicAuth(t *testing.T) {
diff --git a/pkg/gofr/http_server_test.go b/pkg/gofr/http_server_test.go
index b0190fd8..16d5b3c1 100644
--- a/pkg/gofr/http_server_test.go
+++ b/pkg/gofr/http_server_test.go
@@ -10,11 +10,14 @@ import (
 	"testing"
 	"time"
 
+	"github.com/gorilla/mux"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"gofr.dev/pkg/gofr/config"
 	"gofr.dev/pkg/gofr/container"
 	gofrHTTP "gofr.dev/pkg/gofr/http"
+	"gofr.dev/pkg/gofr/http/middleware"
 	"gofr.dev/pkg/gofr/logging"
 	"gofr.dev/pkg/gofr/testutil"
 )
@@ -28,10 +31,24 @@ func TestRun_ServerStartsListening(t *testing.T) {
 		w.WriteHeader(http.StatusOK)
 	}))
 
+	// adding registered routes for applying middlewares
+	var registeredMethods []string
+
+	_ = router.Walk(func(route *mux.Route, _ *mux.Router, _ []*mux.Route) error {
+		met, _ := route.GetMethods()
+		for _, method := range met {
+			if !contains(registeredMethods, method) { // Check for uniqueness before adding
+				registeredMethods = append(registeredMethods, method)
+			}
+		}
+
+		return nil
+	})
+
+	router.RegisteredRoutes = &registeredMethods
+
 	// Create a mock container
-	c := &container.Container{
-		Logger: logging.NewLogger(logging.INFO),
-	}
+	c := container.NewContainer(getConfigs(t))
 
 	// Create an instance of httpServer
 	server := &httpServer{
@@ -40,7 +57,7 @@ func TestRun_ServerStartsListening(t *testing.T) {
 	}
 
 	// Start the server
-	go server.Run(c)
+	go server.run(c, middleware.GetConfigs(getConfigs(t)))
 
 	// Wait for the server to start listening
 	time.Sleep(100 * time.Millisecond)
@@ -61,6 +78,18 @@ func TestRun_ServerStartsListening(t *testing.T) {
 	resp.Body.Close()
 }
 
+func getConfigs(t *testing.T) config.Config {
+	t.Helper()
+
+	var configLocation string
+
+	if _, err := os.Stat("./configs"); err == nil {
+		configLocation = "./configs"
+	}
+
+	return config.NewEnvFile(configLocation, logging.NewLogger(logging.INFO))
+}
+
 func TestRegisterProfillingRoutes(t *testing.T) {
 	port := testutil.GetFreePort(t)
 
@@ -75,7 +104,7 @@ func TestRegisterProfillingRoutes(t *testing.T) {
 
 	server.RegisterProfilingRoutes()
 
-	go server.Run(c)
+	go server.run(c, middleware.GetConfigs(getConfigs(t)))
 
 	// Test if the expected handlers are registered for the pprof endpoints
 	expectedRoutes := []string{
@@ -114,7 +143,7 @@ func TestShutdown_ServerStopsListening(t *testing.T) {
 	}
 
 	// Start the server
-	go server.Run(c)
+	go server.run(c, middleware.GetConfigs(getConfigs(t)))
 
 	// Create a context with a timeout to test the shutdown
 	ctx, cancel := context.WithTimeout(context.Background(), 150*time.Millisecond)
@@ -150,7 +179,7 @@ func TestShutdown_ServerContextDeadline(t *testing.T) {
 	}
 
 	// Start the server
-	go server.Run(c)
+	go server.run(c, middleware.GetConfigs(getConfigs(t)))
 
 	// Create a context with a timeout to test the shutdown
 	ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)