From c45a54d7cb916c22443f5484284ea72a623b4e12 Mon Sep 17 00:00:00 2001
From: Stefan Sedich <stefan.sedich@gmail.com>
Date: Wed, 18 Nov 2020 07:40:31 -0800
Subject: [PATCH] Add support for configuring gRPC max-connection-age for the
 kiam server (#440)

* Add support for configuring gRPC server parameters for the kiam server.
---
 cmd/kiam/server.go           | 5 +++++
 pkg/server/server.go         | 2 ++
 pkg/server/server_builder.go | 1 +
 3 files changed, 8 insertions(+)

diff --git a/cmd/kiam/server.go b/cmd/kiam/server.go
index f534134a..1379c33e 100644
--- a/cmd/kiam/server.go
+++ b/cmd/kiam/server.go
@@ -59,6 +59,11 @@ func (o *serverOptions) bind(parser parser) {
 	parser.Flag("session-refresh", "How soon STS Tokens should be refreshed before their expiration.").Default("5m").DurationVar(&o.SessionRefresh)
 	parser.Flag("assume-role-arn", "IAM Role to assume before processing requests").Default("").StringVar(&o.AssumeRoleArn)
 	parser.Flag("region", "AWS Region to use for regional STS calls (e.g. us-west-2). Defaults to the global endpoint.").Default("").StringVar(&o.Region)
+	parser.Flag("grpc-keepalive-time-duration", "gRPC keepalive time").Default("10s").DurationVar(&o.KeepaliveParams.Time)
+	parser.Flag("grpc-keepalive-timeout-duration", "gRPC keepalive timeout").Default("2s").DurationVar(&o.KeepaliveParams.Timeout)
+	parser.Flag("grpc-max-connection-idle-duration", "gRPC max connection idle").Default("15m").DurationVar(&o.KeepaliveParams.MaxConnectionIdle)
+	parser.Flag("grpc-max-connection-age-duration", "gRPC max connection age").Default("15m").DurationVar(&o.KeepaliveParams.MaxConnectionAge)
+	parser.Flag("grpc-max-connection-age-grace-duration", "gRPC max connection age grace").Default("15m").DurationVar(&o.KeepaliveParams.MaxConnectionAgeGrace)
 }
 
 func (cmd *serverCommand) Run() {
diff --git a/pkg/server/server.go b/pkg/server/server.go
index f081bf70..c89619ba 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -27,6 +27,7 @@ import (
 	"github.com/uswitch/kiam/pkg/prefetch"
 	pb "github.com/uswitch/kiam/proto"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/keepalive"
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/client-go/tools/record"
@@ -48,6 +49,7 @@ type Config struct {
 	PrefetchBufferSize           int
 	AssumeRoleArn                string
 	Region                       string
+	KeepaliveParams              keepalive.ServerParameters
 }
 
 // TLSConfig controls TLS
diff --git a/pkg/server/server_builder.go b/pkg/server/server_builder.go
index a84982c6..d22e5967 100644
--- a/pkg/server/server_builder.go
+++ b/pkg/server/server_builder.go
@@ -169,6 +169,7 @@ func (b *KiamServerBuilder) WithTLS() (*KiamServerBuilder, error) {
 		grpc.Creds(b.transportCredentials),
 		grpc.StreamInterceptor(grpc_prometheus.StreamServerInterceptor),
 		grpc.UnaryInterceptor(grpc_prometheus.UnaryServerInterceptor),
+		grpc.KeepaliveParams(b.config.KeepaliveParams),
 	)
 
 	return b, nil