add custom connect mux adapt to Connect support

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 {
+func defaultConnectServer(addr string) (s *connectServer) {
-	return &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
 }

pkg/http/http.go

@@ -0,0 +1,5 @@
+// Copyright 2023 Michael Li <alimy@gility.net>. 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

pkg/http/http_suite_test.go

@@ -0,0 +1,17 @@
+// Copyright 2023 Michael Li <alimy@gility.net>. 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")
+}

pkg/http/mux.go

@@ -0,0 +1,130 @@
+// Copyright 2023 Michael Li <alimy@gility.net>. 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}
+}

pkg/http/mux_test.go

@@ -0,0 +1,71 @@
+// Copyright 2023 Michael Li <alimy@gility.net>. 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())
+	})
+})