@@ -30,6 +30,7 @@ import (
30
30
31
31
"google.golang.org/grpc"
32
32
"google.golang.org/grpc/codes"
33
+ "google.golang.org/grpc/credentials"
33
34
"google.golang.org/grpc/health/grpc_health_v1"
34
35
"google.golang.org/grpc/status"
35
36
)
@@ -49,13 +50,19 @@ var (
49
50
delayUnhealthySec int
50
51
service string
51
52
forceUnhealthy * bool
53
+ certFile string
54
+ privKeyFile string
52
55
)
53
56
54
57
func init () {
55
58
CmdGrpcHealthChecking .Flags ().IntVar (& port , "port" , 5000 , "Port number." )
56
59
CmdGrpcHealthChecking .Flags ().IntVar (& httpPort , "http-port" , 8080 , "Port number for the /make-serving and /make-not-serving." )
57
60
CmdGrpcHealthChecking .Flags ().IntVar (& delayUnhealthySec , "delay-unhealthy-sec" , - 1 , "Number of seconds to delay before start reporting NOT_SERVING, negative value indicates never." )
58
61
CmdGrpcHealthChecking .Flags ().StringVar (& service , "service" , "" , "Service name to register the health check for." )
62
+ CmdGrpcHealthChecking .Flags ().StringVar (& certFile , "tls-cert-file" , "" ,
63
+ "File containing an x509 certificate for gRPC TLS. (CA cert, if any, concatenated after server cert)." )
64
+ CmdGrpcHealthChecking .Flags ().StringVar (& privKeyFile , "tls-private-key-file" , "" ,
65
+ "File containing an x509 private key matching --tls-cert-file." )
59
66
forceUnhealthy = nil
60
67
}
61
68
@@ -95,6 +102,13 @@ func NewHealthChecker(started time.Time) *HealthChecker {
95
102
func main (cmd * cobra.Command , args []string ) {
96
103
started := time .Now ()
97
104
105
+ // Validate flags
106
+ //
107
+ // if certFile or privKeyFile are not both set, exit with error
108
+ if (certFile == "" && privKeyFile != "" ) || (certFile != "" && privKeyFile == "" ) {
109
+ log .Fatalf ("Both --tls-cert-file and --tls-private-key-file must be set" )
110
+ }
111
+
98
112
http .HandleFunc ("/make-not-serving" , func (w http.ResponseWriter , r * http.Request ) {
99
113
log .Printf ("Mark as unhealthy" )
100
114
forceUnhealthy = new (bool )
@@ -121,17 +135,29 @@ func main(cmd *cobra.Command, args []string) {
121
135
122
136
serverAdr := fmt .Sprintf (":%d" , port )
123
137
listenAddr , err := net .Listen ("tcp" , serverAdr )
138
+
124
139
if err != nil {
125
- log .Fatal (fmt .Sprintf ("Error while starting the listening service %v" , err .Error ()))
140
+ log .Fatalf ("Error while starting the listening service %v" , err )
141
+ }
142
+
143
+ var grpcServer * grpc.Server
144
+
145
+ if certFile != "" && privKeyFile != "" {
146
+ creds , err := credentials .NewServerTLSFromFile (certFile , privKeyFile )
147
+ if err != nil {
148
+ log .Fatalf ("Failed to generate credentials %v" , err )
149
+ }
150
+ grpcServer = grpc .NewServer (grpc .Creds (creds ))
151
+ } else {
152
+ grpcServer = grpc .NewServer ()
126
153
}
127
154
128
- grpcServer := grpc .NewServer ()
129
155
healthService := NewHealthChecker (started )
130
156
grpc_health_v1 .RegisterHealthServer (grpcServer , healthService )
131
157
132
158
log .Printf ("gRPC server starting to listen on %s" , serverAdr )
133
159
if err = grpcServer .Serve (listenAddr ); err != nil {
134
- log .Fatal ( fmt . Sprintf ( "Error while starting the gRPC server on the %s listen address %v" , listenAddr , err . Error ()) )
160
+ log .Fatalf ( "Error while starting the gRPC server on the %s listen address %v" , listenAddr , err )
135
161
}
136
162
137
163
select {}
0 commit comments