diff --git a/context.go b/context.go new file mode 100644 index 0000000..4516fe7 --- /dev/null +++ b/context.go @@ -0,0 +1,68 @@ +package chi + +import "context" + +type contextKey struct{} + +var RouteCtxKey = &contextKey{} + +type Context struct { + parentCtx context.Context + URLParams RouteParams +} + +type RouteParams struct { + Keys []string + Values []string +} + +func NewRouteContext() *Context { + return &Context{} +} + +func (x *Context) Reset() { + x.parentCtx = nil + x.URLParams.Keys = x.URLParams.Keys[:0] + x.URLParams.Values = x.URLParams.Values[:0] +} + +func (x *Context) WithParent(ctx context.Context) { + x.parentCtx = ctx +} + +func (x *Context) URLParam(key string) string { + for i := len(x.URLParams.Keys) - 1; i >= 0; i-- { + if x.URLParams.Keys[i] == key { + return x.URLParams.Values[i] + } + } + return "" +} + +func (x *Context) pushURLParam(key string, value string) { + x.URLParams.Keys = append(x.URLParams.Keys, key) + x.URLParams.Values = append(x.URLParams.Values, value) +} + +func RouteContext(ctx context.Context) *Context { + if ctx == nil { + return nil + } + if rctx, ok := ctx.Value(RouteCtxKey).(*Context); ok { + return rctx + } + return nil +} + +func contextWithRouteContext(parent context.Context, rctx *Context) context.Context { + rctx.WithParent(parent) + return context.WithValue(parent, RouteCtxKey, rctx) +} + +func URLParam(r interface{ Context() context.Context }, key string) string { + rctx := RouteContext(r.Context()) + if rctx == nil { + return "" + } + return rctx.URLParam(key) +} diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..872b418 --- /dev/null +++ b/go.mod @@ -0,0 +1,3 @@ +module github.com/RoseMark45/chi + +go 1.22 diff --git a/main.go b/main.go index 49f4dee..b924384 100644 --- a/main.go +++ b/main.go @@ -1,7 +1 @@ -package main - -import "fmt" - -func main() { - fmt.Println("Hello, Bounty Hunter!") -} +package chi diff --git a/mux.go b/mux.go new file mode 100644 index 0000000..c0d449f --- /dev/null +++ b/mux.go @@ -0,0 +1,142 @@ +package chi + +import ( + "net/http" + "strings" + "sync" +) + +type Router interface { + Use(middlewares ...func(http.Handler) http.Handler) + Get(pattern string, handlerFn http.HandlerFunc) + Route(pattern string, fn func(r Router)) + http.Handler +} + +type Mux struct { + pool sync.Pool + middleware []func(http.Handler) http.Handler + routes []route + prefix string +} + +type route struct { + method string + pattern string + handler http.Handler +} + +func NewRouter() *Mux { + mx := &Mux{} + mx.pool.New = func() any { + return NewRouteContext() + } + return mx +} + +func (mx *Mux) Use(middlewares ...func(http.Handler) http.Handler) { + mx.middleware = append(mx.middleware, middlewares...) +} + +func (mx *Mux) Get(pattern string, handlerFn http.HandlerFunc) { + mx.handle(http.MethodGet, pattern, handlerFn) +} + +func (mx *Mux) Route(pattern string, fn func(r Router)) { + sub := &Mux{ + pool: mx.pool, + middleware: append([]func(http.Handler) http.Handler{}, mx.middleware...), + prefix: joinPatterns(mx.prefix, pattern), + } + fn(sub) + mx.routes = append(mx.routes, sub.routes...) +} + +func (mx *Mux) ServeHTTP(w http.ResponseWriter, r *http.Request) { + rctx := RouteContext(r.Context()) + ownedRouteContext := false + if rctx == nil { + rctx = mx.pool.Get().(*Context) + rctx.Reset() + ownedRouteContext = true + r = r.WithContext(contextWithRouteContext(r.Context(), rctx)) + } + + if ownedRouteContext { + defer mx.pool.Put(rctx) + } + + for _, route := range mx.routes { + if route.method != r.Method { + continue + } + if !matchPattern(route.pattern, r.URL.Path, rctx) { + continue + } + route.handler.ServeHTTP(w, r) + return + } + + http.NotFound(w, r) +} + +func (mx *Mux) handle(method string, pattern string, handler http.Handler) { + fullPattern := joinPatterns(mx.prefix, pattern) + for i := len(mx.middleware) - 1; i >= 0; i-- { + handler = mx.middleware[i](handler) + } + mx.routes = append(mx.routes, route{ + method: method, + pattern: fullPattern, + handler: handler, + }) +} + +func joinPatterns(left string, right string) string { + switch { + case left == "": + return right + case right == "": + return left + case strings.HasSuffix(left, "/") && strings.HasPrefix(right, "/"): + return left + strings.TrimPrefix(right, "/") + case !strings.HasSuffix(left, "/") && !strings.HasPrefix(right, "/"): + return left + "/" + right + default: + return left + right + } +} + +func matchPattern(pattern string, path string, rctx *Context) bool { + patternParts := splitPath(pattern) + pathParts := splitPath(path) + if len(patternParts) != len(pathParts) { + return false + } + + startLenKeys := len(rctx.URLParams.Keys) + startLenValues := len(rctx.URLParams.Values) + + for i, patternPart := range patternParts { + if strings.HasPrefix(patternPart, "{") && strings.HasSuffix(patternPart, "}") { + key := strings.TrimSuffix(strings.TrimPrefix(patternPart, "{"), "}") + rctx.pushURLParam(key, pathParts[i]) + continue + } + if patternPart != pathParts[i] { + rctx.URLParams.Keys = rctx.URLParams.Keys[:startLenKeys] + rctx.URLParams.Values = rctx.URLParams.Values[:startLenValues] + return false + } + } + + return true +} + +func splitPath(path string) []string { + trimmed := strings.Trim(path, "/") + if trimmed == "" { + return nil + } + return strings.Split(trimmed, "/") +} diff --git a/mux_test.go b/mux_test.go new file mode 100644 index 0000000..0dfb0ff --- /dev/null +++ b/mux_test.go @@ -0,0 +1,75 @@ +package chi + +import ( + "context" + "net/http" + "net/http/httptest" + "testing" +) + +func TestMiddlewareRxCloneURLParams(t *testing.T) { + r := NewRouter() + r.Route("/users/{userID}", func(r Router) { + r.Use(func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + ctx := context.WithValue(req.Context(), "customKey", "customValue") + next.ServeHTTP(w, req.WithContext(ctx)) + }) + }) + r.Get("/profile", func(w http.ResponseWriter, req *http.Request) { + userID := URLParam(req, "userID") + if userID != "123" { + t.Errorf("expected userID to be '123', got '%s'", userID) + } + if RouteContext(req.Context()) == nil { + t.Error("expected active route context after request cloning") + } + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("ok")) + }) + }) + + ts := httptest.NewServer(r) + defer ts.Close() + + res, err := http.Get(ts.URL + "/users/123/profile") + if err != nil { + t.Fatal(err) + } + defer res.Body.Close() + if res.StatusCode != http.StatusOK { + t.Errorf("expected status OK, got %d", res.StatusCode) + } +} + +func TestNestedMiddlewareClonePreservesURLParams(t *testing.T) { + r := NewRouter() + r.Route("/orgs/{orgID}", func(r Router) { + r.Route("/users/{userID}", func(r Router) { + r.Use(func(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) { + ctx := context.WithValue(req.Context(), "traceID", "abc") + next.ServeHTTP(w, req.WithContext(ctx)) + }) + }) + r.Get("/profile", func(w http.ResponseWriter, req *http.Request) { + if got := URLParam(req, "orgID"); got != "acme" { + t.Fatalf("expected orgID acme, got %q", got) + } + if got := URLParam(req, "userID"); got != "123" { + t.Fatalf("expected userID 123, got %q", got) + } + w.WriteHeader(http.StatusOK) + }) + }) + }) + + req := httptest.NewRequest(http.MethodGet, "/orgs/acme/users/123/profile", nil) + rec := httptest.NewRecorder() + + r.ServeHTTP(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("expected status 200, got %d", rec.Code) + } +}