fix rocket to not panic on close when no messages have been sent

Summary:
fix rocket to not panic on close when no messages have been sent
This was initially found in D58828503

"""
we discovered a bug affecting Netstate (with Rocket) on NetOS hosts.

(But the bug is not limited to NetOS hosts - it just happens to manifests itself on NetOS hosts, because they differ from regular hosts.)

The bug causes a panic. It occurs when a client is opened and then immediately closed without making any RPC calls.
The protocol field remains nil and we get a nil dereference upon Close().
""" - echistyakov

Reviewed By: podtserkovskiy

Differential Revision: D58864210

fbshipit-source-id: 3533d808e5b641c802f6eb9acb5c6478983a346a

thrift/lib/go/thrift/rocket_upgrade_protocol.go

@@ -148,3 +148,13 @@ func (p *upgradeToRocketProtocol) SetProtocolID(protoID ProtocolID) error {
 	}
 	return p.Protocol.SetProtocolID(protoID)
 }
+
+func (p *upgradeToRocketProtocol) Close() error {
+	if p.Protocol == nil {
+		if err := p.headerProtocol.Close(); err != nil {
+			return err
+		}
+		return p.rocketProtocol.Close()
+	}
+	return p.Protocol.Close()
+}

thrift/lib/go/thrift/rocket_upgrade_protocol_test.go

@@ -0,0 +1,42 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package thrift
+
+import (
+	"testing"
+)
+
+func TestCloseWithoutSendingMessages(t *testing.T) {
+	serverSocket, err := NewServerSocket("[::]:0")
+	if err != nil {
+		t.Fatalf("could not create server socket: %s", err)
+	}
+	addr := serverSocket.Addr()
+	conn, err := DialHostPort(addr.String())
+	if err != nil {
+		t.Fatalf("could not create client socket: %s", err)
+	}
+	clientSocket, err := NewSocket(SocketConn(conn))
+	if err != nil {
+		t.Fatalf("could not create client socket: %s", err)
+	}
+	proto, err := NewUpgradeToRocketProtocol(clientSocket)
+	if err != nil {
+		t.Fatalf("could not create client protocol: %s", err)
+	}
+	proto.Close()
+}