expose Access-Control-Allow-Private-Network

lqs · Nov 19, 2024 · 8027d0e · 8027d0e
1 parent 4e16a77
commit 8027d0e
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 18 deletions.
diff --git a/handler.go b/handler.go
@@ -12,11 +12,12 @@ import (
 )
 
 type mixHandler struct {
-	writer      http.ResponseWriter
-	request     *http.Request
-	grpcServer  *grpc.Server
-	grpcWeb     *grpcweb.WrappedGrpcServer
-	httpHandler http.Handler
+	writer              http.ResponseWriter
+	request             *http.Request
+	grpcServer          *grpc.Server
+	grpcWeb             *grpcweb.WrappedGrpcServer
+	httpHandler         http.Handler
+	allowPrivateNetwork bool
 }
 
 func (h mixHandler) isGrpc() bool {
@@ -67,6 +68,9 @@ func (h mixHandler) handleGrpcWeb() bool {
 	if (h.grpcWeb.IsGrpcWebRequest(h.request) || h.grpcWeb.IsAcceptableGrpcCorsRequest(h.request)) && h.isRegisteredGrpcPath() {
 		wrapper := wrapBrotli(h.writer, h.request)
 		defer wrapper.Close()
+		if h.allowPrivateNetwork {
+			wrapper.Header().Set("Access-Control-Allow-Private-Network", "true")
+		}
 		h.grpcWeb.ServeHTTP(wrapper, h.request)
 		return true
 	}
@@ -91,8 +95,8 @@ func (h mixHandler) handle() {
 	}
 }
 
-func newHandler(grpcServer *grpc.Server, http2Server *http2.Server, httpHandler http.Handler) (http.Handler, func()) {
-	grpcWeb := grpcweb.WrapServer(grpcServer, grpcweb.WithOriginFunc(func(origin string) bool {
+func (s *server) newHandler(http2Server *http2.Server, httpHandler http.Handler) (http.Handler, func()) {
+	grpcWeb := grpcweb.WrapServer(s.grpcServer, grpcweb.WithOriginFunc(func(origin string) bool {
 		return true
 	}))
 	var wg sync.WaitGroup
@@ -105,11 +109,12 @@ func newHandler(grpcServer *grpc.Server, http2Server *http2.Server, httpHandler
 		wg.Add(1)
 		defer wg.Done()
 		mixHandler{
-			writer:      writer,
-			request:     request,
-			grpcServer:  grpcServer,
-			grpcWeb:     grpcWeb,
-			httpHandler: httpHandler,
+			writer:              writer,
+			request:             request,
+			grpcServer:          s.grpcServer,
+			grpcWeb:             grpcWeb,
+			httpHandler:         httpHandler,
+			allowPrivateNetwork: s.config.AllowPrivateNetwork,
 		}.handle()
 	}), http2Server), wg.Wait
 }
diff --git a/server.go b/server.go
@@ -22,11 +22,12 @@ type Server interface {
 }
 
 type Config struct {
-	Port              int
-	ShutdownDelay     time.Duration
-	MaxHeaderBytes    int
-	GrpcServerOptions []grpc.ServerOption
-	OnStarted         func()
+	Port                int
+	ShutdownDelay       time.Duration
+	MaxHeaderBytes      int
+	GrpcServerOptions   []grpc.ServerOption
+	OnStarted           func()
+	AllowPrivateNetwork bool
 }
 
 type server struct {
@@ -58,7 +59,7 @@ func (s *server) StartAndWait(ctx context.Context) error {
 	var connectionClose atomic.Bool
 
 	http2Server := &http2.Server{}
-	handler, wait := newHandler(s.grpcServer, http2Server, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+	handler, wait := s.newHandler(http2Server, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if connectionClose.Load() {
 			w.Header().Set("Connection", "close")
 		}