use switch/case instead of condition logic

livekit · Jun 4, 2024 · 6261b0b · 6261b0b
1 parent 9cd9cf3
commit 6261b0b
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 11 deletions.
diff --git a/pkg/config/output.go b/pkg/config/output.go
@@ -15,7 +15,7 @@
 package config
 
 import (
-	"strings"
+	"net/url"
 
 	"github.com/livekit/egress/pkg/errors"
 	"github.com/livekit/egress/pkg/types"
@@ -86,9 +86,19 @@ func (p *PipelineConfig) updateEncodedOutputs(req egress.EncodedOutput) error {
 		return errors.ErrInvalidInput("multiple stream outputs")
 	}
 	if stream != nil {
-		outputType := types.OutputTypeRTMP
-		if len(stream.Urls) > 0 && strings.HasPrefix(stream.Urls[0], "srt://") {
+		u, err := url.Parse(stream.Urls[0])
+		if err != nil {
+			return errors.ErrInvalidInput("malformed url")
+		}
+
+		var outputType types.OutputType
+		switch u.Scheme {
+		case "srt":
 			outputType = types.OutputTypeSRT
+		case "rtmp":
+			outputType = types.OutputTypeRTMP
+		default:
+			return errors.ErrInvalidInput("invalid stream type")
 		}
 
 		conf, err := p.getStreamConfig(outputType, stream.Urls)

diff --git a/pkg/pipeline/controller.go b/pkg/pipeline/controller.go
@@ -16,7 +16,7 @@ package pipeline
 
 import (
 	"context"
-	"strings"
+	"net/url"
 	"sync"
 	"time"
 
@@ -264,12 +264,22 @@ func (c *Controller) UpdateStream(ctx context.Context, req *livekit.UpdateStream
 
 	// add stream outputs first
 	for _, rawUrl := range req.AddOutputUrls {
-		outputType := types.OutputTypeRTMP
-		if strings.HasPrefix(rawUrl, "srt://") {
+
+		u, err := url.Parse(rawUrl)
+		if err != nil {
+			return errors.ErrInvalidInput("malformed url")
+		}
+
+		var outputType types.OutputType
+		switch u.Scheme {
+		case "srt":
 			outputType = types.OutputTypeSRT
+		case "rtmp":
+			outputType = types.OutputTypeRTMP
+		default:
+			return errors.ErrInvalidInput("invalid stream type")
 		}
-		// validate and redact url
-		//url, redacted, err := c.ValidateUrl(rawUrl, types.OutputTypeRTMP)
+
 		url, redacted, err := c.ValidateUrl(rawUrl, outputType)
 		if err != nil {
 			errs.AppendErr(err)
@@ -303,12 +313,23 @@ func (c *Controller) UpdateStream(ctx context.Context, req *livekit.UpdateStream
 
 	// remove stream outputs
 	for _, rawUrl := range req.RemoveOutputUrls {
-		outputType := types.OutputTypeRTMP
-		if strings.HasPrefix(rawUrl, "srt://") {
+		u, err := url.Parse(rawUrl)
+		if err != nil {
+			return errors.ErrInvalidInput("malformed url")
+		}
+
+		var outputType types.OutputType
+		switch u.Scheme {
+		case "srt":
 			outputType = types.OutputTypeSRT
+		case "rtmp":
+			outputType = types.OutputTypeRTMP
+		default:
+			return errors.ErrInvalidInput("invalid stream type")
 		}
+
 		url, _, err := c.ValidateUrl(rawUrl, outputType)
-		//url, _, err := c.ValidateUrl(rawUrl, types.OutputTypeRTMP)
+
 		if err != nil {
 			errs.AppendErr(err)
 			continue