Phase 2 completed: Stream Server Auth, Database and Business API

2026-03-16 15:54:41 +08:00
commit c84dd6ea36
16 changed files with 757 additions and 0 deletions
--- a/internal/api/auth.go
+++ b/internal/api/auth.go
@@ -0,0 +1,96 @@
+package api
+
+import (
+	"net/http"
+
+	"github.com/gin-gonic/gin"
+
+	"hightube/internal/db"
+	"hightube/internal/model"
+	"hightube/internal/utils"
+)
+
+type RegisterRequest struct {
+	Username string `json:"username" binding:"required"`
+	Password string `json:"password" binding:"required"`
+}
+
+type LoginRequest struct {
+	Username string `json:"username" binding:"required"`
+	Password string `json:"password" binding:"required"`
+}
+
+func Register(c *gin.Context) {
+	var req RegisterRequest
+	if err := c.ShouldBindJSON(&req); err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
+		return
+	}
+
+	// Check if user exists
+	var existingUser model.User
+	if err := db.DB.Where("username = ?", req.Username).First(&existingUser).Error; err == nil {
+		c.JSON(http.StatusConflict, gin.H{"error": "Username already exists"})
+		return
+	}
+
+	// Hash password
+	hashedPassword, err := utils.HashPassword(req.Password)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to hash password"})
+		return
+	}
+
+	// Create user
+	user := model.User{
+		Username: req.Username,
+		Password: hashedPassword,
+	}
+	if err := db.DB.Create(&user).Error; err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to create user"})
+		return
+	}
+
+	// Create a default live room for the new user
+	room := model.Room{
+		UserID:    user.ID,
+		Title:     user.Username + "'s Live Room",
+		StreamKey: utils.GenerateStreamKey(),
+		IsActive:  false,
+	}
+	if err := db.DB.Create(&room).Error; err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to create room for user"})
+		return
+	}
+
+	c.JSON(http.StatusCreated, gin.H{"message": "User registered successfully", "user_id": user.ID})
+}
+
+func Login(c *gin.Context) {
+	var req LoginRequest
+	if err := c.ShouldBindJSON(&req); err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
+		return
+	}
+
+	var user model.User
+	if err := db.DB.Where("username = ?", req.Username).First(&user).Error; err != nil {
+		c.JSON(http.StatusUnauthorized, gin.H{"error": "Invalid username or password"})
+		return
+	}
+
+	if !utils.CheckPasswordHash(req.Password, user.Password) {
+		c.JSON(http.StatusUnauthorized, gin.H{"error": "Invalid username or password"})
+		return
+	}
+
+	token, err := utils.GenerateToken(user.ID)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to generate token"})
+		return
+	}
+
+	c.JSON(http.StatusOK, gin.H{
+		"token": token,
+	})
+}
--- a/internal/api/middleware.go
+++ b/internal/api/middleware.go
@@ -0,0 +1,42 @@
+package api
+
+import (
+	"net/http"
+	"strconv"
+	"strings"
+
+	"github.com/gin-gonic/gin"
+
+	"hightube/internal/utils"
+)
+
+// AuthMiddleware intercepts requests, validates JWT, and injects user_id into context
+func AuthMiddleware() gin.HandlerFunc {
+	return func(c *gin.Context) {
+		authHeader := c.GetHeader("Authorization")
+		if authHeader == "" {
+			c.JSON(http.StatusUnauthorized, gin.H{"error": "Authorization header is required"})
+			c.Abort()
+			return
+		}
+
+		parts := strings.Split(authHeader, " ")
+		if len(parts) != 2 || parts[0] != "Bearer" {
+			c.JSON(http.StatusUnauthorized, gin.H{"error": "Authorization header format must be Bearer {token}"})
+			c.Abort()
+			return
+		}
+
+		tokenStr := parts[1]
+		userIDStr, err := utils.ParseToken(tokenStr)
+		if err != nil {
+			c.JSON(http.StatusUnauthorized, gin.H{"error": "Invalid or expired token"})
+			c.Abort()
+			return
+		}
+
+		userID, _ := strconv.ParseUint(userIDStr, 10, 32)
+		c.Set("user_id", uint(userID))
+		c.Next()
+	}
+}
--- a/internal/api/room.go
+++ b/internal/api/room.go
@@ -0,0 +1,50 @@
+package api
+
+import (
+	"net/http"
+
+	"github.com/gin-gonic/gin"
+
+	"hightube/internal/db"
+	"hightube/internal/model"
+)
+
+// GetMyRoom returns the room details for the currently authenticated user
+func GetMyRoom(c *gin.Context) {
+	userID, _ := c.Get("user_id")
+
+	var room model.Room
+	if err := db.DB.Where("user_id = ?", userID).First(&room).Error; err != nil {
+		c.JSON(http.StatusNotFound, gin.H{"error": "Room not found"})
+		return
+	}
+
+	c.JSON(http.StatusOK, gin.H{
+		"room_id":    room.ID,
+		"title":      room.Title,
+		"stream_key": room.StreamKey,
+		"is_active":  room.IsActive,
+	})
+}
+
+// GetActiveRooms returns a list of all currently active live rooms
+func GetActiveRooms(c *gin.Context) {
+	var rooms []model.Room
+	// Fetch rooms where is_active is true
+	if err := db.DB.Where("is_active = ?", true).Find(&rooms).Error; err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to fetch active rooms"})
+		return
+	}
+
+	// Return safe information (do not leak stream keys)
+	var result []map[string]interface{}
+	for _, r := range rooms {
+		result = append(result, map[string]interface{}{
+			"room_id": r.ID,
+			"title":   r.Title,
+			"user_id": r.UserID,
+		})
+	}
+
+	c.JSON(http.StatusOK, gin.H{"active_rooms": result})
+}
--- a/internal/api/router.go
+++ b/internal/api/router.go
@@ -0,0 +1,30 @@
+package api
+
+import (
+	"github.com/gin-gonic/gin"
+)
+
+// SetupRouter configures the Gin router and defines API endpoints
+func SetupRouter() *gin.Engine {
+	// 设置为发布模式，消除 "[WARNING] Running in debug mode" 警告
+	gin.SetMode(gin.ReleaseMode)
+	
+	r := gin.Default()
+
+	// 清除代理信任警告 "[WARNING] You trusted all proxies"
+	r.SetTrustedProxies(nil)
+
+	// Public routes
+	r.POST("/api/register", Register)
+	r.POST("/api/login", Login)
+	r.GET("/api/rooms/active", GetActiveRooms)
+
+	// Protected routes (require JWT)
+	authGroup := r.Group("/api")
+	authGroup.Use(AuthMiddleware())
+	{
+		authGroup.GET("/room/my", GetMyRoom)
+	}
+
+	return r
+}
--- a/internal/db/db.go
+++ b/internal/db/db.go
@@ -0,0 +1,45 @@
+package db
+
+import (
+	"log"
+	"os"
+	"time"
+
+	"gorm.io/driver/sqlite"
+	"gorm.io/gorm"
+	"gorm.io/gorm/logger"
+
+	"hightube/internal/model"
+)
+
+var DB *gorm.DB
+
+// InitDB initializes the SQLite database connection and auto-migrates models.
+func InitDB() {
+	newLogger := logger.New(
+		log.New(os.Stdout, "\r\n", log.LstdFlags), // io writer
+		logger.Config{
+			SlowThreshold:             time.Second,   // Slow SQL threshold
+			LogLevel:                  logger.Warn,   // Log level
+			IgnoreRecordNotFoundError: true,          // Ignore ErrRecordNotFound error for logger
+			Colorful:                  true,          // Disable color
+		},
+	)
+
+	var err error
+	// Use SQLite database stored in a local file named "hightube.db"
+	DB, err = gorm.Open(sqlite.Open("hightube.db"), &gorm.Config{
+		Logger: newLogger,
+	})
+	if err != nil {
+		log.Fatalf("Failed to connect database: %v", err)
+	}
+
+	// Auto-migrate the schema
+	err = DB.AutoMigrate(&model.User{}, &model.Room{})
+	if err != nil {
+		log.Fatalf("Failed to migrate database: %v", err)
+	}
+
+	log.Println("Database initialized successfully.")
+}
--- a/internal/model/room.go
+++ b/internal/model/room.go
@@ -0,0 +1,14 @@
+package model
+
+import (
+	"gorm.io/gorm"
+)
+
+// Room represents a user's personal live streaming room.
+type Room struct {
+	gorm.Model
+	UserID    uint   `gorm:"uniqueIndex;not null"`
+	Title     string `gorm:"default:'My Live Room'"`
+	StreamKey string `gorm:"uniqueIndex;not null"` // Secret key for OBS streaming
+	IsActive  bool   `gorm:"default:false"`        // Whether the stream is currently active
+}
--- a/internal/model/user.go
+++ b/internal/model/user.go
@@ -0,0 +1,12 @@
+package model
+
+import (
+	"gorm.io/gorm"
+)
+
+// User represents a registered user in the system.
+type User struct {
+	gorm.Model
+	Username string `gorm:"uniqueIndex;not null"`
+	Password string `gorm:"not null"` // Hashed password
+}
--- a/internal/stream/server.go
+++ b/internal/stream/server.go
@@ -0,0 +1,136 @@
+package stream
+
+import (
+	"fmt"
+	"io"
+	"strings"
+	"sync"
+
+	"github.com/nareix/joy4/av/avutil"
+	"github.com/nareix/joy4/av/pubsub"
+	"github.com/nareix/joy4/format"
+	"github.com/nareix/joy4/format/rtmp"
+
+	"hightube/internal/db"
+	"hightube/internal/model"
+)
+
+func init() {
+	// Register all supported audio/video formats
+	format.RegisterAll()
+}
+
+// RTMPServer manages all active live streams
+type RTMPServer struct {
+	server   *rtmp.Server
+	channels map[string]*pubsub.Queue
+	mutex    sync.RWMutex
+}
+
+// NewRTMPServer creates and initializes a new media server
+func NewRTMPServer() *RTMPServer {
+	s := &RTMPServer{
+		channels: make(map[string]*pubsub.Queue),
+		server:   &rtmp.Server{},
+	}
+
+	// Triggered when a broadcaster (e.g., OBS) starts publishing
+	s.server.HandlePublish = func(conn *rtmp.Conn) {
+		streamPath := conn.URL.Path // Expected format: /live/{stream_key}
+		fmt.Printf("[INFO] OBS is attempting to publish to: %s\n", streamPath)
+
+		// Extract stream key from path
+		parts := strings.Split(streamPath, "/")
+		if len(parts) < 3 || parts[1] != "live" {
+			fmt.Printf("[WARN] Invalid publish path format: %s\n", streamPath)
+			return
+		}
+		streamKey := parts[2]
+
+		// Authenticate stream key
+		var room model.Room
+		if err := db.DB.Where("stream_key = ?", streamKey).First(&room).Error; err != nil {
+			fmt.Printf("[WARN] Authentication failed, invalid stream key: %s\n", streamKey)
+			return // Reject connection
+		}
+
+		fmt.Printf("[INFO] Stream authenticated for Room ID: %d\n", room.ID)
+
+		// 1. Get audio/video stream metadata
+		streams, err := conn.Streams()
+		if err != nil {
+			fmt.Printf("[ERROR] Failed to parse stream headers: %v\n", err)
+			return
+		}
+
+		// 2. Map the active stream by Room ID so viewers can use /live/{room_id}
+		roomLivePath := fmt.Sprintf("/live/%d", room.ID)
+		
+		s.mutex.Lock()
+		q := pubsub.NewQueue()
+		q.WriteHeader(streams)
+		s.channels[roomLivePath] = q
+		s.mutex.Unlock()
+
+		// Mark room as active in DB
+		db.DB.Model(&room).Update("is_active", true)
+
+		// 3. Cleanup on end
+		defer func() {
+			s.mutex.Lock()
+			delete(s.channels, roomLivePath)
+			s.mutex.Unlock()
+			q.Close()
+			db.DB.Model(&room).Update("is_active", false) // Mark room as inactive
+			fmt.Printf("[INFO] Publishing ended for Room ID: %d\n", room.ID)
+		}()
+
+		// 4. Continuously copy data packets to our broadcast queue
+		avutil.CopyPackets(q, conn)
+	}
+
+	// Triggered when a viewer (e.g., VLC) requests playback
+	s.server.HandlePlay = func(conn *rtmp.Conn) {
+		streamPath := conn.URL.Path // Expected format: /live/{room_id}
+		fmt.Printf("[INFO] VLC is pulling stream from: %s\n", streamPath)
+
+		// 1. Look for the requested room's data queue
+		s.mutex.RLock()
+		q, ok := s.channels[streamPath]
+		s.mutex.RUnlock()
+
+		if !ok {
+			fmt.Printf("[WARN] Stream not found or inactive: %s\n", streamPath)
+			return
+		}
+
+		// 2. Get the cursor from the latest position and notify client of stream format
+		cursor := q.Latest()
+		streams, _ := cursor.Streams()
+		conn.WriteHeader(streams)
+
+		// 3. Cleanup on end
+		defer fmt.Printf("[INFO] Playback ended: %s\n", streamPath)
+
+		// 4. Continuously copy data packets to the viewer
+		err := avutil.CopyPackets(conn, cursor)
+		if err != nil && err != io.EOF {
+			// 如果是客户端主动断开连接引起的错误，不将其作为严重错误打印
+			errStr := err.Error()
+			if strings.Contains(errStr, "broken pipe") || strings.Contains(errStr, "connection reset by peer") {
+				fmt.Printf("[INFO] Viewer disconnected normally: %s\n", streamPath)
+			} else {
+				fmt.Printf("[ERROR] Error occurred during playback: %v\n", err)
+			}
+		}
+	}
+
+	return s
+}
+
+// Start launches the RTMP server
+func (s *RTMPServer) Start(addr string) error {
+	s.server.Addr = addr
+	fmt.Printf("[INFO] RTMP Server is listening on %s...\n", addr)
+	return s.server.ListenAndServe()
+}
--- a/internal/utils/utils.go
+++ b/internal/utils/utils.go
@@ -0,0 +1,55 @@
+package utils
+
+import (
+	"crypto/rand"
+	"encoding/hex"
+	"fmt"
+	"time"
+
+	"github.com/golang-jwt/jwt/v5"
+	"golang.org/x/crypto/bcrypt"
+)
+
+// In production, load this from environment variables
+var jwtKey = []byte("hightube_super_secret_key_MVP_only")
+
+// GenerateToken generates a JWT token for a given user ID
+func GenerateToken(userID uint) (string, error) {
+	claims := &jwt.RegisteredClaims{
+		Subject:   fmt.Sprintf("%d", userID),
+		ExpiresAt: jwt.NewNumericDate(time.Now().Add(24 * time.Hour)),
+	}
+	token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
+	return token.SignedString(jwtKey)
+}
+
+// ParseToken parses the JWT string and returns the user ID (Subject)
+func ParseToken(tokenStr string) (string, error) {
+	claims := &jwt.RegisteredClaims{}
+	token, err := jwt.ParseWithClaims(tokenStr, claims, func(t *jwt.Token) (interface{}, error) {
+		return jwtKey, nil
+	})
+	if err != nil || !token.Valid {
+		return "", err
+	}
+	return claims.Subject, nil
+}
+
+// HashPassword creates a bcrypt hash of the password
+func HashPassword(password string) (string, error) {
+	bytes, err := bcrypt.GenerateFromPassword([]byte(password), 14)
+	return string(bytes), err
+}
+
+// CheckPasswordHash compares a bcrypt hashed password with its possible plaintext equivalent
+func CheckPasswordHash(password, hash string) bool {
+	err := bcrypt.CompareHashAndPassword([]byte(hash), []byte(password))
+	return err == nil
+}
+
+// GenerateStreamKey generates a random string to be used as a stream key
+func GenerateStreamKey() string {
+	bytes := make([]byte, 16)
+	rand.Read(bytes)
+	return hex.EncodeToString(bytes)
+}