Webrtc推视频到html上播放

package main

import (
    "context"
    "encoding/base64"
    "encoding/json"
    "errors"
    "io"
    "log"
    "net/http"
    "os"
    "sync"
    "time"

    "github.com/gin-gonic/gin"
    "github.com/gorilla/websocket"
    "github.com/pion/webrtc/v4"
    "github.com/pion/webrtc/v4/pkg/media"
    "github.com/pion/webrtc/v4/pkg/media/ivfreader"
)

const (
    videoFileName = "output.ivf"
)

var (
    upgrader = websocket.Upgrader{
        CheckOrigin: func(r *http.Request) bool {
            return true
        },
    }
    clients     = make(map[*websocket.Conn]bool)
    clientsLock sync.Mutex
)

func main() {
    // 检查视频文件是否存在
    _, err := os.Stat(videoFileName)
    haveVideoFile := !os.IsNotExist(err)

    if !haveVideoFile {
        log.Fatalf("找不到视频文件 `%s`", videoFileName)
    }

    // 创建Gin路由
    r := gin.Default()

    // 静态文件服务
    r.Static("/static", "./static")

    // WebSocket路由
    r.GET("/ws", func(c *gin.Context) {
        wsHandler(c.Writer, c.Request)
    })

    // 首页路由
    r.GET("/", func(c *gin.Context) {
        c.File("./static/index.html")
    })

    // 启动服务器
    log.Println("服务器已启动，访问 http://localhost:8080")
    if err := r.Run(":8080"); err != nil {
        log.Fatal("服务器启动失败: ", err)
    }
}

func wsHandler(w http.ResponseWriter, r *http.Request) {
    conn, err := upgrader.Upgrade(w, r, nil)
    if err != nil {
        log.Println("WebSocket升级失败:", err)
        return
    }
    defer conn.Close()

    clientsLock.Lock()
    clients[conn] = true
    clientsLock.Unlock()

    for {
        _, message, err := conn.ReadMessage()
        if err != nil {
            log.Println("读取消息失败:", err)
            clientsLock.Lock()
            delete(clients, conn)
            clientsLock.Unlock()
            break
        }

        // 处理客户端发来的SDP offer
        offer := webrtc.SessionDescription{}
        if err := decode(string(message), &offer); err != nil {
            log.Println("解码SDP失败:", err)
            continue
        }

        // 创建新的PeerConnection
        peerConnection, err := webrtc.NewPeerConnection(webrtc.Configuration{
            ICEServers: []webrtc.ICEServer{
                {URLs: []string{"stun:stun.l.google.com:19302"}},
            },
        })
        if err != nil {
            log.Println("创建PeerConnection失败:", err)
            continue
        }
        defer peerConnection.Close()

        // 创建上下文并确保在任何情况下都能调用 Cancel
        iceConnectedCtx, iceConnectedCtxCancel := context.WithCancel(context.Background())
        defer iceConnectedCtxCancel() // 确保在函数结束时调用 Cancel

        // 打开视频文件
        file, err := os.Open(videoFileName)
        if err != nil {
            log.Println("打开视频文件失败:", err)
            continue
        }
        defer file.Close()

        ivf, header, err := ivfreader.NewWith(file)
        if err != nil {
            log.Println("创建IVF reader失败:", err)
            continue
        }

        // 创建视频轨道
        videoTrack, err := webrtc.NewTrackLocalStaticSample(webrtc.RTPCodecCapability{
            MimeType: webrtc.MimeTypeVP8,
        }, "video", "pion")
        if err != nil {
            log.Println("创建视频轨道失败:", err)
            continue
        }

        rtpSender, err := peerConnection.AddTrack(videoTrack)
        if err != nil {
            log.Println("添加视频轨道失败:", err)
            continue
        }

        // 读取RTCP包
        go func() {
            rtcpBuf := make([]byte, 1500)
            for {
                if _, _, rtcpErr := rtpSender.Read(rtcpBuf); rtcpErr != nil {
                    return
                }
            }
        }()

        // 发送视频帧
        go func() {
            <-iceConnectedCtx.Done()

            ticker := time.NewTicker(time.Millisecond * time.Duration((float32(header.TimebaseNumerator)/float32(header.TimebaseDenominator))*1000))
            defer ticker.Stop()
            for ; true; <-ticker.C {
                frame, _, err := ivf.ParseNextFrame()
                if errors.Is(err, io.EOF) {
                    log.Println("所有视频帧已发送")
                    return
                }

                if err != nil {
                    log.Println("解析视频帧失败:", err)
                    return
                }

                if err = videoTrack.WriteSample(media.Sample{
                    Data:     frame,
                    Duration: time.Second,
                }); err != nil {
                    log.Println("写入视频样本失败:", err)
                    return
                }
            }
        }()

        // ICE连接状态变化处理
        peerConnection.OnICEConnectionStateChange(func(connectionState webrtc.ICEConnectionState) {
            log.Printf("ICE连接状态已更改: %s\n", connectionState.String())
            if connectionState == webrtc.ICEConnectionStateConnected {
                iceConnectedCtxCancel()
            } else if connectionState == webrtc.ICEConnectionStateFailed || connectionState == webrtc.ICEConnectionStateDisconnected {
                peerConnection.Close()
            }
        })

        // 设置远程描述
        if err = peerConnection.SetRemoteDescription(offer); err != nil {
            log.Println("设置远程描述失败:", err)
            continue
        }

        // 创建SDP answer
        answer, err := peerConnection.CreateAnswer(nil)
        if err != nil {
            log.Println("创建answer失败:", err)
            continue
        }

        // 设置本地描述前添加ICE候选收集处理
        peerConnection.OnICECandidate(func(candidate *webrtc.ICECandidate) {
            if candidate == nil {
                return
            }
            log.Printf("收集到ICE候选: %s\n", candidate.ToJSON().Candidate)
        })

        // 设置本地描述
        if err = peerConnection.SetLocalDescription(answer); err != nil {
            log.Println("设置本地描述失败:", err)
            continue
        }

        // 等待ICE收集完成
        gatherComplete := webrtc.GatheringCompletePromise(peerConnection)
        <-gatherComplete

        // 发送answer给客户端
        if err := conn.WriteMessage(websocket.TextMessage, []byte(encode(peerConnection.LocalDescription()))); err != nil {
            log.Println("发送answer失败:", err)
            continue
        }
    }
}

func encode(obj *webrtc.SessionDescription) string {
    b, err := json.Marshal(obj)
    if err != nil {
        log.Panic("编码失败:", err)
    }
    return base64.StdEncoding.EncodeToString(b)
}

func decode(in string, obj *webrtc.SessionDescription) error {
    b, err := base64.StdEncoding.DecodeString(in)
    if err != nil {
        return err
    }
    return json.Unmarshal(b, obj)
}