diff --git a/internal/service/connect_server.go b/internal/service/connect_server.go index 6293efd2..b5a19b6a 100644 --- a/internal/service/connect_server.go +++ b/internal/service/connect_server.go @@ -9,6 +9,7 @@ import ( "net/http" "github.com/bufbuild/connect-go" + hx "github.com/rocboss/paopao-ce/pkg/http" ) var ( @@ -24,7 +25,12 @@ type connectServer struct { keyFile string handlerOpts []connect.HandlerOption server *http.Server - mux *http.ServeMux + mux connectMux +} + +type connectMux interface { + http.Handler + Handle(string, http.Handler) } func (s *connectServer) start() error { @@ -44,12 +50,18 @@ func (s *connectServer) register(path string, handler http.Handler) { s.mux.Handle(path, handler) } -func defaultConnectServer(addr string) *connectServer { - return &connectServer{ +func defaultConnectServer(addr string) (s *connectServer) { + s = &connectServer{ baseServer: newBaseServe(), server: &http.Server{ Addr: addr, }, mux: &http.ServeMux{}, } + // TODO: custom value from config + var useConnectMux bool + if useConnectMux { + s.mux = hx.NewConnectMux() + } + return } diff --git a/pkg/http/http.go b/pkg/http/http.go new file mode 100644 index 00000000..efb096d7 --- /dev/null +++ b/pkg/http/http.go @@ -0,0 +1,5 @@ +// Copyright 2023 Michael Li . All rights reserved. +// Use of this source code is governed by Apache License 2.0 that +// can be found in the LICENSE file. + +package http diff --git a/pkg/http/http_suite_test.go b/pkg/http/http_suite_test.go new file mode 100644 index 00000000..ef72f36b --- /dev/null +++ b/pkg/http/http_suite_test.go @@ -0,0 +1,17 @@ +// Copyright 2023 Michael Li . All rights reserved. +// Use of this source code is governed by Apache License 2.0 that +// can be found in the LICENSE file. + +package http_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +func TestHttp(t *testing.T) { + RegisterFailHandler(Fail) + RunSpecs(t, "Http Suite") +} diff --git a/pkg/http/mux.go b/pkg/http/mux.go new file mode 100644 index 00000000..19d1195e --- /dev/null +++ b/pkg/http/mux.go @@ -0,0 +1,130 @@ +// Copyright 2023 Michael Li . All rights reserved. +// Use of this source code is governed by Apache License 2.0 that +// can be found in the LICENSE file. + +package http + +import ( + "net/http" + "strings" + "sync" +) + +// ConnectMux mux used for Connect +type ConnectMux struct { + mu sync.RWMutex + m muxMap[http.Handler] +} + +type muxMap[T any] interface { + set(path string, val T) bool + get(path string) (T, bool) + match(pattern string) (T, bool) +} + +type simpleMuxMap[T any] map[string]T + +type prefixMuxMap[T any] struct { + prefix string + m map[string]T +} + +func (m simpleMuxMap[T]) set(path string, val T) bool { + if _, exist := m[path]; exist { + return false + } + m[path] = val + return true +} + +func (m simpleMuxMap[T]) get(path string) (val T, exist bool) { + val, exist = m[path] + return +} + +// match assume pattern like `/core.v1.AuthenticateService/login` +func (m simpleMuxMap[T]) match(pattern string) (val T, exist bool) { + idx := strings.IndexByte(pattern[1:], '/') + if idx < 0 { + return + } + return m.get(pattern[:idx+2]) +} + +func (m *prefixMuxMap[T]) set(path string, val T) bool { + if _, exist := m.m[path]; exist { + return false + } + m.m[path] = val + return true +} + +func (m *prefixMuxMap[T]) get(path string) (val T, exist bool) { + val, exist = m.m[path] + return +} + +// match assume pattern like `/core.v1.AuthenticateService/login` +func (m *prefixMuxMap[T]) match(pattern string) (val T, exist bool) { + path, _ := strings.CutPrefix(pattern, m.prefix) + idx := strings.IndexByte(path[1:], '/') + if idx < 0 { + return + } + return m.get(path[:idx+2]) +} + +// ServeHTTP dispatches the request to the handler whose +// pattern most closely matches the request URL. +func (mux *ConnectMux) ServeHTTP(w http.ResponseWriter, r *http.Request) { + if r.RequestURI == "*" { + if r.ProtoAtLeast(1, 1) { + w.Header().Set("Connection", "close") + } + w.WriteHeader(http.StatusBadRequest) + return + } + h := mux.handler(r.URL.Path) + h.ServeHTTP(w, r) +} + +// Handle registers the handler for the given path. +// If a handler already exists for path, Handle panics. +func (mux *ConnectMux) Handle(path string, handler http.Handler) { + mux.mu.Lock() + defer mux.mu.Unlock() + + if path == "" { + panic("http: invalid pattern") + } + if handler == nil { + panic("http: nil handler") + } + if ok := mux.m.set(path, handler); !ok { + panic("http: multiple registrations for " + path) + } +} + +// handler is the main implementation of Handler. +// The path is known to be in canonical form, except for CONNECT methods. +func (mux *ConnectMux) handler(path string) (h http.Handler) { + mux.mu.RLock() + defer mux.mu.RUnlock() + + if h, _ = mux.m.match(path); h == nil { + h = http.NotFoundHandler() + } + return +} + +// NewConnectMux allocates and returns a new ConnectMux. +func NewConnectMux(pathPrefix ...string) *ConnectMux { + var m muxMap[http.Handler] = make(simpleMuxMap[http.Handler]) + if len(pathPrefix) > 0 { + m = &prefixMuxMap[http.Handler]{ + m: make(map[string]http.Handler), + prefix: pathPrefix[0], + } + } + return &ConnectMux{m: m} +} diff --git a/pkg/http/mux_test.go b/pkg/http/mux_test.go new file mode 100644 index 00000000..d3f9b3c9 --- /dev/null +++ b/pkg/http/mux_test.go @@ -0,0 +1,71 @@ +// Copyright 2023 Michael Li . All rights reserved. +// Use of this source code is governed by Apache License 2.0 that +// can be found in the LICENSE file. + +package http + +import ( + g "github.com/onsi/ginkgo/v2" + m "github.com/onsi/gomega" +) + +var _ = g.Describe("Mux", g.Ordered, func() { + var smm muxMap[int] + var pmm muxMap[int] + + g.BeforeAll(func() { + smm = make(simpleMuxMap[int]) + pmm = &prefixMuxMap[int]{ + m: make(map[string]int), + prefix: "/connect", + } + }) + + g.It("simple mux map", func() { + ok := smm.set("/core.v1.AuthenticateService/", 1) + m.Expect(ok).To(m.BeTrue()) + + ok = smm.set("/core.v1.AuthenticateService/", 2) + m.Expect(ok).To(m.BeFalse()) + + smm.set("/greet.v1.GreetService/", 2) + val, exist := smm.get("/greet.v1.GreetService/") + m.Expect(val).To(m.Equal(2)) + m.Expect(exist).To(m.BeTrue()) + + _, exist = smm.get("/greet.v1.OtherService/") + m.Expect(exist).To(m.BeFalse()) + + val, exist = smm.match("/core.v1.AuthenticateService/login") + m.Expect(val).To(m.Equal(1)) + m.Expect(exist).To(m.BeTrue()) + + val, exist = smm.match("/core.v1.AuthenticateService/logout") + m.Expect(val).To(m.Equal(1)) + m.Expect(exist).To(m.BeTrue()) + }) + + g.It("prefix mux map", func() { + ok := pmm.set("/core.v1.AuthenticateService/", 1) + m.Expect(ok).To(m.BeTrue()) + + ok = pmm.set("/core.v1.AuthenticateService/", 2) + m.Expect(ok).To(m.BeFalse()) + + pmm.set("/greet.v1.GreetService/", 2) + val, exist := pmm.get("/greet.v1.GreetService/") + m.Expect(val).To(m.Equal(2)) + m.Expect(exist).To(m.BeTrue()) + + _, exist = pmm.get("/greet.v1.OtherService/") + m.Expect(exist).To(m.BeFalse()) + + val, exist = pmm.match("/connect/core.v1.AuthenticateService/login") + m.Expect(val).To(m.Equal(1)) + m.Expect(exist).To(m.BeTrue()) + + val, exist = pmm.match("/connect/core.v1.AuthenticateService/logout") + m.Expect(val).To(m.Equal(1)) + m.Expect(exist).To(m.BeTrue()) + }) +})