Repository: nats-io/nats.go
Pull request #1433 (2023-10-06 20:43:19)

[ADDED] TLSHandshakeFirst option
This allows to connect to a server that is also configured to perform the TLS handshake first, that is, before sending the INFO protocol.

Signed-off-by: Ivan Kozlovic <ivan@synadia.com>


Patch:
diff --git a/jetstream/jetstream.go b/jetstream/jetstream.go
index 576e150..927ac46 100644
--- a/jetstream/jetstream.go
+++ b/jetstream/jetstream.go
@@ -69,6 +69,8 @@ type (
 		CreateStream(ctx context.Context, cfg StreamConfig) (Stream, error)
 		// UpdateStream updates an existing stream
 		UpdateStream(ctx context.Context, cfg StreamConfig) (Stream, error)
+		// CreateOrUpdateStream creates a stream with given config. If stream already exists, it will be updated (if possible).
+		CreateOrUpdateStream(ctx context.Context, cfg StreamConfig) (Stream, error)
 		// Stream returns a [Stream] hook for a given stream name
 		Stream(ctx context.Context, stream string) (Stream, error)
 		// StreamNameBySubject returns a stream name stream listening on given subject
@@ -460,6 +462,18 @@ func (js *jetStream) UpdateStream(ctx context.Context, cfg StreamConfig) (Stream
 	}, nil
 }
 
+func (js *jetStream) CreateOrUpdateStream(ctx context.Context, cfg StreamConfig) (Stream, error) {
+	s, err := js.UpdateStream(ctx, cfg)
+	if err != nil {
+		if !errors.Is(err, ErrStreamNotFound) {
+			return nil, err
+		}
+		return js.CreateStream(ctx, cfg)
+	}
+
+	return s, nil
+}
+
 // Stream returns a [Stream] hook for a given stream name
 func (js *jetStream) Stream(ctx context.Context, name string) (Stream, error) {
 	if err := validateStreamName(name); err != nil {
diff --git a/nats.go b/nats.go
index 82b7973..73709db 100644
--- a/nats.go
+++ b/nats.go
@@ -311,6 +311,13 @@ type Options struct {
 	// TLSCertCB is used to fetch and return custom tls certificate.
 	TLSCertCB TLSCertHandler
 
+	// TLSHandshakeFirst is used to instruct the library perform
+	// the TLS handshake right after the connect and before receiving
+	// the INFO protocol from the server. If this option is enabled
+	// but the server is not configured to perform the TLS handshake
+	// first, the connection will fail.
+	TLSHandshakeFirst bool
+
 	// RootCAsCB is used to fetch and return a set of root certificate
 	// authorities that clients use when verifying server certificates.
 	RootCAsCB RootCAsHandler
@@ -1315,6 +1322,17 @@ func SkipHostLookup() Option {
 	}
 }
 
+// TLSHandshakeFirst is an Option to perform the TLS handshake first, that is
+// before receiving the INFO protocol. This requires the server to also be
+// configured with such option, otherwise the connection will fail.
+func TLSHandshakeFirst() Option {
+	return func(o *Options) error {
+		o.TLSHandshakeFirst = true
+		o.Secure = true
+		return nil
+	}
+}
+
 // Handler processing
 
 // SetDisconnectHandler will set the disconnect event handler.
@@ -1481,6 +1499,12 @@ func (o Options) Connect() (*Conn, error) {
 		}
 	}
 
+	// If the TLSHandshakeFirst option is specified, make sure that
+	// the Secure boolean is true.
+	if nc.Opts.TLSHandshakeFirst {
+		nc.Opts.Secure = true
+	}
+
 	if err := nc.setupServerPool(); err != nil {
 		return nil, err
 	}
@@ -2235,6 +2259,14 @@ func (nc *Conn) processConnectInit() error {
 	// Set our status to connecting.
 	nc.changeConnStatus(CONNECTING)
 
+	// If we need to have a TLS connection and want the TLS handshake to occur
+	// first, do it now.
+	if nc.Opts.Secure && nc.Opts.TLSHandshakeFirst {
+		if err := nc.makeTLSConn(); err != nil {
+			return err
+		}
+	}
+
 	// Process the INFO protocol received from the server
 	err := nc.processExpectedInfo()
 	if err != nil {
@@ -2351,8 +2383,13 @@ func (nc *Conn) checkForSecure() error {
 		o.Secure = true
 	}
 
-	// Need to rewrap with bufio
 	if o.Secure {
+		// If TLS handshake first is true, we have already done
+		// the handshake, so we are done here.
+		if o.TLSHandshakeFirst {
+			return nil
+		}
+		// Need to rewrap with bufio
 		if err := nc.makeTLSConn(); err != nil {
 			return err
 		}


Test patch:
diff --git a/jetstream/test/jetstream_test.go b/jetstream/test/jetstream_test.go
index 3168bfa..410c7f7 100644
--- a/jetstream/test/jetstream_test.go
+++ b/jetstream/test/jetstream_test.go
@@ -455,6 +455,136 @@ func TestCreateStreamMirrorCrossDomains(t *testing.T) {
 	}
 }
 
+func TestCreateOrUpdateStream(t *testing.T) {
+	tests := []struct {
+		name          string
+		stream        string
+		subject       string
+		timeout       time.Duration
+		withError     error
+		withInfoCheck bool
+	}{
+		{
+			name:          "create stream ok",
+			stream:        "foo",
+			timeout:       10 * time.Second,
+			subject:       "FOO.1",
+			withInfoCheck: false,
+		},
+		{
+			name:          "create stream empty context",
+			stream:        "foo-o",
+			subject:       "FOO.12",
+			withInfoCheck: false,
+		},
+		{
+			name:          "create stream invalid stream name",
+			stream:        "foo.123",
+			subject:       "FOO-123",
+			timeout:       10 * time.Second,
+			withError:     jetstream.ErrInvalidStreamName,
+			withInfoCheck: false,
+		},
+		{
+			name:          "create stream stream name required",
+			stream:        "",
+			subject:       "FOO-1234",
+			timeout:       10 * time.Second,
+			withError:     jetstream.ErrStreamNameRequired,
+			withInfoCheck: false,
+		},
+		{
+			name:          "update stream ok",
+			stream:        "foo",
+			subject:       "BAR-123",
+			timeout:       10 * time.Second,
+			withInfoCheck: true,
+		},
+		{
+			name:          "create stream context timeout",
+			stream:        "foo",
+			subject:       "BAR-1234",
+			timeout:       1 * time.Microsecond,
+			withError:     context.DeadlineExceeded,
+			withInfoCheck: false,
+		},
+		{
+			name:          "update stream with empty context",
+			stream:        "sample-foo-1",
+			subject:       "SAMPLE-FOO-123",
+			withInfoCheck: true,
+		},
+		{
+			name:          "update stream invalid stream name",
+			stream:        "sample-foo.123",
+			subject:       "SAMPLE-FOO-1234",
+			timeout:       10 * time.Second,
+			withError:     jetstream.ErrInvalidStreamName,
+			withInfoCheck: true,
+		},
+		{
+			name:          "update stream stream name required",
+			stream:        "",
+			subject:       "SAMPLE-FOO-123",
+			timeout:       10 * time.Second,
+			withError:     jetstream.ErrStreamNameRequired,
+			withInfoCheck: true,
+		},
+		{
+			name:          "update stream context timeout",
+			stream:        "sample-foo-2",
+			subject:       "SAMPLE-FOO-123456",
+			timeout:       1 * time.Microsecond,
+			withError:     context.DeadlineExceeded,
+			withInfoCheck: true,
+		},
+	}
+
+	srv := RunBasicJetStreamServer()
+	defer shutdownJSServerAndRemoveStorage(t, srv)
+	nc, err := nats.Connect(srv.ClientURL())
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+
+	js, err := jetstream.New(nc)
+	if err != nil {
+		t.Fatalf("Unexpected error: %v", err)
+	}
+	defer nc.Close()
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			ctx := context.Background()
+			if test.timeout > 0 {
+				var cancel context.CancelFunc
+				ctx, cancel = context.WithTimeout(context.Background(), test.timeout)
+				defer cancel()
+			}
+			s, err := js.CreateOrUpdateStream(ctx, jetstream.StreamConfig{Name: test.stream, Subjects: []string{test.subject}})
+			if test.withError != nil {
+				if !errors.Is(err, test.withError) {
+					t.Fatalf("Expected error: %v; got: %v", test.withError, err)
+				}
+				return
+			}
+			if err != nil {
+				t.Fatalf("Unexpected error: %v", err)
+			}
+
+			if test.withInfoCheck {
+				info, err := s.Info(ctx)
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+				if len(info.Config.Subjects) != 1 || info.Config.Subjects[0] != test.subject {
+					t.Fatalf("Invalid stream subjects after update: %v", info.Config.Subjects)
+				}
+			}
+		})
+	}
+}
+
 func TestUpdateStream(t *testing.T) {
 	tests := []struct {
 		name      string
diff --git a/test/conn_test.go b/test/conn_test.go
index 995b3d1..f4bdd87 100644
--- a/test/conn_test.go
+++ b/test/conn_test.go
@@ -2863,3 +2863,112 @@ func TestConnStatusChangedEvents(t *testing.T) {
 		time.Sleep(100 * time.Millisecond)
 	})
 }
+
+func TestTLSHandshakeFirst(t *testing.T) {
+	s, opts := RunServerWithConfig("./configs/tls.conf")
+	defer s.Shutdown()
+
+	secureURL := fmt.Sprintf("tls://derek:porkchop@localhost:%d", opts.Port)
+	nc, err := nats.Connect(secureURL,
+		nats.RootCAs("./configs/certs/ca.pem"),
+		nats.TLSHandshakeFirst())
+	if err == nil || !strings.Contains(err.Error(), "TLS handshake") {
+		if err == nil {
+			nc.Close()
+		}
+		t.Fatalf("Expected error about not being a TLS handshake, got %v", err)
+	}
+
+	tc := &server.TLSConfigOpts{
+		CertFile: "./configs/certs/server.pem",
+		KeyFile:  "./configs/certs/key.pem",
+	}
+	tlsConf, err := server.GenTLSConfig(tc)
+	if err != nil {
+		t.Fatalf("Can't build TLCConfig: %v", err)
+	}
+	tlsConf.ServerName = "localhost"
+
+	// Start a mockup server that will do the TLS handshake first
+	// and then send the INFO protcol.
+	l, e := net.Listen("tcp", ":0")
+	if e != nil {
+		t.Fatal("Could not listen on an ephemeral port")
+	}
+	tl := l.(*net.TCPListener)
+	defer tl.Close()
+
+	addr := tl.Addr().(*net.TCPAddr)
+
+	errCh := make(chan error, 1)
+	doneCh := make(chan struct{})
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		conn, err := l.Accept()
+		if err != nil {
+			errCh <- fmt.Errorf("error accepting client connection: %v", err)
+			return
+		}
+		defer conn.Close()
+
+		// Do the TLS handshake now.
+		conn = tls.Server(conn, tlsConf)
+		tlsconn := conn.(*tls.Conn)
+		if err := tlsconn.Handshake(); err != nil {
+			errCh <- fmt.Errorf("Server error during handshake: %v", err)
+			return
+		}
+
+		// Send back the INFO
+		info := fmt.Sprintf("INFO {\"server_id\":\"foobar\",\"host\":\"localhost\",\"port\":%d,\"auth_required\":false,\"tls_required\":true,\"tls_available\":true,\"tls_verify\":true,\"max_payload\":1048576}\r\n", addr.Port)
+		tlsconn.Write([]byte(info))
+
+		// Read connect and ping commands sent from the client
+		line := make([]byte, 256)
+		_, err = tlsconn.Read(line)
+		if err != nil {
+			errCh <- fmt.Errorf("expected CONNECT and PING from client, got: %s", err)
+			return
+		}
+		tlsconn.Write([]byte("PONG\r\n"))
+
+		// Wait for the signal that client is ok
+		<-doneCh
+		// Server is done now.
+		errCh <- nil
+	}()
+
+	time.Sleep(100 * time.Millisecond)
+
+	secureURL = fmt.Sprintf("tls://derek:porkchop@localhost:%d", addr.Port)
+	nc, err = nats.Connect(secureURL,
+		nats.RootCAs("./configs/certs/ca.pem"),
+		nats.TLSHandshakeFirst())
+	if err != nil {
+		wg.Wait()
+		e := <-errCh
+		t.Fatalf("Unexpected error: %v (server error=%s)", err, e.Error())
+	}
+
+	state, err := nc.TLSConnectionState()
+	if err != nil {
+		t.Fatalf("Expected connection state: %v", err)
+	}
+	if !state.HandshakeComplete {
+		t.Fatalf("Expected valid connection state")
+	}
+	nc.Close()
+
+	close(doneCh)
+	wg.Wait()
+	select {
+	case e := <-errCh:
+		if e != nil {
+			t.Fatalf("Error from server: %v", err)
+		}
+	case <-time.After(2 * time.Second):
+		t.Fatal("Server did not exit")
+	}
+}