github.com/cnotch/ipchub@v1.1.0/service/rtsp/multicast_proxy.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 rtsp
     6  
     7  import (
     8  	"fmt"
     9  	"io"
    10  	"net"
    11  	"sync"
    12  
    13  	"github.com/cnotch/ipchub/media"
    14  	"github.com/cnotch/ipchub/network"
    15  	"github.com/cnotch/xlog"
    16  )
    17  
    18  // 组播代理
    19  type multicastProxy struct {
    20  	// 创建时设置
    21  	logger     *xlog.Logger
    22  	path       string
    23  	bufferSize int
    24  
    25  	multicastIP string
    26  	ports       [rtpChannelCount]int
    27  	ttl         int
    28  	sourceIP    string
    29  
    30  	closed   bool
    31  	udpConn  *net.UDPConn
    32  	destAddr [rtpChannelCount]*net.UDPAddr
    33  	cid      media.CID
    34  
    35  	multicastLock sync.Mutex
    36  	members       []io.Closer
    37  }
    38  
    39  func (proxy *multicastProxy) AddMember(m io.Closer) {
    40  	proxy.multicastLock.Lock()
    41  	defer proxy.multicastLock.Unlock()
    42  
    43  	if len(proxy.members) == 0 {
    44  		stream := media.Get(proxy.path)
    45  		if stream == nil {
    46  			proxy.logger.Error("start multicast proxy failed.")
    47  			return
    48  		}
    49  
    50  		udpConn, err := net.ListenUDP("udp", &net.UDPAddr{})
    51  		if err != nil {
    52  			proxy.logger.Errorf("start multicast proxy failed. %s", err.Error())
    53  			return
    54  		}
    55  
    56  		proxy.udpConn = udpConn
    57  		err = udpConn.SetWriteBuffer(proxy.bufferSize)
    58  
    59  		for i, port := range proxy.ports {
    60  			if port > 0 {
    61  				proxy.destAddr[i], _ = net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", proxy.multicastIP, proxy.ports[i]))
    62  			}
    63  		}
    64  
    65  		proxy.members = append(proxy.members, m)
    66  		proxy.cid = stream.StartConsume(proxy, media.RTPPacket,
    67  			"net = rtsp-multicast, "+proxy.multicastIP)
    68  		proxy.closed = false
    69  
    70  		proxy.logger.Info("multicast proxy started.")
    71  	}
    72  }
    73  
    74  func (proxy *multicastProxy) ReleaseMember(m io.Closer) {
    75  	proxy.multicastLock.Lock()
    76  	defer proxy.multicastLock.Unlock()
    77  	for i, m2 := range proxy.members {
    78  		if m == m2 {
    79  			proxy.members = append(proxy.members[:i], proxy.members[i+1:]...)
    80  			break
    81  		}
    82  	}
    83  
    84  	if len(proxy.members) == 0 {
    85  		// 停止组播代理
    86  		proxy.close()
    87  	}
    88  }
    89  
    90  func (proxy *multicastProxy) MulticastIP() string {
    91  	return proxy.multicastIP
    92  }
    93  
    94  func (proxy *multicastProxy) Port(index int) int {
    95  	if index < 0 || index > len(proxy.ports) {
    96  		return 0
    97  	}
    98  	return proxy.ports[index]
    99  }
   100  
   101  func (proxy *multicastProxy) TTL() int {
   102  	return proxy.ttl
   103  }
   104  
   105  func (proxy *multicastProxy) SourceIP() string {
   106  	if len(proxy.sourceIP) == 0 {
   107  		addrs := network.GetLocalIP()
   108  		if len(addrs) == 0 {
   109  			proxy.sourceIP = "Unknown"
   110  		} else {
   111  			proxy.sourceIP = addrs[0]
   112  		}
   113  	}
   114  	return proxy.sourceIP
   115  }
   116  
   117  func (proxy *multicastProxy) Consume(p Pack) {
   118  	if proxy.closed {
   119  		return
   120  	}
   121  
   122  	p2 := p.(*RTPPack)
   123  	addr := proxy.destAddr[int(p2.Channel)]
   124  	if addr != nil {
   125  		_, err := proxy.udpConn.WriteToUDP(p2.Data, addr)
   126  		if err != nil {
   127  			proxy.logger.Error(err.Error())
   128  			return
   129  		}
   130  	}
   131  }
   132  
   133  func (proxy *multicastProxy) Close() error {
   134  	proxy.multicastLock.Lock()
   135  	defer proxy.multicastLock.Unlock()
   136  
   137  	proxy.close()
   138  	return nil
   139  }
   140  
   141  func (proxy *multicastProxy) close() {
   142  	if proxy.closed {
   143  		return
   144  	}
   145  	proxy.closed = true
   146  
   147  	stream := media.Get(proxy.path)
   148  	if stream != nil {
   149  		stream.StopConsume(proxy.cid)
   150  	}
   151  
   152  	if proxy.udpConn != nil {
   153  		proxy.udpConn.Close()
   154  		proxy.udpConn = nil
   155  	}
   156  
   157  	// 关闭所有的组播客户端
   158  	for _, m := range proxy.members {
   159  		m.Close()
   160  	}
   161  	proxy.members = nil
   162  
   163  	for i := range proxy.destAddr {
   164  		proxy.destAddr[i] = nil
   165  	}
   166  
   167  	proxy.logger.Info("multicast proxy stopped.")
   168  }