func (d *Diagnostics) getDiagnosticFromPeer(ctx context.Context, p peer.ID, pmes *pb.Message) (<-chan *DiagInfo, error) { s, err := d.host.NewStream(ctx, ProtocolDiag, p) if err != nil { return nil, err } cr := ctxio.NewReader(ctx, s) // ok to use. we defer close stream in this func cw := ctxio.NewWriter(ctx, s) // ok to use. we defer close stream in this func r := ggio.NewDelimitedReader(cr, inet.MessageSizeMax) w := ggio.NewDelimitedWriter(cw) start := time.Now() if err := w.WriteMsg(pmes); err != nil { return nil, err } out := make(chan *DiagInfo) go func() { defer func() { close(out) s.Close() rtt := time.Since(start) log.Infof("diagnostic request took: %s", rtt.String()) }() for { rpmes := new(pb.Message) if err := r.ReadMsg(rpmes); err != nil { log.Debugf("Error reading diagnostic from stream: %s", err) return } if rpmes == nil { log.Debug("got no response back from diag request") return } di, err := decodeDiagJson(rpmes.GetData()) if err != nil { log.Debug(err) return } select { case out <- di: case <-ctx.Done(): return } } }() return out, nil }
func (d *Diagnostics) HandleMessage(ctx context.Context, s inet.Stream) error { cr := ctxio.NewReader(ctx, s) cw := ctxio.NewWriter(ctx, s) r := ggio.NewDelimitedReader(cr, inet.MessageSizeMax) // maxsize w := ggio.NewDelimitedWriter(cw) // deserialize msg pmes := new(pb.Message) if err := r.ReadMsg(pmes); err != nil { log.Debugf("Failed to decode protobuf message: %v", err) return nil } // Print out diagnostic log.Infof("[peer: %s] Got message from [%s]\n", d.self.Pretty(), s.Conn().RemotePeer()) // Make sure we havent already handled this request to prevent loops if err := d.startDiag(pmes.GetDiagID()); err != nil { return nil } resp := newMessage(pmes.GetDiagID()) resp.Data = d.getDiagInfo().Marshal() if err := w.WriteMsg(resp); err != nil { log.Debugf("Failed to write protobuf message over stream: %s", err) return err } timeout := pmes.GetTimeoutDuration() if timeout < HopTimeoutDecrement { return fmt.Errorf("timeout too short: %s", timeout) } ctx, cancel := context.WithTimeout(ctx, timeout) defer cancel() pmes.SetTimeoutDuration(timeout - HopTimeoutDecrement) dpeers, err := d.getDiagnosticFromPeers(ctx, d.getPeers(), pmes) if err != nil { log.Debugf("diagnostic from peers err: %s", err) return err } for b := range dpeers { resp := newMessage(pmes.GetDiagID()) resp.Data = b.Marshal() if err := w.WriteMsg(resp); err != nil { log.Debugf("Failed to write protobuf message over stream: %s", err) return err } } return nil }
func (dht *IpfsDHT) handleNewMessage(s inet.Stream) { defer s.Close() ctx := dht.Context() cr := ctxio.NewReader(ctx, s) // ok to use. we defer close stream in this func cw := ctxio.NewWriter(ctx, s) // ok to use. we defer close stream in this func r := ggio.NewDelimitedReader(cr, inet.MessageSizeMax) w := ggio.NewDelimitedWriter(cw) mPeer := s.Conn().RemotePeer() for { // receive msg pmes := new(pb.Message) if err := r.ReadMsg(pmes); err != nil { log.Debugf("Error unmarshaling data: %s", err) return } // update the peer (on valid msgs only) dht.updateFromMessage(ctx, mPeer, pmes) // get handler for this msg type. handler := dht.handlerForMsgType(pmes.GetType()) if handler == nil { log.Debug("got back nil handler from handlerForMsgType") return } // dispatch handler. rpmes, err := handler(ctx, mPeer, pmes) if err != nil { log.Debugf("handle message error: %s", err) return } // if nil response, return it before serializing if rpmes == nil { log.Debug("got back nil response from request") continue } // send out response msg if err := w.WriteMsg(rpmes); err != nil { log.Debugf("send response error: %s", err) return } } }