Fiber example with proper context cancel

listener.go

package http

import (
	"net"
	"sync"
)

// listen is a wrapper over [net.Listener]
type listen struct {
	net.Listener
}

// Accept wraps connection
func (l *listen) Accept() (net.Conn, error) {
	conn, err := l.Listener.Accept()
	if err != nil {
		return nil, err
	}
	return &connection{
		Conn: conn,
	}, nil
}

// connection is a wrapper over [net.Conn]
type connection struct {
	net.Conn
	sync.Mutex

	n int
	b []byte
}

// Read reads data from connection and handles a special case for connection status check
func (a *connection) Read(b []byte) (n int, err error) {
	a.Lock()
	defer a.Unlock()

	// Special case to check if connection is still active
	if len(b) == 0 {
		bp := make([]byte, 1)

		n, err = a.Conn.Read(bp)
		if err != nil || n == 0 {
			return 0, err
		}

		if a.b == nil {
			a.b = bp
		} else {
			a.b = append(a.b, bp[0])
		}
		a.n += n

		return 0, nil
	}

	// Return buffered data back
	if a.n > 0 {
		n = copy(b, a.b)
		a.b = a.b[n:]
		a.n -= n
		return n, nil
	}

	return a.Conn.Read(b)
}

middleware_context.go

package http

import (
	"context"

	"github.com/gofiber/fiber/v2"
	"github.com/gofiber/websocket/v2"
)

// middlewareContext cancels context if connection was closed
func middlewareContext(noCancel bool) fiber.Handler {
	return func(fCtx *fiber.Ctx) error {
		if websocket.IsWebSocketUpgrade(fCtx) || noCancel {
			fCtx.SetUserContext(fCtx.Context())
			return fCtx.Next()
		}

		conn := fCtx.Context().Conn()

		ctx, cancel := context.WithCancel(fCtx.Context())
		go func() {
			_, _ = conn.Read(nil)
			cancel()
		}()
		fCtx.SetUserContext(ctx)

		return fCtx.Next()
	}
}