stream-server/srt/tunnel.go

package srt

import (
	"net"
	"fmt"
)

type Tunnel struct {
	socket net.PacketConn
	peer net.Addr
	queue chan []byte
	broken bool
	state *SRTManager
}

func (tunnel *Tunnel) Start() {
	defer func(a *bool) {
		if r := recover(); r != nil {
			fmt.Println(r)
		}
		*a = true
	}(&(tunnel.broken)) // force mark tunnel for deletion if any error occurs
	tunnel.state = NewSRTManager(tunnel.socket)
	// central tunnel loop, read incoming, process and generate response
	// write response if any
	for {
		packet, err := tunnel.ReadPacket()
		if err != nil {
			fmt.Println(err)
			tunnel.broken = true
		}
		response, err := tunnel.state.Process(packet)
		if err != nil {
			fmt.Println(err)
			tunnel.broken = true
		}
		if response != nil {
			tunnel.WritePacket(response)
		}
	}
}

// send a shutdown command, for use when tunnel gets broken
// not ideal but works
func (tunnel *Tunnel) Shutdown() {
	if tunnel.state != nil && tunnel.state.state > 1 {
		packet := tunnel.state.create_basic_header()
		packet.packet_type = SHUTDOWN
		info := new(ControlHeader)
		info.ctrl_type = 5
		packet.header_info = info
		tunnel.WritePacket(packet)
		if tunnel.state.output != nil {
			tunnel.state.output.Close()
		}
	}
}

func (tunnel *Tunnel) WritePacket(packet *Packet) {
	buffer, err := MarshallPacket(packet, tunnel.state)
	if err != nil {
		tunnel.broken = true
		fmt.Println(err)
		return
	}
	tunnel.socket.WriteTo(buffer, tunnel.peer)
}

func (tunnel *Tunnel) ReadPacket() (*Packet, error) {
	packet := <- tunnel.queue // blocking read, should add timeout here
	return ParsePacket(packet)
}