Add a handshaker option to server

This commit is contained in:
iamqizhao 2015-04-30 16:30:48 -07:00
Родитель 58e0450a3c
Коммит c6f6e4701d
1 изменённых файлов: 15 добавлений и 1 удалений

Просмотреть файл

@ -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()