rkKVM/stream/video.go

package stream

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"rkkvm/config"
	"rkkvm/external/ffmpeg"
	"rkkvm/http/reqrsp"
	"rkkvm/http/rtc"
	"strconv"

	"github.com/gin-gonic/gin"
	"github.com/pion/webrtc/v4"
	log "github.com/sirupsen/logrus"
)

/*
{
  "ok": true,
  "result": {
    "instance_id": "",
    "encoder": {
      "type": "CPU",
      "quality": 80
    },
    "source": {
      "resolution": {
        "width": 1920,
        "height": 1080
      },
      "online": true,
      "desired_fps": 30,
      "captured_fps": 50
    },
    "stream": {
      "queued_fps": 26,
      "clients": 1,
      "clients_stat": {
        "2d90d210e837a19c": {
          "fps": 26,
          "extra_headers": false,
          "advance_headers": false,
          "dual_final_frames": false,
          "zero_data": false,
          "key": "0"
        }
      }
    }
  }
}
*/

type State struct {
	Ok     bool `json:"ok"`
	Result struct {
		InstanceID string `json:"instance_id"`
		Encoder    struct {
			Type    string `json:"type"`
			Quality int    `json:"quality"`
		} `json:"encoder"`
		Source struct {
			Resolution struct {
				Width  int `json:"width"`
				Height int `json:"height"`
			} `json:"resolution"`
			Online      bool `json:"online"`
			DesiredFPS  int  `json:"desired_fps"`
			CapturedFPS int  `json:"captured_fps"`
		} `json:"source"`
		Stream struct {
			QueuedFPS   int `json:"queued_fps"`
			Clients     int `json:"clients"`
			ClientsStat map[string]struct {
				FPS             int    `json:"fps"`
				ExtraHeaders    bool   `json:"extra_headers"`
				AdvanceHeaders  bool   `json:"advance_headers"`
				DualFinalFrames bool   `json:"dual_final_frames"`
				ZeroData        bool   `json:"zero_data"`
				Key             string `json:"key"`
			}
		} `json:"stream"`
	} `json:"result"`
}

func FormatHttpUrl() string {
	cfg := config.Get()
	return fmt.Sprintf("http://%s:%d", cfg.UStreamer.Host, cfg.UStreamer.Port)
}

func GetState() (State, error) {
	url := FormatHttpUrl() + "/state"

	rsp, err := http.Get(url)
	if err != nil {
		log.Errorf("Failed to get state: %v", err)
		return State{}, err
	}
	defer rsp.Body.Close()

	if rsp.StatusCode != http.StatusOK {
		log.Errorf("Invalid status code: %d", rsp.StatusCode)
		return State{}, err
	}

	var state State
	if err := json.NewDecoder(rsp.Body).Decode(&state); err != nil {
		log.Errorf("Failed to decode state: %v", err)
		return State{}, err
	}

	return state, nil
}

func MjpegHandler(c *gin.Context) {
	url := FormatHttpUrl() + "/stream"

	rsp, err := http.Get(url)
	if err != nil {
		c.String(http.StatusInternalServerError, "Failed to get stream")
		log.Errorf("Failed to get stream: %v", err)
		return
	}
	defer rsp.Body.Close()

	if rsp.StatusCode != http.StatusOK {
		c.String(rsp.StatusCode, "Invalid status code")
		return
	}

	c.Header("Content-Type", rsp.Header.Get("Content-Type"))
	buf := make([]byte, 1024*1024)
	_, err = io.CopyBuffer(c.Writer, rsp.Body, buf)
	if err != nil {
		c.String(http.StatusInternalServerError, "Failed to passthrough stream")
		log.Errorf("Failed to passthrough stream: %v", err)
		return
	}
}

func WebRTCPeerConnect(c *gin.Context) {
	var offer webrtc.SessionDescription
	if err := c.ShouldBindJSON(&offer); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": "invalid offer"})
		return
	}

	peer, err := rtc.NewPeer()
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to create peer connection"})
		return
	}

	answer, err := rtc.Get().AddPeer(peer, offer)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to add peer"})
		return
	}

	c.JSON(http.StatusOK, reqrsp.NanoKVMRsp{
		Msg:  reqrsp.MsgSuccess,
		Data: answer,
	})
}

func WebRTCSettings(c *gin.Context) {
	bitrateStr := c.Query("bitrate")
	fpsStr := c.Query("fps")

	if fpsStr == "" && bitrateStr == "" {
		c.String(http.StatusBadGateway, "please specify bitrate or fps parameters")
		return
	}

	ffmpeg := ffmpeg.GetFFmpeg()

	if len(bitrateStr) > 0 {
		bitrate, err := strconv.Atoi(bitrateStr)
		if err != nil {
			c.String(http.StatusBadRequest, "wrong bitrate value")
			return
		}
		ffmpeg.SetBitrate(bitrate)
	}

	if len(fpsStr) > 0 {
		fps, err := strconv.Atoi(fpsStr)
		if err != nil {
			c.String(http.StatusBadRequest, "wrong fps value")
			return
		}
		ffmpeg.SetFPS(fps)
	}

	ffmpeg.Stop()
	ffmpeg.Start()
}