close body from primay thread to avoid crazy errors

sudorandom · Aug 21, 2024 · 6d724b4 · 6d724b4
1 parent b5cc7da
commit 6d724b4
Show file tree

Hide file tree

Showing 5 changed files with 65 additions and 56 deletions.
diff --git a/cmd/fauxrpc/cmd_run.go b/cmd/fauxrpc/cmd_run.go
@@ -23,8 +23,10 @@ import (
 )
 
 type RunCmd struct {
-	Schema []string `help:"The modules to use for the RPC schema. It can be protobuf descriptors (binpb, json, yaml), a URL for reflection or a directory of descriptors."`
-	Addr   string   `short:"a" help:"Address to bind to." default:"127.0.0.1:6660"`
+	Schema       []string `help:"The modules to use for the RPC schema. It can be protobuf descriptors (binpb, json, yaml), a URL for reflection or a directory of descriptors."`
+	Addr         string   `short:"a" help:"Address to bind to." default:"127.0.0.1:6660"`
+	NoReflection bool     `help:"Disables the server reflection service."`
+	NoDocPage    bool     `help:"Disables the documentation page."`
 }
 
 func (c *RunCmd) Run(globals *Globals) error {
@@ -39,9 +41,6 @@ func (c *RunCmd) Run(globals *Globals) error {
 	if registry.ServiceCount() == 0 {
 		return errors.New("no services found in the given schemas")
 	}
-
-	// TODO: Add --no-reflection option
-	// TODO: Add --no-openapi option
 	// TODO: Load descriptors from stdin (assume protocol descriptors in binary format)
 	// TODO: add a stub service for registering stubs
 
@@ -60,20 +59,25 @@ func (c *RunCmd) Run(globals *Globals) error {
 	if err != nil {
 		log.Fatalf("err: %s", err)
 	}
-	reflector := grpcreflect.NewReflector(&staticNames{names: serviceNames}, grpcreflect.WithDescriptorResolver(registry.Files()))
 
 	mux := http.NewServeMux()
 	mux.Handle("/", transcoder)
-	mux.Handle(grpcreflect.NewHandlerV1(reflector))
-	mux.Handle(grpcreflect.NewHandlerV1Alpha(reflector))
+
+	if !c.NoReflection {
+		reflector := grpcreflect.NewReflector(&staticNames{names: serviceNames}, grpcreflect.WithDescriptorResolver(registry.Files()))
+		mux.Handle(grpcreflect.NewHandlerV1(reflector))
+		mux.Handle(grpcreflect.NewHandlerV1Alpha(reflector))
+	}
 
 	// OpenAPI Stuff
-	resp, err := convertToOpenAPISpec(registry)
-	if err != nil {
-		return err
+	if !c.NoDocPage {
+		resp, err := convertToOpenAPISpec(registry)
+		if err != nil {
+			return err
+		}
+		mux.Handle("GET /fauxrpc.openapi.html", singleFileHandler(openapiHTML))
+		mux.Handle("GET /fauxrpc.openapi.yaml", singleFileHandler(resp.File[0].GetContent()))
 	}
-	mux.Handle("GET /fauxrpc.openapi.html", singleFileHandler(openapiHTML))
-	mux.Handle("GET /fauxrpc.openapi.yaml", singleFileHandler(resp.File[0].GetContent()))
 
 	server := &http.Server{
 		Addr:    c.Addr,
@@ -82,10 +86,14 @@ func (c *RunCmd) Run(globals *Globals) error {
 
 	fmt.Printf("FauxRPC (%s)\n", fullVersion())
 	fmt.Printf("Listening on http://%s\n", c.Addr)
-	fmt.Printf("OpenAPI documentation: http://%s/fauxrpc.openapi.html\n", c.Addr)
+	if !c.NoDocPage {
+		fmt.Printf("OpenAPI documentation: http://%s/fauxrpc.openapi.html\n", c.Addr)
+	}
 	fmt.Println()
 	fmt.Println("Example Commands:")
-	fmt.Printf("$ buf curl --http2-prior-knowledge http://%s --list-methods\n", c.Addr)
+	if !c.NoReflection {
+		fmt.Printf("$ buf curl --http2-prior-knowledge http://%s --list-methods\n", c.Addr)
+	}
 	fmt.Printf("$ buf curl --http2-prior-knowledge http://%s/[METHOD_NAME]\n", c.Addr)
 	return server.ListenAndServe()
 }
@@ -119,7 +127,7 @@ func convertToOpenAPISpec(registry *protobuf.ServiceRegistry) (*pluginpb.CodeGen
 	base := openAPIBase{
 		OpenAPI: "3.1.0",
 		Info: openAPIBaseInfo{
-			Title:       "FauxRPC Generated Documentation",
+			Title:       "FauxRPC Documentation",
 			Description: descBuilder.String(),
 			Version:     strings.TrimPrefix(version, "v"),
 		},
@@ -139,39 +147,7 @@ func convertToOpenAPISpec(registry *protobuf.ServiceRegistry) (*pluginpb.CodeGen
 
 func singleFileHandler(content string) http.HandlerFunc {
 	return func(w http.ResponseWriter, r *http.Request) {
+
 		fmt.Fprint(w, content)
 	}
 }
-
-const openapiHTML = `
-<!doctype html>
-<html lang="en">
-  <head>
-    <meta charset="utf-8">
-    <meta name="viewport" content="width=device-width, initial-scale=1, shrink-to-fit=no">
-    <title>Elements in HTML</title>
-    <!-- Embed elements Elements via Web Component -->
-    <script src="https://unpkg.com/@stoplight/elements/web-components.min.js"></script>
-    <link rel="stylesheet" href="https://unpkg.com/@stoplight/elements/styles.min.css">
-  </head>
-  <body>
-
-    <elements-api
-      apiDescriptionUrl="/fauxrpc.openapi.yaml"
-      router="hash"
-      layout="sidebar"
-    />
-
-  </body>
-</html>
-`
-
-type openAPIBaseInfo struct {
-	Description string `yaml:"description"`
-	Title       string `yaml:"title"`
-	Version     string `yaml:"version"`
-}
-type openAPIBase struct {
-	OpenAPI string          `yaml:"openapi"`
-	Info    openAPIBaseInfo `yaml:"info"`
-}
diff --git a/cmd/fauxrpc/openapi.go b/cmd/fauxrpc/openapi.go
@@ -0,0 +1,33 @@
+package main
+
+const openapiHTML = `
+<!doctype html>
+<html lang="en">
+  <head>
+    <meta charset="utf-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1, shrink-to-fit=no">
+    <title>FauxRPC Documentation</title>
+    <script src="https://unpkg.com/@stoplight/elements@8.3.4/web-components.min.js"></script>
+	<link rel="stylesheet" href="https://unpkg.com/@stoplight/elements@8.3.4/styles.min.css">
+  </head>
+  <body>
+
+    <elements-api
+      apiDescriptionUrl="/fauxrpc.openapi.yaml"
+      router="hash"
+      layout="sidebar"
+    />
+
+  </body>
+</html>
+`
+
+type openAPIBaseInfo struct {
+	Description string `yaml:"description"`
+	Title       string `yaml:"title"`
+	Version     string `yaml:"version"`
+}
+type openAPIBase struct {
+	OpenAPI string          `yaml:"openapi"`
+	Info    openAPIBaseInfo `yaml:"info"`
+}
diff --git a/go.mod b/go.mod
@@ -17,7 +17,7 @@ require (
 	github.com/bufbuild/protoyaml-go v0.1.11
 	github.com/golang/protobuf v1.5.4
 	github.com/stretchr/testify v1.9.0
-	github.com/sudorandom/protoc-gen-connect-openapi v0.8.10
+	github.com/sudorandom/protoc-gen-connect-openapi v0.8.11
 	golang.org/x/net v0.28.0
 	google.golang.org/protobuf v1.34.2
 	gopkg.in/yaml.v3 v3.0.1
@@ -41,6 +41,6 @@ require (
 	github.com/wk8/go-ordered-map/v2 v2.1.8 // indirect
 	golang.org/x/exp v0.0.0-20240808152545-0cdaa3abc0fa // indirect
 	golang.org/x/text v0.17.0 // indirect
-	google.golang.org/genproto/googleapis/api v0.0.0-20240814211410-ddb44dafa142 // indirect
+	google.golang.org/genproto/googleapis/api v0.0.0-20240820151423-278611b39280 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20240814211410-ddb44dafa142 // indirect
 )
diff --git a/go.sum b/go.sum
@@ -924,8 +924,8 @@ github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO
 github.com/stretchr/testify v1.8.1/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
 github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
 github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
-github.com/sudorandom/protoc-gen-connect-openapi v0.8.10 h1:DLhsphX20NjFZoLztqvgWYpRVvNaPngQ4lebnsgMS/E=
-github.com/sudorandom/protoc-gen-connect-openapi v0.8.10/go.mod h1:J0UbccOfevPsNTErI+htxk4CwvZOOzEKzNBafHCx2YI=
+github.com/sudorandom/protoc-gen-connect-openapi v0.8.11 h1:nUEG2qIEyye/QQwyc9yVcdE1Q4oeW8iAa2xrMkvMeRI=
+github.com/sudorandom/protoc-gen-connect-openapi v0.8.11/go.mod h1:J0UbccOfevPsNTErI+htxk4CwvZOOzEKzNBafHCx2YI=
 github.com/vmware-labs/yaml-jsonpath v0.3.2 h1:/5QKeCBGdsInyDCyVNLbXyilb61MXGi9NP674f9Hobk=
 github.com/vmware-labs/yaml-jsonpath v0.3.2/go.mod h1:U6whw1z03QyqgWdgXxvVnQ90zN1BWz5V+51Ewf8k+rQ=
 github.com/wk8/go-ordered-map/v2 v2.1.8 h1:5h/BUHu93oj4gIdvHHHGsScSTMijfx5PeYkE/fJgbpc=
@@ -1511,8 +1511,8 @@ google.golang.org/genproto v0.0.0-20230525234025-438c736192d0/go.mod h1:9ExIQyXL
 google.golang.org/genproto v0.0.0-20230526161137-0005af68ea54/go.mod h1:zqTuNwFlFRsw5zIts5VnzLQxSRqh+CGOTVMlYbY0Eyk=
 google.golang.org/genproto/googleapis/api v0.0.0-20230525234020-1aefcd67740a/go.mod h1:ts19tUU+Z0ZShN1y3aPyq2+O3d5FUNNgT6FtOzmrNn8=
 google.golang.org/genproto/googleapis/api v0.0.0-20230525234035-dd9d682886f9/go.mod h1:vHYtlOoi6TsQ3Uk2yxR7NI5z8uoV+3pZtR4jmHIkRig=
-google.golang.org/genproto/googleapis/api v0.0.0-20240814211410-ddb44dafa142 h1:wKguEg1hsxI2/L3hUYrpo1RVi48K+uTyzKqprwLXsb8=
-google.golang.org/genproto/googleapis/api v0.0.0-20240814211410-ddb44dafa142/go.mod h1:d6be+8HhtEtucleCbxpPW9PA9XwISACu8nvpPqF0BVo=
+google.golang.org/genproto/googleapis/api v0.0.0-20240820151423-278611b39280 h1:YDFM9oOjiFhaMAVgbDxfxW+66nRrsvzQzJ51wp3OxC0=
+google.golang.org/genproto/googleapis/api v0.0.0-20240820151423-278611b39280/go.mod h1:fO8wJzT2zbQbAjbIoos1285VfEIYKDDY+Dt+WpTkh6g=
 google.golang.org/genproto/googleapis/rpc v0.0.0-20230525234015-3fc162c6f38a/go.mod h1:xURIpW9ES5+/GZhnV6beoEtxQrnkRGIfP5VQG2tCBLc=
 google.golang.org/genproto/googleapis/rpc v0.0.0-20230525234030-28d5490b6b19/go.mod h1:66JfowdXAEgad5O9NnYcsNPLCPZJD++2L9X0PCMODrA=
 google.golang.org/genproto/googleapis/rpc v0.0.0-20240814211410-ddb44dafa142 h1:e7S5W7MGGLaSu8j3YjdezkZ+m1/Nm0uRVRMEMGk26Xs=

diff --git a/private/protobuf/handler.go b/private/protobuf/handler.go
@@ -38,11 +38,11 @@ func NewHandler(service protoreflect.ServiceDescriptor) http.Handler {
 			w.Header().Set("Grpc-Message", "method not found")
 			return
 		}
+		defer r.Body.Close()
 
 		// completely ignore the body. Maybe later we'll need it as input to the response message
 		go func() {
 			_, _ = io.Copy(io.Discard, r.Body)
-			_ = r.Body.Close()
 		}()
 
 		slog.Info("MethodCalled", slog.String("service", serviceName), slog.String("method", methodName))