Merge pull request #2 from tekdoc/script

Add option to run a script when client connects or disconnects (Linux…

README.md

@@ -14,6 +14,7 @@ Features:
 * Each stream can have multiple video and audio tracks
 * Supports the RTP/RTCP streaming protocol
 * Optional authentication mechanism for publishers
+* Run a script when a client connects or disconnects
 * Compatible with Linux and Windows, does not require any dependency or interpreter, it's a single executable
 
 ## Installation
@@ -73,6 +74,8 @@ Flags:
   --rtp-port=8000        port of the RTP UDP listener
   --rtcp-port=8001       port of the RTCP UDP listener
   --publish-key=""       optional authentication key required to publish
+  --pre-script=""        optional script to run on client connect
+  --post-script=""       optional script to run on client disconnect
 ```
 
 ## Links

client.go

@@ -6,6 +6,7 @@ import (
 	"log"
 	"net"
 	"net/url"
+	"os/exec"
 	"strconv"
 	"strings"
 	"time"
@@ -162,7 +163,18 @@ func (c *client) log(format string, args ...interface{}) {
 }
 
 func (c *client) run() {
+	defer func() {
+		if c.p.postScript != "" {
+			postScript := exec.Command(c.p.postScript)
+			err := postScript.Run()
+			if err != nil {
+				c.log("ERR: %s", err)
+			}
+		}
+	}()
+
 	defer c.log("disconnected")
+
 	defer func() {
 		c.p.mutex.Lock()
 		defer c.p.mutex.Unlock()
@@ -174,6 +186,14 @@ func (c *client) run() {
 
 	c.log("connected")
 
+	if c.p.preScript != "" {
+		preScript := exec.Command(c.p.preScript)
+		err := preScript.Run()
+		if err != nil {
+			c.log("ERR: %s", err)
+		}
+	}
+
 	for {
 		req, err := c.conn.ReadRequest()
 		if err != nil {

main.go

@@ -53,6 +53,8 @@ type program struct {
 	rtpPort    int
 	rtcpPort   int
 	publishKey string
+	preScript  string
+	postScript string
 	mutex      sync.RWMutex
 	rtspl      *serverTcpListener
 	rtpl       *serverUdpListener
@@ -61,7 +63,7 @@ type program struct {
 	publishers map[string]*client
 }
 
-func newProgram(protocolsStr string, rtspPort int, rtpPort int, rtcpPort int, publishKey string) (*program, error) {
+func newProgram(protocolsStr string, rtspPort int, rtpPort int, rtcpPort int, publishKey string, preScript string, postScript string) (*program, error) {
 
 	if rtspPort == 0 {
 		return nil, fmt.Errorf("rtsp port not provided")
@@ -114,6 +116,8 @@ func newProgram(protocolsStr string, rtspPort int, rtpPort int, rtcpPort int, pu
 		rtpPort:    rtpPort,
 		rtcpPort:   rtcpPort,
 		publishKey: publishKey,
+		preScript:  preScript,
+		postScript: postScript,
 		clients:    make(map[*client]struct{}),
 		publishers: make(map[string]*client),
 	}
@@ -189,6 +193,8 @@ func main() {
 	rtpPort := kingpin.Flag("rtp-port", "port of the RTP UDP listener").Default("8000").Int()
 	rtcpPort := kingpin.Flag("rtcp-port", "port of the RTCP UDP listener").Default("8001").Int()
 	publishKey := kingpin.Flag("publish-key", "optional authentication key required to publish").Default("").String()
+	preScript := kingpin.Flag("pre-script", "optional script to run on client connect").Default("").String()
+	postScript := kingpin.Flag("post-script", "optional script to run on client disconnect").Default("").String()
 
 	kingpin.Parse()
 
@@ -197,7 +203,7 @@ func main() {
 		os.Exit(0)
 	}
 
-	p, err := newProgram(*protocols, *rtspPort, *rtpPort, *rtcpPort, *publishKey)
+	p, err := newProgram(*protocols, *rtspPort, *rtpPort, *rtcpPort, *publishKey, *preScript, *postScript)
 	if err != nil {
 		log.Fatal("ERR: ", err)
 	}