Repository: nats-io/nats.go
Pull request #1428 (2023-10-03 17:54:40)

[FIXED] Invalid handling of heartbeats in Consume and Messages


Patch:
diff --git a/dependencies.md b/dependencies.md
index cc986b2..d54bcf9 100644
--- a/dependencies.md
+++ b/dependencies.md
@@ -5,9 +5,11 @@ This file lists the dependencies used in this repository.
 | Dependency | License |
 |-|-|
 | Go | BSD 3-Clause "New" or "Revised" License |
-| github.com/nats-io/nats.go | Apache License 2.0 |
 | github.com/golang/protobuf v1.4.2 | BSD 3-Clause "New" or "Revised" License |
-| github.com/nats-io/nats-server/v2 v2.1.8-0.20201115145023-f61fa8529a0f | Apache License 2.0 |
-| github.com/nats-io/nkeys v0.2.0 | Apache License 2.0 |
+| github.com/nats-io/nats-server/v2 v2.10.0 | Apache License 2.0 |
+| github.com/nats-io/nkeys v0.4.5 | Apache License 2.0 |
 | github.com/nats-io/nuid v1.0.1 | Apache License 2.0 |
 | google.golang.org/protobuf v1.23.0 | BSD 3-Clause License |
+| github.com/klauspost/compress v1.17.0 | Apache License 2.0 |
+| go.uber.org/goleak v1.2.1 | MIT License |
+| golang.org/x/text |  BSD 3-Clause License |
diff --git a/jetstream/consumer.go b/jetstream/consumer.go
index 1fb6fb5..324327f 100644
--- a/jetstream/consumer.go
+++ b/jetstream/consumer.go
@@ -77,6 +77,9 @@ func (p *pullConsumer) Info(ctx context.Context) (*ConsumerInfo, error) {
 		}
 		return nil, resp.Error
 	}
+	if resp.Error == nil && resp.ConsumerInfo == nil {
+		return nil, ErrConsumerNotFound
+	}
 
 	p.info = resp.ConsumerInfo
 	return resp.ConsumerInfo, nil
@@ -188,6 +191,9 @@ func getConsumer(ctx context.Context, js *jetStream, stream, name string) (Consu
 		}
 		return nil, resp.Error
 	}
+	if resp.Error == nil && resp.ConsumerInfo == nil {
+		return nil, ErrConsumerNotFound
+	}
 
 	cons := &pullConsumer{
 		jetStream:     js,
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/jetstream/options.go b/jetstream/options.go
index 566f0dc..b47268e 100644
--- a/jetstream/options.go
+++ b/jetstream/options.go
@@ -121,18 +121,20 @@ func (max PullMaxMessages) configureMessages(opts *consumeOpts) error {
 type PullExpiry time.Duration
 
 func (exp PullExpiry) configureConsume(opts *consumeOpts) error {
-	if exp < 0 {
-		return fmt.Errorf("%w: expires value must be positive", ErrInvalidOption)
+	expiry := time.Duration(exp)
+	if expiry < 1*time.Second {
+		return fmt.Errorf("%w: expires value must be at least 1s", ErrInvalidOption)
 	}
-	opts.Expires = time.Duration(exp)
+	opts.Expires = expiry
 	return nil
 }
 
 func (exp PullExpiry) configureMessages(opts *consumeOpts) error {
-	if exp < 0 {
-		return fmt.Errorf("%w: expires value must be positive", ErrInvalidOption)
+	expiry := time.Duration(exp)
+	if expiry < 0 {
+		return fmt.Errorf("%w: expires value must be at least 1s", ErrInvalidOption)
 	}
-	opts.Expires = time.Duration(exp)
+	opts.Expires = expiry
 	return nil
 }
 
@@ -191,8 +193,8 @@ type PullHeartbeat time.Duration
 
 func (hb PullHeartbeat) configureConsume(opts *consumeOpts) error {
 	hbTime := time.Duration(hb)
-	if hbTime < 1*time.Second || hbTime > 30*time.Second {
-		return fmt.Errorf("%w: idle_heartbeat value must be within 1s-30s range", ErrInvalidOption)
+	if hbTime < 500*time.Millisecond || hbTime > 30*time.Second {
+		return fmt.Errorf("%w: idle_heartbeat value must be within 500ms-30s range", ErrInvalidOption)
 	}
 	opts.Heartbeat = hbTime
 	return nil
@@ -200,8 +202,8 @@ func (hb PullHeartbeat) configureConsume(opts *consumeOpts) error {
 
 func (hb PullHeartbeat) configureMessages(opts *consumeOpts) error {
 	hbTime := time.Duration(hb)
-	if hbTime < 1*time.Second || hbTime > 30*time.Second {
-		return fmt.Errorf("%w: idle_heartbeat value must be within 1s-30s range", ErrInvalidOption)
+	if hbTime < 500*time.Millisecond || hbTime > 30*time.Second {
+		return fmt.Errorf("%w: idle_heartbeat value must be within 500ms-30s range", ErrInvalidOption)
 	}
 	opts.Heartbeat = hbTime
 	return nil
diff --git a/jetstream/ordered.go b/jetstream/ordered.go
index 296ea6b..e462981 100644
--- a/jetstream/ordered.go
+++ b/jetstream/ordered.go
@@ -519,6 +519,9 @@ func (c *orderedConsumer) Info(ctx context.Context) (*ConsumerInfo, error) {
 		}
 		return nil, resp.Error
 	}
+	if resp.Error == nil && resp.ConsumerInfo == nil {
+		return nil, ErrConsumerNotFound
+	}
 
 	c.currentConsumer.info = resp.ConsumerInfo
 	return resp.ConsumerInfo, nil
diff --git a/jetstream/pull.go b/jetstream/pull.go
index 7029bb0..a3601ec 100644
--- a/jetstream/pull.go
+++ b/jetstream/pull.go
@@ -132,7 +132,6 @@ type (
 const (
 	DefaultMaxMessages = 500
 	DefaultExpires     = 30 * time.Second
-	DefaultHeartbeat   = 5 * time.Second
 	unset              = -1
 )
 
@@ -192,12 +191,17 @@ func (p *pullConsumer) Consume(handler MessageHandler, opts ...PullConsumeOpt) (
 		}
 		userMsg, msgErr := checkMsg(msg)
 		if !userMsg && msgErr == nil {
+			if sub.hbMonitor != nil {
+				sub.hbMonitor.Reset(2 * consumeOpts.Heartbeat)
+			}
 			return
 		}
 		defer func() {
 			sub.Lock()
 			sub.checkPending()
-			sub.hbMonitor.Reset(2 * consumeOpts.Heartbeat)
+			if sub.hbMonitor != nil {
+				sub.hbMonitor.Reset(2 * consumeOpts.Heartbeat)
+			}
 			sub.Unlock()
 		}()
 		if !userMsg {
@@ -305,6 +309,9 @@ func (p *pullConsumer) Consume(handler MessageHandler, opts ...PullConsumeOpt) (
 							MaxBytes:  sub.consumeOpts.MaxBytes,
 							Heartbeat: sub.consumeOpts.Heartbeat,
 						}
+						if sub.hbMonitor != nil {
+							sub.hbMonitor.Reset(2 * sub.consumeOpts.Heartbeat)
+						}
 						sub.resetPendingMsgs()
 					}
 					sub.Unlock()
@@ -325,6 +332,9 @@ func (p *pullConsumer) Consume(handler MessageHandler, opts ...PullConsumeOpt) (
 						MaxBytes:  sub.consumeOpts.MaxBytes,
 						Heartbeat: sub.consumeOpts.Heartbeat,
 					}
+					if sub.hbMonitor != nil {
+						sub.hbMonitor.Reset(2 * sub.consumeOpts.Heartbeat)
+					}
 					sub.resetPendingMsgs()
 				}
 				sub.Unlock()
@@ -465,7 +475,7 @@ func (s *pullSubscription) Next() (Msg, error) {
 	if atomic.LoadUint32(&s.closed) == 1 {
 		return nil, ErrMsgIteratorClosed
 	}
-	hbMonitor := s.scheduleHeartbeatCheck(s.consumeOpts.Heartbeat)
+	hbMonitor := s.scheduleHeartbeatCheck(2 * s.consumeOpts.Heartbeat)
 	defer func() {
 		if hbMonitor != nil {
 			hbMonitor.Stop()
@@ -509,6 +519,9 @@ func (s *pullSubscription) Next() (Msg, error) {
 				if s.consumeOpts.ReportMissingHeartbeats {
 					return nil, err
 				}
+				if hbMonitor != nil {
+					hbMonitor.Reset(2 * s.consumeOpts.Heartbeat)
+				}
 			}
 			if errors.Is(err, errConnected) {
 				if !isConnected {
@@ -531,12 +544,14 @@ func (s *pullSubscription) Next() (Msg, error) {
 					}
 					s.pending.msgCount = 0
 					s.pending.byteCount = 0
-					hbMonitor = s.scheduleHeartbeatCheck(s.consumeOpts.Heartbeat)
+					if hbMonitor != nil {
+						hbMonitor.Reset(2 * s.consumeOpts.Heartbeat)
+					}
 				}
 			}
 			if errors.Is(err, errDisconnected) {
 				if hbMonitor != nil {
-					hbMonitor.Stop()
+					hbMonitor.Reset(2 * s.consumeOpts.Heartbeat)
 				}
 				isConnected = false
 			}
diff --git a/js.go b/js.go
index 7fdb013..444278e 100644
--- a/js.go
+++ b/js.go
@@ -3223,6 +3223,9 @@ func (js *js) getConsumerInfoContext(ctx context.Context, stream, consumer strin
 		}
 		return nil, info.Error
 	}
+	if info.Error == nil && info.ConsumerInfo == nil {
+		return nil, ErrConsumerNotFound
+	}
 	return info.ConsumerInfo, 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/helper_test.go b/jetstream/test/helper_test.go
index ce7909d..a9dbae2 100644
--- a/jetstream/test/helper_test.go
+++ b/jetstream/test/helper_test.go
@@ -39,6 +39,14 @@ type jsServer struct {
 	restart sync.Mutex
 }
 
+// Restart can be used to start again a server
+// using the same listen address as before.
+func (srv *jsServer) Restart() {
+	srv.restart.Lock()
+	defer srv.restart.Unlock()
+	srv.Server = natsserver.RunServer(srv.myopts)
+}
+
 // Dumb wait program to sync on callbacks, etc... Will timeout
 func Wait(ch chan bool) error {
 	return WaitTime(ch, 5*time.Second)
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/jetstream/test/pull_test.go b/jetstream/test/pull_test.go
index 981a44f..2035412 100644
--- a/jetstream/test/pull_test.go
+++ b/jetstream/test/pull_test.go
@@ -2045,141 +2045,220 @@ func TestPullConsumerConsume(t *testing.T) {
 func TestPullConsumerConsume_WithCluster(t *testing.T) {
 	testSubject := "FOO.123"
 	testMsgs := []string{"m1", "m2", "m3", "m4", "m5"}
-	publishTestMsgs := func(t *testing.T, nc *nats.Conn) {
+	publishTestMsgs := func(t *testing.T, js jetstream.JetStream) {
 		for _, msg := range testMsgs {
-			if err := nc.Publish(testSubject, []byte(msg)); err != nil {
+			if _, err := js.Publish(context.Background(), testSubject, []byte(msg)); err != nil {
 				t.Fatalf("Unexpected error during publish: %s", err)
 			}
 		}
 	}
 
 	name := "cluster"
-	stream := jetstream.StreamConfig{
+	singleStream := jetstream.StreamConfig{
 		Name:     name,
 		Replicas: 1,
 		Subjects: []string{"FOO.*"},
 	}
 
-	t.Run("no options", func(t *testing.T) {
-		withJSClusterAndStream(t, name, 3, stream, func(t *testing.T, subject string, srvs ...*jsServer) {
-			nc, err := nats.Connect(srvs[0].ClientURL())
-			if err != nil {
-				t.Fatalf("Unexpected error: %v", err)
-			}
+	streamWithReplicas := jetstream.StreamConfig{
+		Name:     name,
+		Replicas: 3,
+		Subjects: []string{"FOO.*"},
+	}
 
-			js, err := jetstream.New(nc)
-			if err != nil {
-				t.Fatalf("Unexpected error: %v", err)
-			}
-			defer nc.Close()
+	for _, stream := range []jetstream.StreamConfig{singleStream, streamWithReplicas} {
+		t.Run(fmt.Sprintf("num replicas: %d, no options", stream.Replicas), func(t *testing.T) {
+			withJSClusterAndStream(t, name, 3, stream, func(t *testing.T, subject string, srvs ...*jsServer) {
+				nc, err := nats.Connect(srvs[0].ClientURL())
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
 
-			ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
-			defer cancel()
-			s, err := js.Stream(ctx, stream.Name)
-			if err != nil {
-				t.Fatalf("Unexpected error: %v", err)
-			}
-			c, err := s.CreateOrUpdateConsumer(ctx, jetstream.ConsumerConfig{AckPolicy: jetstream.AckExplicitPolicy})
-			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()
 
-			msgs := make([]jetstream.Msg, 0)
-			wg := &sync.WaitGroup{}
-			wg.Add(len(testMsgs))
-			l, err := c.Consume(func(msg jetstream.Msg) {
-				msgs = append(msgs, msg)
-				wg.Done()
-			})
-			if err != nil {
-				t.Fatalf("Unexpected error: %v", err)
-			}
-			defer l.Stop()
+				ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+				defer cancel()
+				s, err := js.Stream(ctx, stream.Name)
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+				c, err := s.CreateOrUpdateConsumer(ctx, jetstream.ConsumerConfig{AckPolicy: jetstream.AckExplicitPolicy})
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
 
-			publishTestMsgs(t, nc)
-			wg.Wait()
-			if len(msgs) != len(testMsgs) {
-				t.Fatalf("Unexpected received message count; want %d; got %d", len(testMsgs), len(msgs))
-			}
-			for i, msg := range msgs {
-				if string(msg.Data()) != testMsgs[i] {
-					t.Fatalf("Invalid msg on index %d; expected: %s; got: %s", i, testMsgs[i], string(msg.Data()))
+				msgs := make([]jetstream.Msg, 0)
+				wg := &sync.WaitGroup{}
+				wg.Add(len(testMsgs))
+				l, err := c.Consume(func(msg jetstream.Msg) {
+					msgs = append(msgs, msg)
+					wg.Done()
+				})
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
 				}
-			}
+				defer l.Stop()
+
+				publishTestMsgs(t, js)
+				wg.Wait()
+				if len(msgs) != len(testMsgs) {
+					t.Fatalf("Unexpected received message count; want %d; got %d", len(testMsgs), len(msgs))
+				}
+				for i, msg := range msgs {
+					if string(msg.Data()) != testMsgs[i] {
+						t.Fatalf("Invalid msg on index %d; expected: %s; got: %s", i, testMsgs[i], string(msg.Data()))
+					}
+				}
+			})
 		})
-	})
 
-	t.Run("subscribe, cancel subscription, then subscribe again", func(t *testing.T) {
-		withJSClusterAndStream(t, name, 3, stream, func(t *testing.T, subject string, srvs ...*jsServer) {
-			nc, err := nats.Connect(srvs[0].ClientURL())
-			if err != nil {
-				t.Fatalf("Unexpected error: %v", err)
-			}
+		t.Run(fmt.Sprintf("num replicas: %d, subscribe, cancel subscription, then subscribe again", stream.Replicas), func(t *testing.T) {
+			withJSClusterAndStream(t, name, 3, stream, func(t *testing.T, subject string, srvs ...*jsServer) {
+				nc, err := nats.Connect(srvs[0].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()
+				js, err := jetstream.New(nc)
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+				defer nc.Close()
 
-			ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
-			s, err := js.Stream(ctx, stream.Name)
-			if err != nil {
-				t.Fatalf("Unexpected error: %v", err)
-			}
-			c, err := s.CreateOrUpdateConsumer(ctx, jetstream.ConsumerConfig{AckPolicy: jetstream.AckExplicitPolicy})
-			if err != nil {
-				t.Fatalf("Unexpected error: %v", err)
-			}
+				ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
+				defer cancel()
+				s, err := js.Stream(ctx, stream.Name)
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+				c, err := s.CreateOrUpdateConsumer(ctx, jetstream.ConsumerConfig{AckPolicy: jetstream.AckExplicitPolicy})
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
 
-			wg := sync.WaitGroup{}
-			wg.Add(len(testMsgs))
-			msgs := make([]jetstream.Msg, 0)
-			l, err := c.Consume(func(msg jetstream.Msg) {
-				if err := msg.Ack(); err != nil {
+				wg := sync.WaitGroup{}
+				wg.Add(len(testMsgs))
+				msgs := make([]jetstream.Msg, 0)
+				l, err := c.Consume(func(msg jetstream.Msg) {
+					if err := msg.Ack(); err != nil {
+						t.Fatalf("Unexpected error: %v", err)
+					}
+					msgs = append(msgs, msg)
+					if len(msgs) == 5 {
+						cancel()
+					}
+					wg.Done()
+				})
+				if err != nil {
 					t.Fatalf("Unexpected error: %v", err)
 				}
-				msgs = append(msgs, msg)
-				if len(msgs) == 5 {
-					cancel()
+
+				publishTestMsgs(t, js)
+				wg.Wait()
+				l.Stop()
+
+				time.Sleep(10 * time.Millisecond)
+				wg.Add(len(testMsgs))
+				l, err = c.Consume(func(msg jetstream.Msg) {
+					if err := msg.Ack(); err != nil {
+						t.Fatalf("Unexpected error: %v", err)
+					}
+					msgs = append(msgs, msg)
+					wg.Done()
+				})
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+				defer l.Stop()
+				publishTestMsgs(t, js)
+				wg.Wait()
+				if len(msgs) != 2*len(testMsgs) {
+					t.Fatalf("Unexpected received message count; want %d; got %d", len(testMsgs), len(msgs))
+				}
+				expectedMsgs := append(testMsgs, testMsgs...)
+				for i, msg := range msgs {
+					if string(msg.Data()) != expectedMsgs[i] {
+						t.Fatalf("Invalid msg on index %d; expected: %s; got: %s", i, testMsgs[i], string(msg.Data()))
+					}
 				}
-				wg.Done()
 			})
-			if err != nil {
-				t.Fatalf("Unexpected error: %v", err)
-			}
+		})
 
-			publishTestMsgs(t, nc)
-			wg.Wait()
-			l.Stop()
+		t.Run(fmt.Sprintf("num replicas: %d, recover consume after server restart", stream.Replicas), func(t *testing.T) {
+			withJSClusterAndStream(t, name, 3, stream, func(t *testing.T, subject string, srvs ...*jsServer) {
+				nc, err := nats.Connect(srvs[0].ClientURL())
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
 
-			time.Sleep(10 * time.Millisecond)
-			wg.Add(len(testMsgs))
-			defer cancel()
-			l, err = c.Consume(func(msg jetstream.Msg) {
-				if err := msg.Ack(); err != nil {
+				js, err := jetstream.New(nc)
+				if err != nil {
 					t.Fatalf("Unexpected error: %v", err)
 				}
-				msgs = append(msgs, msg)
-				wg.Done()
-			})
-			if err != nil {
-				t.Fatalf("Unexpected error: %v", err)
-			}
-			defer l.Stop()
-			publishTestMsgs(t, nc)
-			wg.Wait()
-			if len(msgs) != 2*len(testMsgs) {
-				t.Fatalf("Unexpected received message count; want %d; got %d", len(testMsgs), len(msgs))
-			}
-			expectedMsgs := append(testMsgs, testMsgs...)
-			for i, msg := range msgs {
-				if string(msg.Data()) != expectedMsgs[i] {
-					t.Fatalf("Invalid msg on index %d; expected: %s; got: %s", i, testMsgs[i], string(msg.Data()))
+				defer nc.Close()
+
+				ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
+				defer cancel()
+				s, err := js.Stream(ctx, streamWithReplicas.Name)
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
 				}
-			}
+				c, err := s.CreateOrUpdateConsumer(ctx, jetstream.ConsumerConfig{AckPolicy: jetstream.AckExplicitPolicy, InactiveThreshold: 10 * time.Second})
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+
+				wg := sync.WaitGroup{}
+				wg.Add(len(testMsgs))
+				msgs := make([]jetstream.Msg, 0)
+				l, err := c.Consume(func(msg jetstream.Msg) {
+					if err := msg.Ack(); err != nil {
+						t.Fatalf("Unexpected error: %v", err)
+					}
+					msgs = append(msgs, msg)
+					wg.Done()
+				}, jetstream.PullExpiry(1*time.Second), jetstream.PullHeartbeat(500*time.Millisecond))
+				if err != nil {
+					t.Fatalf("Unexpected error: %v", err)
+				}
+				defer l.Stop()
+
+				publishTestMsgs(t, js)
+				wg.Wait()
+
+				time.Sleep(10 * time.Millisecond)
+				srvs[0].Shutdown()
+				srvs[1].Shutdown()
+				srvs[0].Restart()
+				srvs[1].Restart()
+				wg.Add(len(testMsgs))
+
+				for i := 0; i < 10; i++ {
+					time.Sleep(500 * time.Millisecond)
+					if _, err := js.Stream(context.Background(), stream.Name); err == nil {
+						break
+					} else if i == 9 {
+						t.Fatal("JetStream not recovered: ", err)
+					}
+				}
+				publishTestMsgs(t, js)
+				wg.Wait()
+				if len(msgs) != 2*len(testMsgs) {
+					t.Fatalf("Unexpected received message count; want %d; got %d", len(testMsgs), len(msgs))
+				}
+				expectedMsgs := append(testMsgs, testMsgs...)
+				for i, msg := range msgs {
+					if string(msg.Data()) != expectedMsgs[i] {
+						t.Fatalf("Invalid msg on index %d; expected: %s; got: %s", i, testMsgs[i], string(msg.Data()))
+					}
+				}
+			})
 		})
-	})
+	}
 }
 
 func TestPullConsumerNext(t *testing.T) {
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")
+	}
+}