fix: lots of fixes and improvements

Author: ConsoleTVs

contract/session.go

@@ -58,6 +58,10 @@ type Session interface {
 	// HasRegenerated returns true if the session has been regenerated (i.e., the session ID
 	// has changed). This is useful for sending updated session identifiers to the client.
 	HasRegenerated() bool
+
+	// MarkAsUnchanged sets the session as if nothing has changed, therefore avoiding saving
+	// the session when the request finishes.
+	MarkAsUnchanged()
 }
 
 // SessionDriver defines the interface for persisting and retrieving session data.

framework/README.md

@@ -55,11 +55,7 @@ func handler(w http.ResponseWriter, r *http.Request) error {
 func main() {
 	app := framework.New()
 
-	app.Use(middleware.ErrorHandler(middleware.ErrorHandlerOptions{
-		Logger: slog.Default(),
-		IsDev:  true,
-	}))
-
+	app.Use(middleware.Logger(slog.Default()))
 	app.Use(middleware.Recover())
 
 	app.Get("/", handler)

framework/cache/redis/cache.go

@@ -20,8 +20,8 @@ func New(options *Options) *Client {
 }
 
 // Get retrieves a value by key or returns contract.ErrNotFound if missing.
-func (c *Client) Get(ctx context.Context, key string) (v any, e error) {
-	encoded, err := (*redis.Client)(c).Get(ctx, key).Result()
+func (c *Client) Get(ctx context.Context, key string) (any, error) {
+	v, err := (*redis.Client)(c).Get(ctx, key).Result()
 
 	if errors.Is(err, redis.Nil) {
 		return nil, fmt.Errorf("%w: %s", contract.ErrCacheKeyNotFound, key)
@@ -31,22 +31,12 @@ func (c *Client) Get(ctx context.Context, key string) (v any, e error) {
 		return nil, err
 	}
 
-	if err := json.Unmarshal([]byte(encoded), &v); err != nil {
-		return nil, err
-	}
-
 	return v, nil
 }
 
 // Put sets a key with value and TTL.
 func (c *Client) Put(ctx context.Context, key string, value any, ttl time.Duration) error {
-	encoded, err := json.Marshal(value)
-
-	if err != nil {
-		return err
-	}
-
-	return (*redis.Client)(c).Set(ctx, key, encoded, ttl).Err()
+	return (*redis.Client)(c).Set(ctx, key, value, ttl).Err()
 }
 
 // Delete removes a key.

framework/go.mod

@@ -20,6 +20,7 @@ require (
 	github.com/davecgh/go-spew v1.1.2-0.20180830191138-d8f796af33cc // indirect
 	github.com/dgryski/go-rendezvous v0.0.0-20200823014737-9f7001d12a5f // indirect
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
+	github.com/stretchr/objx v0.5.2 // indirect
 	golang.org/x/sys v0.38.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )

framework/go.sum

@@ -28,6 +28,8 @@ github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 h1:Jamvg5psRI
 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/redis/go-redis/v9 v9.17.2 h1:P2EGsA4qVIM3Pp+aPocCJ7DguDHhqrXNhVcEp4ViluI=
 github.com/redis/go-redis/v9 v9.17.2/go.mod h1:u410H11HMLoB+TP67dz8rL9s6QW2j76l0//kSOd3370=
+github.com/stretchr/objx v0.5.2 h1:xuMeJ0Sdp5ZMRXx/aWO6RZxdr3beISkG5/G/aIRr3pY=
+github.com/stretchr/objx v0.5.2/go.mod h1:FRsXN1f5AsAjCGJKqEizvkpNtU+EGNCLh3NxZ/8L+MA=
 github.com/stretchr/testify v1.11.1 h1:7s2iGBzp5EwR7/aIZr8ao5+dra3wiQyKjjFuvgVKu7U=
 github.com/stretchr/testify v1.11.1/go.mod h1:wZwfW3scLgRK+23gO65QZefKpKQRnfz6sD981Nm4B6U=
 github.com/studiolambda/cosmos/contract v0.5.0 h1:wGJj6k3Yhv1jeZvmxcDFsuSUyYTxV88S6n/5Gh4njBQ=

framework/handler.go

@@ -5,6 +5,9 @@ import (
 	"errors"
 	"net/http"
 	"net/http/httptest"
+
+	"github.com/studiolambda/cosmos/framework/hook"
+	"github.com/studiolambda/cosmos/problem"
 )
 
 // Handler defines the function signature for HTTP request handlers in Cosmos.
@@ -50,6 +53,28 @@ type HTTPStatus interface {
 	HTTPStatus() int
 }
 
+func handleError(w http.ResponseWriter, r *http.Request, err error) {
+	status := http.StatusInternalServerError
+
+	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+		status = 499 // A non-standard status code: 499 Client Closed Request
+	}
+
+	if s, ok := err.(HTTPStatus); ok {
+		status = s.HTTPStatus()
+	}
+
+	// We can check if the error can be directly
+	// handled by using a [http.Handler], in the case
+	// we'll simply handle it using ServeHTTP.
+	if h, ok := err.(http.Handler); ok {
+		h.ServeHTTP(w, r)
+		return
+	}
+
+	problem.NewProblem(err, status).ServeHTTP(w, r)
+}
+
 // ServeHTTP implements the http.Handler interface, allowing Cosmos handlers
 // to be used with the standard HTTP server. It bridges the gap between
 // Cosmos's error-returning handlers and Go's standard http.Handler interface.
@@ -59,18 +84,21 @@ type HTTPStatus interface {
 // response to the client. However, if the response has already been partially
 // written (e.g., during streaming), the error response may not be deliverable.
 func (handler Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
-	if err := handler(w, r); err != nil {
-		status := http.StatusInternalServerError
+	hooks := hook.NewManager()
+	wrapped := hook.NewResponseWriter(w, hooks)
+	ctx := context.WithValue(r.Context(), hook.Key, hooks)
+	err := handler(wrapped, r.WithContext(ctx))
 
-		if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
-			status = 499 // A non-standard status code: 499 Client Closed Request
-		}
+	if err != nil {
+		handleError(w, r, err)
+	}
 
-		if s, ok := err.(HTTPStatus); ok {
-			status = s.HTTPStatus()
-		}
+	if !wrapped.WriteHeaderCalled() {
+		wrapped.WriteHeader(http.StatusNoContent)
+	}
 
-		http.Error(w, err.Error(), status)
+	for _, callback := range hooks.AfterResponseFuncs() {
+		callback(err)
 	}
 }
 

framework/hook/manager.go

@@ -0,0 +1,85 @@
+package hook
+
+import (
+	"net/http"
+	"slices"
+	"sync"
+)
+
+// Manager is a structure that can be used to
+// attach specific hooks to the request / response
+// lifecycle. It requires the use of the hooks middleware.
+//
+// It is safe for concurrent use.
+type Manager struct {
+	mutex                  sync.Mutex
+	afterResponseFuncs     []AfterResponseFunc
+	beforeWriteHeaderFuncs []BeforeWriteHeaderFunc
+	beforeWriteFuncs       []BeforeWriteFunc
+}
+
+type AfterResponseFunc func(err error)
+type BeforeWriteHeaderFunc func(w http.ResponseWriter, status int)
+type BeforeWriteFunc func(w http.ResponseWriter, content []byte)
+type key struct{}
+
+var Key = key{}
+
+func NewManager() *Manager {
+	return &Manager{
+		mutex:                  sync.Mutex{},
+		beforeWriteHeaderFuncs: []BeforeWriteHeaderFunc{},
+		beforeWriteFuncs:       []BeforeWriteFunc{},
+	}
+}
+
+func (h *Manager) BeforeWriteHeader(callbacks ...BeforeWriteHeaderFunc) {
+	h.mutex.Lock()
+	defer h.mutex.Unlock()
+
+	h.beforeWriteHeaderFuncs = append(h.beforeWriteHeaderFuncs, callbacks...)
+}
+
+func (h *Manager) BeforeWriteHeaderFuncs() []BeforeWriteHeaderFunc {
+	h.mutex.Lock()
+	defer h.mutex.Unlock()
+
+	clone := slices.Clone(h.beforeWriteHeaderFuncs)
+	slices.Reverse(clone)
+
+	return clone
+}
+
+func (h *Manager) BeforeWrite(callbacks ...BeforeWriteFunc) {
+	h.mutex.Lock()
+	defer h.mutex.Unlock()
+
+	h.beforeWriteFuncs = append(h.beforeWriteFuncs, callbacks...)
+}
+
+func (h *Manager) BeforeWriteFuncs() []BeforeWriteFunc {
+	h.mutex.Lock()
+	defer h.mutex.Unlock()
+
+	clone := slices.Clone(h.beforeWriteFuncs)
+	slices.Reverse(clone)
+
+	return clone
+}
+
+func (h *Manager) AfterResponse(callbacks ...AfterResponseFunc) {
+	h.mutex.Lock()
+	defer h.mutex.Unlock()
+
+	h.afterResponseFuncs = append(h.afterResponseFuncs, callbacks...)
+}
+
+func (h *Manager) AfterResponseFuncs() []AfterResponseFunc {
+	h.mutex.Lock()
+	defer h.mutex.Unlock()
+
+	clone := slices.Clone(h.afterResponseFuncs)
+	slices.Reverse(clone)
+
+	return clone
+}

framework/hook/writer.go

@@ -0,0 +1,65 @@
+package hook
+
+import "net/http"
+
+type ResponseWriter struct {
+	http.ResponseWriter
+	*Manager
+	writeHeaderCalled bool
+}
+
+type ResponseWriterFlusher struct {
+	*ResponseWriter
+	http.Flusher
+}
+
+type WrappedResponseWriter interface {
+	http.ResponseWriter
+	WriteHeaderCalled() bool
+}
+
+func NewResponseWriter(w http.ResponseWriter, m *Manager) WrappedResponseWriter {
+	wrapped := &ResponseWriter{
+		ResponseWriter: w,
+		Manager:        m,
+	}
+
+	if f, ok := w.(http.Flusher); ok {
+		return &ResponseWriterFlusher{
+			ResponseWriter: wrapped,
+			Flusher:        f,
+		}
+	}
+
+	return wrapped
+}
+
+func (w *ResponseWriter) WriteHeaderCalled() bool {
+	return w.writeHeaderCalled
+}
+
+func (w *ResponseWriter) WriteHeader(status int) {
+	if w.WriteHeaderCalled() {
+		return
+	}
+
+	for _, hook := range w.Manager.BeforeWriteHeaderFuncs() {
+		hook(w.ResponseWriter, status)
+	}
+
+	w.ResponseWriter.WriteHeader(status)
+	w.writeHeaderCalled = true
+}
+
+func (w *ResponseWriter) Write(content []byte) (int, error) {
+	if !w.WriteHeaderCalled() {
+		// Same behaviour as the [http.ResponseWriter]
+		w.WriteHeader(http.StatusOK)
+	}
+
+	for _, hook := range w.Manager.BeforeWriteFuncs() {
+		hook(w.ResponseWriter, content)
+	}
+
+	return w.ResponseWriter.Write(content)
+}