From 4f2c902f4b856dbf59c6aee80a5f67810ed15079 Mon Sep 17 00:00:00 2001
From: Doria Keung <doriable@users.noreply.github.com>
Date: Mon, 27 Feb 2023 15:26:56 -0500
Subject: [PATCH] Handle private network preflight requests in Studio agent
 with flag `--private-network` (#1863)

This PR adds a `--private-network` boolean flag that allows the Studio
agent to handle
preflight requests from Studio in private networks that set the
`Access-Control-Request-Private-Network`
header set.
We still want to preserve the default behaviour of rejecting preflight
requests if the flag has not been set.
---
 CHANGELOG.md                                   |  2 ++
 .../command/beta/studioagent/studioagent.go    |  9 +++++++++
 .../bufpkg/bufstudioagent/bufstudioagent.go    | 18 +++++++++++++++---
 .../bufstudioagent/bufstudioagent_test.go      |  2 ++
 4 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3c2ce0f2bf..c5ddfa559a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,8 @@
 - Update built-in Well-Known Types to Protobuf v22.0.
 - Fixes a bug in `buf format` where C-style block comments in which every
   line includes a prefix (usually "*") would be incorrectly indented.
+- Add `--private-network` flag to `buf beta studio-agent` to support handling CORS requests
+  from Studio on private networks that set the `Access-Control-Request-Private-Network` header.
 
 ## [v1.14.0] - 2023-02-09
 
diff --git a/private/buf/cmd/buf/command/beta/studioagent/studioagent.go b/private/buf/cmd/buf/command/beta/studioagent/studioagent.go
index cdc435c881..a10b9dd4d9 100644
--- a/private/buf/cmd/buf/command/beta/studioagent/studioagent.go
+++ b/private/buf/cmd/buf/command/beta/studioagent/studioagent.go
@@ -42,6 +42,7 @@ const (
 	clientKeyFlagName         = "client-key"
 	serverCertFlagName        = "server-cert"
 	serverKeyFlagName         = "server-key"
+	privateNetworkFlagName    = "private-network"
 )
 
 // NewCommand returns a new Command.
@@ -75,6 +76,7 @@ type flags struct {
 	ClientKey         string
 	ServerCert        string
 	ServerKey         string
+	PrivateNetwork    bool
 }
 
 func newFlags() *flags {
@@ -142,6 +144,12 @@ func (f *flags) Bind(flagSet *pflag.FlagSet) {
 		"",
 		"The key to be used in the server TLS configuration",
 	)
+	flagSet.BoolVar(
+		&f.PrivateNetwork,
+		privateNetworkFlagName,
+		false,
+		`Use the agent with private network CORS`,
+	)
 }
 
 func run(
@@ -180,6 +188,7 @@ func run(
 		clientTLSConfig,
 		stringutil.SliceToMap(flags.DisallowedHeaders),
 		flags.ForwardHeaders,
+		flags.PrivateNetwork,
 	)
 	var httpListenConfig net.ListenConfig
 	httpListener, err := httpListenConfig.Listen(ctx, "tcp", fmt.Sprintf("%s:%s", flags.BindAddress, flags.Port))
diff --git a/private/bufpkg/bufstudioagent/bufstudioagent.go b/private/bufpkg/bufstudioagent/bufstudioagent.go
index 54a4661803..54a1d11fce 100644
--- a/private/bufpkg/bufstudioagent/bufstudioagent.go
+++ b/private/bufpkg/bufstudioagent/bufstudioagent.go
@@ -30,12 +30,17 @@ func NewHandler(
 	tlsClientConfig *tls.Config,
 	disallowedHeaders map[string]struct{},
 	forwardHeaders map[string]string,
+	privateNetwork bool,
 ) http.Handler {
-	corsHandler := cors.New(cors.Options{
+	corsHandlerOptions := cors.Options{
 		AllowedOrigins:   []string{origin},
-		AllowedMethods:   []string{http.MethodPost},
+		AllowedMethods:   []string{http.MethodPost, http.MethodOptions},
 		AllowCredentials: true,
-	})
+	}
+	if privateNetwork {
+		corsHandlerOptions.AllowPrivateNetwork = true
+	}
+	corsHandler := cors.New(corsHandlerOptions)
 	plainHandler := corsHandler.Handler(newPlainPostHandler(logger, disallowedHeaders, forwardHeaders, tlsClientConfig))
 	mux := http.NewServeMux()
 	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
@@ -46,6 +51,13 @@ func NewHandler(
 		case http.MethodPost:
 			plainHandler.ServeHTTP(w, r)
 			return
+		case http.MethodOptions:
+			if privateNetwork {
+				corsHandler.HandlerFunc(w, r)
+				return
+			}
+			// If the private network flag is not used, fall through to the default
+			fallthrough
 		default:
 			http.Error(w, "", http.StatusMethodNotAllowed)
 			return
diff --git a/private/bufpkg/bufstudioagent/bufstudioagent_test.go b/private/bufpkg/bufstudioagent/bufstudioagent_test.go
index bb2e60d5c0..91a7e5e779 100644
--- a/private/bufpkg/bufstudioagent/bufstudioagent_test.go
+++ b/private/bufpkg/bufstudioagent/bufstudioagent_test.go
@@ -65,6 +65,7 @@ func testPlainPostHandler(t *testing.T, upstreamServer *httptest.Server) {
 			upstreamServer.TLS,
 			nil,
 			map[string]string{"foo": "bar"},
+			false,
 		),
 	)
 	defer agentServer.Close()
@@ -142,6 +143,7 @@ func testPlainPostHandlerErrors(t *testing.T, upstreamServer *httptest.Server) {
 			upstreamServer.TLS,
 			map[string]struct{}{"forbidden-header": {}},
 			nil,
+			false,
 		),
 	)
 	defer agentServer.Close()