mpegps.go 3.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159
  1. package mpegps
  2. import (
  3. "encoding/binary"
  4. "errors"
  5. "io"
  6. "m7s.live/engine/v4/util"
  7. )
  8. const (
  9. StartCodePS = 0x000001ba
  10. StartCodeSYS = 0x000001bb
  11. StartCodeMAP = 0x000001bc
  12. StartCodeVideo = 0x000001e0
  13. StartCodeAudio = 0x000001c0
  14. PrivateStreamCode = 0x000001bd
  15. MEPGProgramEndCode = 0x000001b9
  16. )
  17. type EsHandler interface {
  18. ReceiveAudio(MpegPsEsStream)
  19. ReceiveVideo(MpegPsEsStream)
  20. ReceivePSM(util.Buffer, bool, bool)
  21. }
  22. type MpegPsEsStream struct {
  23. Type byte
  24. util.Buffer
  25. PTS uint32
  26. DTS uint32
  27. }
  28. type MpegPsStream struct {
  29. buffer util.Buffer
  30. EsHandler
  31. audio MpegPsEsStream
  32. video MpegPsEsStream
  33. }
  34. func (ps *MpegPsStream) Reset() {
  35. ps.buffer.Reset()
  36. ps.audio.Reset()
  37. if ps.video.Buffer.CanRead() {
  38. ps.ReceiveVideo(ps.video)
  39. ps.video.Buffer = make(util.Buffer, 0)
  40. } else {
  41. ps.video.Reset()
  42. }
  43. }
  44. func (ps *MpegPsStream) Feed(data util.Buffer) (err error) {
  45. reader := &data
  46. if ps.buffer.CanRead() {
  47. ps.buffer.Write(data)
  48. reader = &ps.buffer
  49. }
  50. var begin util.Buffer
  51. var payload []byte
  52. var frame MpegPsEsStream
  53. defer func() {
  54. if err != nil && begin.CanRead() {
  55. ps.buffer.Reset()
  56. ps.buffer.Write(begin)
  57. }
  58. }()
  59. for err == nil && reader.CanReadN(4) {
  60. begin = *reader
  61. code := reader.ReadUint32()
  62. switch code {
  63. case StartCodePS:
  64. if reader.CanReadN(9) {
  65. reader.ReadN(9)
  66. if reader.CanRead() {
  67. psl := reader.ReadByte() & 0x07
  68. if reader.CanReadN(int(psl)) {
  69. reader.ReadN(int(psl))
  70. continue
  71. }
  72. }
  73. }
  74. err = io.ErrShortBuffer
  75. case StartCodeSYS, PrivateStreamCode:
  76. _, err = ps.ReadPayload(reader)
  77. case StartCodeMAP:
  78. err = ps.decProgramStreamMap(reader)
  79. case StartCodeVideo:
  80. payload, err = ps.ReadPayload(reader)
  81. if err == nil {
  82. frame, err = ps.video.parsePESPacket(payload)
  83. if err == nil {
  84. ps.ReceiveVideo(frame)
  85. }
  86. }
  87. case StartCodeAudio:
  88. payload, err = ps.ReadPayload(reader)
  89. if err == nil {
  90. frame, err = ps.audio.parsePESPacket(payload)
  91. if err == nil {
  92. ps.ReceiveAudio(frame)
  93. }
  94. }
  95. case MEPGProgramEndCode:
  96. return
  97. default:
  98. err = errors.New("start code error")
  99. }
  100. }
  101. return
  102. }
  103. func (ps *MpegPsStream) ReadPayload(data *util.Buffer) (payload []byte, err error) {
  104. if !data.CanReadN(2) {
  105. return nil, io.ErrShortBuffer
  106. }
  107. payloadlen := data.ReadUint16()
  108. if data.CanReadN(int(payloadlen)) {
  109. payload = data.ReadN(int(payloadlen))
  110. } else {
  111. err = io.ErrShortBuffer
  112. }
  113. return
  114. }
  115. func (ps *MpegPsStream) decProgramStreamMap(data *util.Buffer) error {
  116. psm, err := ps.ReadPayload(data)
  117. if err != nil {
  118. return err
  119. }
  120. l := len(psm)
  121. index := 2
  122. programStreamInfoLen := binary.BigEndian.Uint16(psm[index:])
  123. index += 2
  124. index += int(programStreamInfoLen)
  125. programStreamMapLen := binary.BigEndian.Uint16(psm[index:])
  126. index += 2
  127. for programStreamMapLen > 0 {
  128. if l <= index+1 {
  129. break
  130. }
  131. streamType := psm[index]
  132. index++
  133. elementaryStreamID := psm[index]
  134. index++
  135. if elementaryStreamID >= 0xe0 && elementaryStreamID <= 0xef {
  136. ps.video.Type = streamType
  137. } else if elementaryStreamID >= 0xc0 && elementaryStreamID <= 0xdf {
  138. ps.audio.Type = streamType
  139. }
  140. if l <= index+1 {
  141. break
  142. }
  143. elementaryStreamInfoLength := binary.BigEndian.Uint16(psm[index:])
  144. index += 2
  145. index += int(elementaryStreamInfoLength)
  146. programStreamMapLen -= 4 + elementaryStreamInfoLength
  147. }
  148. return nil
  149. }