1/*
2 *
3 * Copyright 2017 gRPC authors.
4 *
5 * Licensed under the Apache License, Version 2.0 (the "License");
6 * you may not use this file except in compliance with the License.
7 * You may obtain a copy of the License at
8 *
9 *     http://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
16 *
17 */
18
19package test
20
21import (
22	"context"
23	"fmt"
24	"net"
25	"sync"
26	"testing"
27	"time"
28
29	"google.golang.org/grpc"
30	"google.golang.org/grpc/codes"
31	"google.golang.org/grpc/status"
32	testpb "google.golang.org/grpc/test/grpc_testing"
33)
34
35type delayListener struct {
36	net.Listener
37	closeCalled  chan struct{}
38	acceptCalled chan struct{}
39	allowCloseCh chan struct{}
40	dialed       bool
41}
42
43func (d *delayListener) Accept() (net.Conn, error) {
44	select {
45	case <-d.acceptCalled:
46		// On the second call, block until closed, then return an error.
47		<-d.closeCalled
48		<-d.allowCloseCh
49		return nil, fmt.Errorf("listener is closed")
50	default:
51		close(d.acceptCalled)
52		conn, err := d.Listener.Accept()
53		if err != nil {
54			return nil, err
55		}
56		// Allow closing of listener only after accept.
57		// Note: Dial can return successfully, yet Accept
58		// might now have finished.
59		d.allowClose()
60		return conn, nil
61	}
62}
63
64func (d *delayListener) allowClose() {
65	close(d.allowCloseCh)
66}
67func (d *delayListener) Close() error {
68	close(d.closeCalled)
69	go func() {
70		<-d.allowCloseCh
71		d.Listener.Close()
72	}()
73	return nil
74}
75
76func (d *delayListener) Dial(ctx context.Context) (net.Conn, error) {
77	if d.dialed {
78		// Only hand out one connection (net.Dial can return more even after the
79		// listener is closed).  This is not thread-safe, but Dial should never be
80		// called concurrently in this environment.
81		return nil, fmt.Errorf("no more conns")
82	}
83	d.dialed = true
84	return (&net.Dialer{}).DialContext(ctx, "tcp", d.Listener.Addr().String())
85}
86
87func (s) TestGracefulStop(t *testing.T) {
88	// This test ensures GracefulStop causes new connections to fail.
89	//
90	// Steps of this test:
91	// 1. Start Server
92	// 2. GracefulStop() Server after listener's Accept is called, but don't
93	//    allow Accept() to exit when Close() is called on it.
94	// 3. Create a new connection to the server after listener.Close() is called.
95	//    Server should close this connection immediately, before handshaking.
96	// 4. Send an RPC on the new connection.  Should see Unavailable error
97	//    because the ClientConn is in transient failure.
98	lis, err := net.Listen("tcp", "localhost:0")
99	if err != nil {
100		t.Fatalf("Error listenening: %v", err)
101	}
102	dlis := &delayListener{
103		Listener:     lis,
104		acceptCalled: make(chan struct{}),
105		closeCalled:  make(chan struct{}),
106		allowCloseCh: make(chan struct{}),
107	}
108	d := func(ctx context.Context, _ string) (net.Conn, error) { return dlis.Dial(ctx) }
109
110	ss := &stubServer{
111		fullDuplexCall: func(stream testpb.TestService_FullDuplexCallServer) error {
112			_, err := stream.Recv()
113			if err != nil {
114				return err
115			}
116			return stream.Send(&testpb.StreamingOutputCallResponse{})
117		},
118	}
119	s := grpc.NewServer()
120	testpb.RegisterTestServiceServer(s, ss)
121
122	// 1. Start Server
123	wg := sync.WaitGroup{}
124	wg.Add(1)
125	go func() {
126		s.Serve(dlis)
127		wg.Done()
128	}()
129
130	// 2. GracefulStop() Server after listener's Accept is called, but don't
131	//    allow Accept() to exit when Close() is called on it.
132	<-dlis.acceptCalled
133	wg.Add(1)
134	go func() {
135		s.GracefulStop()
136		wg.Done()
137	}()
138
139	// 3. Create a new connection to the server after listener.Close() is called.
140	//    Server should close this connection immediately, before handshaking.
141
142	<-dlis.closeCalled // Block until GracefulStop calls dlis.Close()
143
144	// Now dial.  The listener's Accept method will return a valid connection,
145	// even though GracefulStop has closed the listener.
146	ctx, dialCancel := context.WithTimeout(context.Background(), 5*time.Second)
147	defer dialCancel()
148	cc, err := grpc.DialContext(ctx, "", grpc.WithInsecure(), grpc.WithContextDialer(d))
149	if err != nil {
150		t.Fatalf("grpc.DialContext(_, %q, _) = %v", lis.Addr().String(), err)
151	}
152	client := testpb.NewTestServiceClient(cc)
153	defer cc.Close()
154
155	// 4. Send an RPC on the new connection.
156	// The server would send a GOAWAY first, but we are delaying the server's
157	// writes for now until the client writes more than the preface.
158	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
159	if _, err = client.FullDuplexCall(ctx); err == nil || status.Code(err) != codes.Unavailable {
160		t.Fatalf("FullDuplexCall= _, %v; want _, <status code Unavailable>", err)
161	}
162	cancel()
163	wg.Wait()
164}
165