37 lines
1.0 KiB
Go
37 lines
1.0 KiB
Go
package rpc
|
|
|
|
import (
|
|
"net/http"
|
|
|
|
"connectrpc.com/grpcreflect"
|
|
)
|
|
|
|
// ServeMux wraps http.ServeMux to collect registered Connect-RPC service names for automatic reflection.
|
|
type ServeMux struct {
|
|
*http.ServeMux
|
|
services []string
|
|
}
|
|
|
|
// NewServeMux creates a new initialized ServeMux.
|
|
func NewServeMux() *ServeMux {
|
|
return &ServeMux{
|
|
ServeMux: http.NewServeMux(),
|
|
}
|
|
}
|
|
|
|
// RegisterConnect registers a Connect-RPC service handler and tracks the service name for reflection.
|
|
func (sm *ServeMux) RegisterConnect(serviceName string, path string, handler http.Handler) {
|
|
sm.ServeMux.Handle(path, handler)
|
|
sm.services = append(sm.services, serviceName)
|
|
}
|
|
|
|
// MountReflection registers the reflection handlers for all registered Connect services on the ServeMux.
|
|
func (sm *ServeMux) MountReflection() {
|
|
if len(sm.services) == 0 {
|
|
return
|
|
}
|
|
reflector := grpcreflect.NewStaticReflector(sm.services...)
|
|
sm.ServeMux.Handle(grpcreflect.NewHandlerV1(reflector))
|
|
sm.ServeMux.Handle(grpcreflect.NewHandlerV1Alpha(reflector))
|
|
}
|