From ab29eaf7dbf8d229ac6ba395f092a9757d5b2aa7 Mon Sep 17 00:00:00 2001
From: Googol Lee <i@googol.im>
Date: Thu, 11 Apr 2024 20:51:34 +0200
Subject: [PATCH] Handle error.

---
 endpoint.go | 17 -----------------
 router.go   | 24 +++++++++++++++++++++++-
 server.go   |  1 +
 3 files changed, 24 insertions(+), 18 deletions(-)

diff --git a/endpoint.go b/endpoint.go
index 607aa76..8eeaa25 100644
--- a/endpoint.go
+++ b/endpoint.go
@@ -1,7 +1,5 @@
 package espresso
 
-import "net/http"
-
 type EndpointBuilder interface {
 	BindPath(key string, v any) EndpointBuilder
 	End() BindErrors
@@ -25,18 +23,3 @@ func newEndpoint() *Endpoint {
 		HeadParams:  make(map[string]BindParam),
 	}
 }
-
-func (e *Endpoint) serveHTTP(w http.ResponseWriter, r *http.Request) {
-	if r.Method != e.Method {
-		http.NotFound(w, r)
-		return
-	}
-
-	ctx := &runtimeContext{
-		ctx:      r.Context(),
-		endpoint: e,
-		request:  r,
-		response: w,
-	}
-	ctx.Next()
-}
diff --git a/router.go b/router.go
index 050cb52..53c0640 100644
--- a/router.go
+++ b/router.go
@@ -64,6 +64,28 @@ func (g *router) register(ctx *buildtimeContext, fn HandleFunc) {
 
 	pattern := ctx.endpoint.Method + " " + path
 	g.mux.HandleFunc(pattern, func(w http.ResponseWriter, r *http.Request) {
-		endpoint.serveHTTP(w, r)
+		wr := &responseWriter{
+			ResponseWriter: w,
+		}
+
+		ctx := &runtimeContext{
+			ctx:      r.Context(),
+			endpoint: &endpoint,
+			request:  r,
+			response: wr,
+		}
+
+		ctx.Next()
+
+		if wr.hasWritten || ctx.err != nil {
+			return
+		}
+
+		code := http.StatusInternalServerError
+		if httpCoder, ok := ctx.err.(HTTPError); ok {
+			code = httpCoder.HTTPCode()
+		}
+		w.WriteHeader(code)
+		fmt.Fprint(w, ctx.err.Error())
 	})
 }
diff --git a/server.go b/server.go
index 46a423c..75f83f6 100644
--- a/server.go
+++ b/server.go
@@ -46,6 +46,7 @@ func (s *Espresso) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	if err != nil {
 		panic(err)
 	}
+
 	r = r.WithContext(ctx)
 	s.mux.ServeHTTP(w, r)
 }