diff --git a/server.go b/server.go index 79f17f3d..0d9a5915 100644 --- a/server.go +++ b/server.go @@ -85,6 +85,7 @@ type Server struct { } type options struct { + handshaker func(net.Conn) error codec Codec maxConcurrentStreams uint32 } @@ -92,6 +93,12 @@ type options struct { // A ServerOption sets options. type ServerOption func(*options) +func Handshaker(f func(net.Conn) error) ServerOption { + return func(o *options) { + o.handshaker = f + } +} + func CustomCodec(codec Codec) ServerOption { return func(o *options) { o.codec = codec @@ -185,7 +192,14 @@ func (s *Server) Serve(lis net.Listener) error { if err != nil { return err } - + // Perform handshaking if it is required. + if s.opts.handshaker != nil { + if err := s.opts.handshaker(c); err != nil { + log.Println("grpc: Server.Serve failed to complete handshake.") + c.Close() + continue + } + } s.mu.Lock() if s.conns == nil { s.mu.Unlock()