[streaming] move response packet after registration

Prevent a race where a client may attempt to use a stream
before it is registered.

Signed-off-by: Derek McGowan <derek@mcg.dev>
This commit is contained in:
Derek McGowan 2022-12-30 21:56:25 -08:00
parent b0b28f1d8e
commit d7ef6cbfa3
No known key found for this signature in database
GPG Key ID: F58C5D0A4405ACDB

View File

@ -74,20 +74,23 @@ func (s *service) Stream(srv api.Streaming_StreamServer) error {
if err != nil { if err != nil {
return err return err
} }
if err := srv.Send(protobuf.FromAny(response)); err != nil {
return err
}
cc := make(chan struct{}) cc := make(chan struct{})
ss := &serviceStream{ ss := &serviceStream{
s: srv, s: srv,
cc: cc, cc: cc,
} }
log.G(srv.Context()).WithField("stream", i.ID).Debug("registering stream") log.G(srv.Context()).WithField("stream", i.ID).Debug("registering stream")
if err := s.manager.Register(srv.Context(), i.ID, ss); err != nil { if err := s.manager.Register(srv.Context(), i.ID, ss); err != nil {
return err return err
} }
// Send response packet after registering stream
if err := srv.Send(protobuf.FromAny(response)); err != nil {
return err
}
select { select {
case <-srv.Context().Done(): case <-srv.Context().Done():
// TODO: Should return error if not cancelled? // TODO: Should return error if not cancelled?