server: rewrite working principle

examples/server/main.go

@@ -11,154 +11,140 @@ import (
 
 // This example shows how to
 // 1. create a RTSP server which accepts plain connections
-// 2. allow a single client to publish a stream with TCP
-// 3. allow multiple clients to read that stream with TCP
+// 2. allow a single client to publish a stream with TCP or UDP
+// 3. allow multiple clients to read that stream with TCP or UDP
 
-var mutex sync.Mutex
-var publisher *gortsplib.ServerConn
-var readers = make(map[*gortsplib.ServerConn]struct{})
-var sdp []byte
+type serverHandler struct {
+	mutex     sync.Mutex
+	publisher *gortsplib.ServerConn
+	readers   map[*gortsplib.ServerConn]struct{}
+	sdp       []byte
+}
 
-// this is called for each incoming connection
-func handleConn(conn *gortsplib.ServerConn) {
-	defer conn.Close()
+// called when a connection is opened.
+func (sh *serverHandler) OnConnOpen(sc *gortsplib.ServerConn) {
+	log.Printf("conn opened")
+}
 
-	log.Printf("client connected")
+// called when a connection is closed.
+func (sh *serverHandler) OnConnClose(sc *gortsplib.ServerConn, err error) {
+	log.Println("conn closed (%v)", err)
 
-	// called after receiving a DESCRIBE request.
-	onDescribe := func(ctx *gortsplib.ServerConnDescribeCtx) (*base.Response, []byte, error) {
-		mutex.Lock()
-		defer mutex.Unlock()
+	sh.mutex.Lock()
+	defer sh.mutex.Unlock()
 
-		// no one is publishing yet
-		if publisher == nil {
-			return &base.Response{
-				StatusCode: base.StatusNotFound,
-			}, nil, nil
-		}
-
-		return &base.Response{
-			StatusCode: base.StatusOK,
-		}, sdp, nil
-	}
-
-	// called after receiving an ANNOUNCE request.
-	onAnnounce := func(ctx *gortsplib.ServerConnAnnounceCtx) (*base.Response, error) {
-		mutex.Lock()
-		defer mutex.Unlock()
-
-		if publisher != nil {
-			return &base.Response{
-				StatusCode: base.StatusBadRequest,
-			}, fmt.Errorf("someone is already publishing")
-		}
-
-		publisher = conn
-		sdp = ctx.Tracks.Write()
-
-		return &base.Response{
-			StatusCode: base.StatusOK,
-			Header: base.Header{
-				"Session": base.HeaderValue{"12345678"},
-			},
-		}, nil
-	}
-
-	// called after receiving a SETUP request.
-	onSetup := func(ctx *gortsplib.ServerConnSetupCtx) (*base.Response, error) {
-		return &base.Response{
-			StatusCode: base.StatusOK,
-			Header: base.Header{
-				"Session": base.HeaderValue{"12345678"},
-			},
-		}, nil
-	}
-
-	// called after receiving a PLAY request.
-	onPlay := func(ctx *gortsplib.ServerConnPlayCtx) (*base.Response, error) {
-		mutex.Lock()
-		defer mutex.Unlock()
-
-		readers[conn] = struct{}{}
-
-		return &base.Response{
-			StatusCode: base.StatusOK,
-			Header: base.Header{
-				"Session": base.HeaderValue{"12345678"},
-			},
-		}, nil
-	}
-
-	// called after receiving a RECORD request.
-	onRecord := func(ctx *gortsplib.ServerConnRecordCtx) (*base.Response, error) {
-		mutex.Lock()
-		defer mutex.Unlock()
-
-		if conn != publisher {
-			return &base.Response{
-				StatusCode: base.StatusBadRequest,
-			}, fmt.Errorf("someone is already publishing")
-		}
-
-		return &base.Response{
-			StatusCode: base.StatusOK,
-			Header: base.Header{
-				"Session": base.HeaderValue{"12345678"},
-			},
-		}, nil
-	}
-
-	// called after receiving a frame.
-	onFrame := func(trackID int, typ gortsplib.StreamType, buf []byte) {
-		mutex.Lock()
-		defer mutex.Unlock()
-
-		// if we are the publisher, route frames to readers
-		if conn == publisher {
-			for r := range readers {
-				r.WriteFrame(trackID, typ, buf)
-			}
-		}
-	}
-
-	err := <-conn.Read(gortsplib.ServerConnReadHandlers{
-		OnDescribe: onDescribe,
-		OnAnnounce: onAnnounce,
-		OnSetup:    onSetup,
-		OnPlay:     onPlay,
-		OnRecord:   onRecord,
-		OnFrame:    onFrame,
-	})
-	log.Printf("client disconnected (%s)", err)
-
-	mutex.Lock()
-	defer mutex.Unlock()
-
-	if conn == publisher {
-		publisher = nil
-		sdp = nil
+	if sc == sh.publisher {
+		sh.publisher = nil
+		sh.sdp = nil
 	} else {
-		delete(readers, conn)
+		delete(sh.readers, sc)
+	}
+}
+
+// called after receiving a DESCRIBE request.
+func (sh *serverHandler) OnDescribe(ctx *gortsplib.ServerHandlerOnDescribeCtx) (*base.Response, []byte, error) {
+	sh.mutex.Lock()
+	defer sh.mutex.Unlock()
+
+	// no one is publishing yet
+	if sh.publisher == nil {
+		return &base.Response{
+			StatusCode: base.StatusNotFound,
+		}, nil, nil
+	}
+
+	return &base.Response{
+		StatusCode: base.StatusOK,
+	}, sh.sdp, nil
+}
+
+// called after receiving an ANNOUNCE request.
+func (sh *serverHandler) OnAnnounce(ctx *gortsplib.ServerHandlerOnAnnounceCtx) (*base.Response, error) {
+	sh.mutex.Lock()
+	defer sh.mutex.Unlock()
+
+	if sh.publisher != nil {
+		return &base.Response{
+			StatusCode: base.StatusBadRequest,
+		}, fmt.Errorf("someone is already publishing")
+	}
+
+	sh.publisher = ctx.Conn
+	sh.sdp = ctx.Tracks.Write()
+
+	return &base.Response{
+		StatusCode: base.StatusOK,
+		Header: base.Header{
+			"Session": base.HeaderValue{"12345678"},
+		},
+	}, nil
+}
+
+// called after receiving a SETUP request.
+func (sh *serverHandler) OnSetup(ctx *gortsplib.ServerHandlerOnSetupCtx) (*base.Response, error) {
+	return &base.Response{
+		StatusCode: base.StatusOK,
+		Header: base.Header{
+			"Session": base.HeaderValue{"12345678"},
+		},
+	}, nil
+}
+
+// called after receiving a PLAY request.
+func (sh *serverHandler) OnPlay(ctx *gortsplib.ServerHandlerOnPlayCtx) (*base.Response, error) {
+	sh.mutex.Lock()
+	defer sh.mutex.Unlock()
+
+	sh.readers[ctx.Conn] = struct{}{}
+
+	return &base.Response{
+		StatusCode: base.StatusOK,
+		Header: base.Header{
+			"Session": base.HeaderValue{"12345678"},
+		},
+	}, nil
+}
+
+// called after receiving a RECORD request.
+func (sh *serverHandler) OnRecord(ctx *gortsplib.ServerHandlerOnRecordCtx) (*base.Response, error) {
+	sh.mutex.Lock()
+	defer sh.mutex.Unlock()
+
+	if ctx.Conn != sh.publisher {
+		return &base.Response{
+			StatusCode: base.StatusBadRequest,
+		}, fmt.Errorf("someone is already publishing")
+	}
+
+	return &base.Response{
+		StatusCode: base.StatusOK,
+		Header: base.Header{
+			"Session": base.HeaderValue{"12345678"},
+		},
+	}, nil
+}
+
+// called after receiving a frame.
+func (sh *serverHandler) OnFrame(ctx *gortsplib.ServerHandlerOnFrameCtx) {
+	sh.mutex.Lock()
+	defer sh.mutex.Unlock()
+
+	// if we are the publisher, route frames to readers
+	if ctx.Conn == sh.publisher {
+		for r := range sh.readers {
+			r.WriteFrame(ctx.TrackID, ctx.StreamType, ctx.Payload)
+		}
 	}
 }
 
 func main() {
-	// create server
-	s := &gortsplib.Server{}
-	err := s.Serve(":8554")
-	if err != nil {
-		panic(err)
+	// configure server
+	s := &gortsplib.Server{
+		Handler: &serverHandler{},
+		UDPRTPAddress:  ":8000",
+		UDPRTCPAddress: ":8001",
 	}
 
-	log.Printf("server is ready")
-
-	// accept connections
-	for {
-		conn, err := s.Accept()
-		if err != nil {
-			panic(err)
-		}
-
-		go handleConn(conn)
-	}
+	// start server and wait until a fatal error
+	panic(s.StartAndWait(":8554"))
 }