AI: Fix #1: Fix RTMP protocol crash bug.

Author: ossrs-ai

env.go

@@ -81,7 +81,7 @@ func buildDefaultEnvironmentVariables(ctx context.Context) {
 	// The API server of proxy itself.
 	setEnvDefault("PROXY_SYSTEM_API", "12025")
 	// The static directory for web server, optional.
-	setEnvDefault("PROXY_STATIC_FILES", "")
+	setEnvDefault("PROXY_STATIC_FILES", "../srs/trunk/research")
 
 	// The load balancer, use redis or memory.
 	setEnvDefault("PROXY_LOAD_BALANCER_TYPE", "memory")

rtmp.go

@@ -6,12 +6,10 @@ package main
 import (
 	"context"
 	"fmt"
-	"math/rand"
 	"net"
 	"strconv"
 	"strings"
 	"sync"
-	"time"
 
 	"srs-proxy/errors"
 	"srs-proxy/logger"
@@ -24,16 +22,12 @@ import (
 type srsRTMPServer struct {
 	// The TCP listener for RTMP server.
 	listener *net.TCPListener
-	// The random number generator.
-	rd *rand.Rand
 	// The wait group for all goroutines.
 	wg sync.WaitGroup
 }
 
 func NewSRSRTMPServer(opts ...func(*srsRTMPServer)) *srsRTMPServer {
-	v := &srsRTMPServer{
-		rd: rand.New(rand.NewSource(time.Now().UnixNano())),
-	}
+	v := &srsRTMPServer{}
 	for _, opt := range opts {
 		opt(v)
 	}
@@ -97,9 +91,7 @@ func (v *srsRTMPServer) Run(ctx context.Context) error {
 					}
 				}
 
-				rc := NewRTMPConnection(func(client *RTMPConnection) {
-					client.rd = v.rd
-				})
+				rc := NewRTMPConnection()
 				if err := rc.serve(ctx, conn); err != nil {
 					handleErr(err)
 				} else {
@@ -119,8 +111,6 @@ func (v *srsRTMPServer) Run(ctx context.Context) error {
 // then proxy to the corresponding backend server. All state is in the RTMP request, so this
 // connection is stateless.
 type RTMPConnection struct {
-	// The random number generator.
-	rd *rand.Rand
 }
 
 func NewRTMPConnection(opts ...func(*RTMPConnection)) *RTMPConnection {
@@ -151,7 +141,7 @@ func (v *RTMPConnection) serve(ctx context.Context, conn *net.TCPConn) error {
 	}
 
 	// Simple handshake with client.
-	hs := rtmp.NewHandshake(v.rd)
+	hs := rtmp.NewHandshake()
 	if _, err := hs.ReadC0S0(conn); err != nil {
 		return errors.Wrapf(err, "read c0")
 	}
@@ -294,7 +284,7 @@ func (v *RTMPConnection) serve(ctx context.Context, conn *net.TCPConn) error {
 
 	// Find a backend SRS server to proxy the RTMP stream.
 	backend = NewRTMPClientToBackend(func(client *RTMPClientToBackend) {
-		client.rd, client.typ = v.rd, clientType
+		client.typ = clientType
 	})
 	defer backend.Close()
 
@@ -412,8 +402,6 @@ const (
 
 // RTMPClientToBackend is a RTMP client to proxy the RTMP stream to backend.
 type RTMPClientToBackend struct {
-	// The random number generator.
-	rd *rand.Rand
 	// The underlayer tcp client.
 	tcpConn *net.TCPConn
 	// The RTMP protocol client.
@@ -470,7 +458,7 @@ func (v *RTMPClientToBackend) Connect(ctx context.Context, tcUrl, streamName str
 	}
 	v.tcpConn = c
 
-	hs := rtmp.NewHandshake(v.rd)
+	hs := rtmp.NewHandshake()
 	client := rtmp.NewProtocol(c)
 	v.client = client
 

rtmp/rtmp.go

@@ -7,26 +7,24 @@ import (
 	"bufio"
 	"bytes"
 	"context"
+	"crypto/rand"
 	"encoding"
 	"encoding/binary"
 	"fmt"
 	"io"
-	"math/rand"
 	"sync"
 
 	"srs-proxy/errors"
 )
 
 // The handshake implements the RTMP handshake protocol.
 type Handshake struct {
-	// The random number generator.
-	r *rand.Rand
 	// The c1s1 cache.
 	c1s1 []byte
 }
 
-func NewHandshake(r *rand.Rand) *Handshake {
-	return &Handshake{r: r}
+func NewHandshake() *Handshake {
+	return &Handshake{}
 }
 
 func (v *Handshake) C1S1() []byte {
@@ -56,8 +54,9 @@ func (v *Handshake) ReadC0S0(r io.Reader) (c0 []byte, err error) {
 func (v *Handshake) WriteC1S1(w io.Writer) (err error) {
 	p := make([]byte, 1536)
 
-	for i := 8; i < len(p); i++ {
-		p[i] = byte(v.r.Int())
+	// Use crypto/rand for thread-safe random generation
+	if _, err = rand.Read(p[8:]); err != nil {
+		return errors.Wrap(err, "generate random bytes")
 	}
 
 	r := bytes.NewReader(p)

srs.go

@@ -242,7 +242,8 @@ func (v *srsMemoryLoadBalancer) Pick(ctx context.Context, streamURL string) (*SR
 		return nil, fmt.Errorf("no server available for %v", streamURL)
 	}
 
-	// Pick a server randomly from servers.
+	// Pick a server randomly from servers. Use global rand which is thread-safe since Go 1.20.
+	// For older Go versions, this is still safe as we're only reading from the servers slice.
 	server := servers[rand.Intn(len(servers))]
 	v.picked.Store(streamURL, server)
 	return server, nil
@@ -420,7 +421,7 @@ func (v *srsRedisLoadBalancer) Pick(ctx context.Context, streamURL string) (*SRS
 	}
 
 	// All server should be alive, if not, should have been removed by redis. So we only
-	// random pick one that is always available.
+	// random pick one that is always available. Use global rand which is thread-safe since Go 1.20.
 	var serverKey string
 	var server SRSServer
 	for i := 0; i < 3; i++ {