Merge pull request #114 from joestarzxh/master

[opt] rtmp: 对ping request回应ping response(包含ClientSession和ServerSession)

pkg/base/rtmp_t.go

@@ -24,8 +24,10 @@ const (
 	RtmpTypeIdAggregateMessage   uint8 = 22
 
 	// user control message type
-	RtmpUserControlStreamBegin uint8 = 0
+	RtmpUserControlStreamBegin  uint8 = 0
-	RtmpUserControlRecorded    uint8 = 4
+	RtmpUserControlRecorded     uint8 = 4
+	RtmpUserControlPingRequest  uint8 = 6
+	RtmpUserControlPingResponse uint8 = 7
 
 	// spec-video_file_format_spec_v10.pdf
 	// Video tags

pkg/rtmp/client_session.go

@@ -367,6 +367,7 @@ func (s *ClientSession) doMsg(stream *Stream) error {
 			nazalog.Warnf("[%s] read user control message, ignore. buf=%s",
 				s.uniqueKey, hex.Dump(stream.msg.buff.Peek(32)))
 		}
+		s.doUserControl(stream)
 	case base.RtmpTypeIdAudio:
 		fallthrough
 	case base.RtmpTypeIdVideo:
@@ -383,6 +384,15 @@ func (s *ClientSession) doAck(stream *Stream) error {
 	nazalog.Infof("[%s] < R Acknowledgement. ignore. sequence number=%d.", s.uniqueKey, seqNum)
 	return nil
 }
+func (s *ClientSession) doUserControl(stream *Stream) error {
+	userControlType := bele.BeUint16(stream.msg.buff.Bytes())
+	if userControlType == uint16(base.RtmpUserControlPingRequest) {
+		stream.msg.buff.Skip(2)
+		timeStamp := bele.BeUint32(stream.msg.buff.Bytes())
+		s.packer.writePingResponse(s.conn, timeStamp)
+	}
+	return nil
+}
 
 func (s *ClientSession) doDataMessageAmf0(stream *Stream) error {
 	val, err := stream.msg.peekStringWithType()

pkg/rtmp/message_packer.go

@@ -253,6 +253,26 @@ func (packer *MessagePacker) writeStreamBegin(writer io.Writer, streamid uint32)
 	return packer.ChunkAndWrite(writer, csidProtocolControl, base.RtmpTypeIdUserControl, 0)
 }
 
+func (packer *MessagePacker) writePingRequest(writer io.Writer, timeStamp uint32) error {
+	packer.b.ModWritePos(12)
+
+	// 6
+	_ = bele.WriteBe(packer.b, uint16(base.RtmpUserControlPingRequest))
+	_ = bele.WriteBe(packer.b, uint32(timeStamp))
+
+	return packer.ChunkAndWrite(writer, csidProtocolControl, base.RtmpTypeIdUserControl, 0)
+}
+
+func (packer *MessagePacker) writePingResponse(writer io.Writer, timeStamp uint32) error {
+	packer.b.ModWritePos(12)
+
+	// 6
+	_ = bele.WriteBe(packer.b, uint16(base.RtmpUserControlPingResponse))
+	_ = bele.WriteBe(packer.b, uint32(timeStamp))
+
+	return packer.ChunkAndWrite(writer, csidProtocolControl, base.RtmpTypeIdUserControl, 0)
+}
+
 // ---------------------------------------------------------------------------------------------------------------------
 
 // TODO(chef): 整理所有的buffer

pkg/rtmp/server_session.go

@@ -219,6 +219,8 @@ func (s *ServerSession) doMsg(stream *Stream) error {
 		return s.doDataMessageAmf0(stream)
 	case base.RtmpTypeIdAck:
 		return s.doAck(stream)
+	case base.RtmpTypeIdUserControl:
+		s.doUserControl(stream)
 	case base.RtmpTypeIdAudio:
 		fallthrough
 	case base.RtmpTypeIdVideo:
@@ -238,7 +240,15 @@ func (s *ServerSession) doAck(stream *Stream) error {
 	nazalog.Infof("[%s] < R Acknowledgement. ignore. sequence number=%d.", s.uniqueKey, seqNum)
 	return nil
 }
-
+func (s *ServerSession) doUserControl(stream *Stream) error {
+	userControlType := bele.BeUint16(stream.msg.buff.Bytes())
+	if userControlType == uint16(base.RtmpUserControlPingRequest) {
+		stream.msg.buff.Skip(2)
+		timeStamp := bele.BeUint32(stream.msg.buff.Bytes())
+		s.packer.writePingResponse(s.conn, timeStamp)
+	}
+	return nil
+}
 func (s *ServerSession) doDataMessageAmf0(stream *Stream) error {
 	if s.t != ServerSessionTypePub {
 		return nazaerrors.Wrap(base.ErrRtmpUnexpectedMsg)