github.com/cnotch/ipchub@v1.1.0/service/flv/wsflv.go (about)

     1  // Copyright (c) 2019,CAOHONGJU All rights reserved.
     2  // Use of this source code is governed by a MIT-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package flv
     6  
     7  import (
     8  	"runtime/debug"
     9  	"time"
    10  
    11  	"github.com/cnotch/ipchub/media"
    12  	"github.com/cnotch/ipchub/network/websocket"
    13  	"github.com/cnotch/ipchub/av/format/flv"
    14  	"github.com/cnotch/ipchub/stats"
    15  	"github.com/cnotch/xlog"
    16  )
    17  
    18  type wsFlvConsumer struct {
    19  	logger *xlog.Logger
    20  	w      *flv.Writer
    21  	conn   websocket.Conn
    22  	closed bool
    23  }
    24  
    25  func (c *wsFlvConsumer) Consume(pack Pack) {
    26  	if c.closed {
    27  		return
    28  	}
    29  
    30  	err := c.w.WriteFlvTag(pack.(*flv.Tag))
    31  
    32  	if err != nil {
    33  		c.logger.Errorf("ws-flv: send tag failed; %v", err)
    34  		c.Close()
    35  		return
    36  	}
    37  }
    38  
    39  func (c *wsFlvConsumer) Close() (err error) {
    40  	if c.closed {
    41  		return
    42  	}
    43  
    44  	c.closed = true
    45  	c.conn.Close()
    46  	return nil
    47  }
    48  
    49  func (c *wsFlvConsumer) Type() string {
    50  	return "websocket-flv"
    51  }
    52  
    53  // ConsumeByWebsocket 处理 websocket 方式访问流媒体
    54  func ConsumeByWebsocket(logger *xlog.Logger, path string, addr string, conn websocket.Conn) {
    55  	logger = logger.With(xlog.Fields(
    56  		xlog.F("path", path),xlog.F("ext", "flv"),
    57  		xlog.F("addr", addr)))
    58  
    59  	stream := media.GetOrCreate(path)
    60  	if stream == nil {
    61  		conn.Close()
    62  		logger.Errorf("ws-flv: no stream found")
    63  		return
    64  	}
    65  
    66  	typeFlags := stream.FlvTypeFlags()
    67  	if typeFlags == 0 {
    68  		conn.Close()
    69  		logger.Errorf("ws-flv: stream not support flv")
    70  		return
    71  	}
    72  
    73  	var cid media.CID
    74  
    75  	defer func() {
    76  		if r := recover(); r != nil {
    77  			xlog.Errorf("ws-flv: panic; %v \n %s", r, debug.Stack())
    78  		}
    79  		stream.StopConsume(cid)
    80  		conn.Close()
    81  		stats.FlvConns.Release()
    82  		logger.Info("stop websocket-flv consume")
    83  	}()
    84  
    85  	logger.Info("start websocket-flv consume")
    86  	stats.FlvConns.Add()
    87  
    88  	// 启动 pack 消费,必须 StartConsume 前写入 Header
    89  	flvWriter, err := flv.NewWriter(conn, typeFlags)
    90  	if err != nil {
    91  		logger.Error("ws-flv: send flv header failed.")
    92  		return
    93  	}
    94  
    95  	c := &wsFlvConsumer{
    96  		logger: logger,
    97  		conn:   conn,
    98  		w:      flvWriter,
    99  	}
   100  
   101  	cid = stream.StartConsume(c, media.FLVPacket, "net=websocket-flv,"+addr)
   102  
   103  	for !c.closed {
   104  		deadLine := time.Time{}
   105  
   106  		if err := conn.SetReadDeadline(deadLine); err != nil {
   107  			break
   108  		}
   109  		var temp [1]byte
   110  		if _, err := conn.Read(temp[:]); err != nil {
   111  			if !c.closed {
   112  				logger.Errorf("websocket error; %v.", err)
   113  			}
   114  			break
   115  		}
   116  	}
   117  }