Implementation for server enforcement of keepalive policy.

keepalive/keepalive.go

@@ -67,3 +67,12 @@ type ServerParameters struct {
 	// the connection is closed.
 	Timeout time.Duration
 }
+
+// EnforcementPolicy is used to set keepalive enforcement policy on the server-side.
+// Server will close connection with a client that violates this policy.
+type EnforcementPolicy struct {
+	// MinTime is the minimum amount of time a client should wait before sending a keepalive ping.
+	MinTime time.Duration
+	// If true, server expects keepalive pings even when there are no active streams(RPCs).
+	PermitWithoutStream bool
+}

server.go

@@ -119,6 +119,7 @@ type options struct {
 	useHandlerImpl       bool // use http.Handler-based server
 	unknownStreamDesc    *StreamDesc
 	keepaliveParams      keepalive.ServerParameters
+	keepalivePolicy      keepalive.EnforcementPolicy
 }
 
 var defaultMaxMsgSize = 1024 * 1024 * 4 // use 4MB as the default message size limit
@@ -133,6 +134,13 @@ func KeepaliveParams(kp keepalive.ServerParameters) ServerOption {
 	}
 }
 
+// KeepaliveEnforcementPolicy returns a ServerOption that sets keepalive enforcement policy for the server.
+func KeepaliveEnforcementPolicy(kep keepalive.EnforcementPolicy) ServerOption {
+	return func(o *options) {
+		o.keepalivePolicy = kep
+	}
+}
+
 // CustomCodec returns a ServerOption that sets a codec for message marshaling and unmarshaling.
 func CustomCodec(codec Codec) ServerOption {
 	return func(o *options) {
@@ -479,6 +487,7 @@ func (s *Server) serveHTTP2Transport(c net.Conn, authInfo credentials.AuthInfo)
 		InTapHandle:     s.opts.inTapHandle,
 		StatsHandler:    s.opts.statsHandler,
 		KeepaliveParams: s.opts.keepaliveParams,
+		KeepalivePolicy: s.opts.keepalivePolicy,
 	}
 	st, err := transport.NewServerTransport("http2", c, config)
 	if err != nil {

transport/control.go

@@ -57,6 +57,7 @@ const (
 	defaultMaxConnectionAgeGrace  = infinity
 	defaultServerKeepaliveTime    = time.Duration(2 * time.Hour)
 	defaultServerKeepaliveTimeout = time.Duration(20 * time.Second)
+	defaultKeepalivePolicyMinTime = time.Duration(5 * time.Minute)
 )
 
 // The following defines various control items which could flow through
@@ -84,6 +85,8 @@ type resetStream struct {
 func (*resetStream) item() {}
 
 type goAway struct {
+	code      http2.ErrCode
+	debugData []byte
 }
 
 func (*goAway) item() {}

transport/http2_server.go

@@ -100,6 +100,17 @@ type http2Server struct {
 	// Keepalive and max-age parameters for the server.
 	kp keepalive.ServerParameters
 
+	// Keepalive enforcement policy.
+	kep keepalive.EnforcementPolicy
+	// The time instance last ping was received.
+	lastPingAt time.Time
+	// Number of times the client has violated keepalive ping policy so far.
+	pingStrikes uint8
+	// Flag to signify that number of ping strikes should be reset to 0.
+	// This is set whenever data or header frames are sent.
+	// 1 means yes.
+	resetPingStrikes uint32 // Accessed atomically.
+
 	mu            sync.Mutex // guard the following
 	state         transportState
 	activeStreams map[uint32]*Stream
@@ -161,6 +172,10 @@ func newHTTP2Server(conn net.Conn, config *ServerConfig) (_ ServerTransport, err
 	if kp.Timeout == 0 {
 		kp.Timeout = defaultServerKeepaliveTimeout
 	}
+	kep := config.KeepalivePolicy
+	if kep.MinTime == 0 {
+		kep.MinTime = defaultKeepalivePolicyMinTime
+	}
 	var buf bytes.Buffer
 	t := &http2Server{
 		ctx:             context.Background(),
@@ -184,6 +199,7 @@ func newHTTP2Server(conn net.Conn, config *ServerConfig) (_ ServerTransport, err
 		stats:           config.StatsHandler,
 		kp:              kp,
 		idle:            time.Now(),
+		kep:             kep,
 	}
 	if t.stats != nil {
 		t.ctx = t.stats.TagConn(t.ctx, &stats.ConnTagInfo{
@@ -504,13 +520,49 @@ func (t *http2Server) handleSettings(f *http2.SettingsFrame) {
 	t.controlBuf.put(&settings{ack: true, ss: ss})
 }
 
+const (
+	maxPingStrikes = 2
+)
+
 func (t *http2Server) handlePing(f *http2.PingFrame) {
 	if f.IsAck() { // Do nothing.
 		return
 	}
 	pingAck := &ping{ack: true}
 	copy(pingAck.data[:], f.Data[:])
 	t.controlBuf.put(pingAck)
+
+	now := time.Now()
+	defer func() {
+		t.lastPingAt = now
+	}()
+	// A reset ping strikes means that we don't need to check for policy
+	// violation for this ping and the pingStrikes counter should be set
+	// to 0.
+	if atomic.CompareAndSwapUint32(&t.resetPingStrikes, 1, 0) {
+		t.pingStrikes = 0
+		return
+	}
+	t.mu.Lock()
+	ns := len(t.activeStreams)
+	t.mu.Unlock()
+	if ns < 1 && !t.kep.PermitWithoutStream {
+		// Keepalive shouldn't be active thus, this new ping should
+		// have come after atleast two hours.
+		if t.lastPingAt.Add(2 * time.Hour).After(now) {
+			t.pingStrikes++
+		}
+	} else {
+		// Check if keepalive policy is respected.
+		if t.lastPingAt.Add(t.kep.MinTime).After(now) {
+			t.pingStrikes++
+		}
+	}
+
+	if t.pingStrikes > maxPingStrikes {
+		// Send goaway and close the connection.
+		t.controlBuf.put(&goAway{code: http2.ErrCodeEnhanceYourCalm, debugData: []byte("too_many_pings")})
+	}
 }
 
 func (t *http2Server) handleWindowUpdate(f *http2.WindowUpdateFrame) {
@@ -529,6 +581,13 @@ func (t *http2Server) writeHeaders(s *Stream, b *bytes.Buffer, endStream bool) e
 	first := true
 	endHeaders := false
 	var err error
+	defer func() {
+		if err == nil {
+			// Reset ping strikes when seding headers since that might cause the
+			// peer to send ping.
+			atomic.StoreUint32(&t.resetPingStrikes, 1)
+		}
+	}()
 	// Sends the headers in a single batch.
 	for !endHeaders {
 		size := t.hBuf.Len()
@@ -672,7 +731,7 @@ func (t *http2Server) WriteStatus(s *Stream, statusCode codes.Code, statusDesc s
 
 // Write converts the data into HTTP2 data frame and sends it out. Non-nil error
 // is returns if it fails (e.g., framing error, transport error).
-func (t *http2Server) Write(s *Stream, data []byte, opts *Options) error {
+func (t *http2Server) Write(s *Stream, data []byte, opts *Options) (err error) {
 	// TODO(zhaoq): Support multi-writers for a single stream.
 	var writeHeaderFrame bool
 	s.mu.Lock()
@@ -687,6 +746,13 @@ func (t *http2Server) Write(s *Stream, data []byte, opts *Options) error {
 	if writeHeaderFrame {
 		t.WriteHeader(s, nil)
 	}
+	defer func() {
+		if err == nil {
+			// Reset ping strikes when sending data since this might cause
+			// the peer to send ping.
+			atomic.StoreUint32(&t.resetPingStrikes, 1)
+		}
+	}()
 	r := bytes.NewBuffer(data)
 	for {
 		if r.Len() == 0 {
@@ -892,7 +958,10 @@ func (t *http2Server) controller() {
 					sid := t.maxStreamID
 					t.state = draining
 					t.mu.Unlock()
-					t.framer.writeGoAway(true, sid, http2.ErrCodeNo, nil)
+					t.framer.writeGoAway(true, sid, i.code, i.debugData)
+					if i.code == http2.ErrCodeEnhanceYourCalm {
+						t.Close()
+					}
 				case *flushIO:
 					t.framer.flushWrite()
 				case *ping:
@@ -972,7 +1041,7 @@ func (t *http2Server) RemoteAddr() net.Addr {
 }
 
 func (t *http2Server) Drain() {
-	t.controlBuf.put(&goAway{})
+	t.controlBuf.put(&goAway{code: http2.ErrCodeNo})
 }
 
 var rgen = rand.New(rand.NewSource(time.Now().UnixNano()))

transport/transport.go

@@ -370,6 +370,7 @@ type ServerConfig struct {
 	InTapHandle     tap.ServerInHandle
 	StatsHandler    stats.Handler
 	KeepaliveParams keepalive.ServerParameters
+	KeepalivePolicy keepalive.EnforcementPolicy
 }
 
 // NewServerTransport creates a ServerTransport with conn or non-nil error

transport/transport_test.go

@@ -323,6 +323,9 @@ func TestMaxConnectionIdle(t *testing.T) {
 	timeout := time.NewTimer(time.Second * 4)
 	select {
 	case <-client.GoAway():
+		if !timeout.Stop() {
+			<-timeout.C
+		}
 	case <-timeout.C:
 		t.Fatalf("Test timed out, expected a GoAway from the server.")
 	}
@@ -345,6 +348,9 @@ func TestMaxConnectionIdleNegative(t *testing.T) {
 	timeout := time.NewTimer(time.Second * 4)
 	select {
 	case <-client.GoAway():
+		if !timeout.Stop() {
+			<-timeout.C
+		}
 		t.Fatalf("A non-idle client received a GoAway.")
 	case <-timeout.C:
 	}
@@ -369,6 +375,9 @@ func TestMaxConnectionAge(t *testing.T) {
 	timeout := time.NewTimer(4 * time.Second)
 	select {
 	case <-client.GoAway():
+		if !timeout.Stop() {
+			<-timeout.C
+		}
 	case <-timeout.C:
 		t.Fatalf("Test timer out, expected a GoAway from the server.")
 	}
@@ -523,6 +532,126 @@ func TestKeepaliveClientStaysHealthyWithResponsiveServer(t *testing.T) {
 	}
 }
 
+func TestKeepaliveServerEnforcementWithAbusiveClientNoRPC(t *testing.T) {
+	serverConfig := &ServerConfig{
+		KeepalivePolicy: keepalive.EnforcementPolicy{
+			MinTime: 2 * time.Second,
+		},
+	}
+	clientOptions := ConnectOptions{
+		KeepaliveParams: keepalive.ClientParameters{
+			Time:                50 * time.Millisecond,
+			Timeout:             50 * time.Millisecond,
+			PermitWithoutStream: true,
+		},
+	}
+	server, client := setUpWithOptions(t, 0, serverConfig, normal, clientOptions)
+	defer server.stop()
+	defer client.Close()
+
+	timeout := time.NewTimer(2 * time.Second)
+	select {
+	case <-client.GoAway():
+		if !timeout.Stop() {
+			<-timeout.C
+		}
+	case <-timeout.C:
+		t.Fatalf("Test failed: Expected a GoAway from server.")
+	}
+}
+
+func TestKeepaliveServerEnforcementWithAbusiveClientWithRPC(t *testing.T) {
+	serverConfig := &ServerConfig{
+		KeepalivePolicy: keepalive.EnforcementPolicy{
+			MinTime: 2 * time.Second,
+		},
+	}
+	clientOptions := ConnectOptions{
+		KeepaliveParams: keepalive.ClientParameters{
+			Time:    50 * time.Millisecond,
+			Timeout: 50 * time.Millisecond,
+		},
+	}
+	server, client := setUpWithOptions(t, 0, serverConfig, suspended, clientOptions)
+	defer server.stop()
+	defer client.Close()
+
+	if _, err := client.NewStream(context.Background(), &CallHdr{Flush: true}); err != nil {
+		t.Fatalf("Client failed to create stream.")
+	}
+	timeout := time.NewTimer(2 * time.Second)
+	select {
+	case <-client.GoAway():
+		if !timeout.Stop() {
+			<-timeout.C
+		}
+	case <-timeout.C:
+		t.Fatalf("Test failed: Expected a GoAway from server.")
+	}
+
+}
+
+func TestKeepaliveServerEnforcementWithObeyingClientNoRPC(t *testing.T) {
+	serverConfig := &ServerConfig{
+		KeepalivePolicy: keepalive.EnforcementPolicy{
+			MinTime:             100 * time.Millisecond,
+			PermitWithoutStream: true,
+		},
+	}
+	clientOptions := ConnectOptions{
+		KeepaliveParams: keepalive.ClientParameters{
+			Time:                100 * time.Millisecond,
+			Timeout:             50 * time.Millisecond,
+			PermitWithoutStream: true,
+		},
+	}
+	server, client := setUpWithOptions(t, 0, serverConfig, normal, clientOptions)
+	defer server.stop()
+	defer client.Close()
+
+	// Give keepalive enough time.
+	time.Sleep(2 * time.Second)
+	// Asser that connection is healthy
+	ct := client.(*http2Client)
+	ct.mu.Lock()
+	defer ct.mu.Unlock()
+	if ct.state != reachable {
+		t.Fatalf("Test failed: Expected connection to be healthy.")
+	}
+}
+
+func TestKeepaliveServerEnforcementWithObeyingClientWithRPC(t *testing.T) {
+	serverConfig := &ServerConfig{
+		KeepalivePolicy: keepalive.EnforcementPolicy{
+			MinTime: 100 * time.Millisecond,
+		},
+	}
+	clientOptions := ConnectOptions{
+		KeepaliveParams: keepalive.ClientParameters{
+			Time:    100 * time.Millisecond,
+			Timeout: 50 * time.Millisecond,
+		},
+	}
+	server, client := setUpWithOptions(t, 0, serverConfig, suspended, clientOptions)
+	defer server.stop()
+	defer client.Close()
+
+	if _, err := client.NewStream(context.Background(), &CallHdr{Flush: true}); err != nil {
+		t.Fatalf("Client failed to create stream.")
+	}
+
+	// Give keepalive enough time.
+	time.Sleep(2 * time.Second)
+	// Asser that connection is healthy
+	ct := client.(*http2Client)
+	ct.mu.Lock()
+	defer ct.mu.Unlock()
+	if ct.state != reachable {
+		t.Fatalf("Test failed: Expected connection to be healthy.")
+	}
+
+}
+
 func TestClientSendAndReceive(t *testing.T) {
 	server, ct := setUp(t, 0, math.MaxUint32, normal)
 	callHdr := &CallHdr{