monibuca/plugin/cascade/server.go

package plugin_cascade

import (
	"bufio"
	"database/sql"
	"net/http"
	"strconv"
	"strings"
	"sync"

	"google.golang.org/protobuf/types/known/timestamppb"
	"m7s.live/v5"
	"m7s.live/v5/pkg"
	"m7s.live/v5/pkg/task"
	"m7s.live/v5/pkg/util"

	"context"

	"github.com/quic-go/quic-go"
	"m7s.live/v5/plugin/cascade/pb"
	cascade "m7s.live/v5/plugin/cascade/pkg"
)

type CascadeServerPlugin struct {
	m7s.Plugin
	pb.UnimplementedServerServer
	AutoRegister bool                   `default:"true" desc:"下级自动注册"`
	RelayAPI     cascade.RelayAPIConfig `desc:"访问控制"`
	clients      util.Collection[uint, *cascade.Instance]
}

func (c *CascadeServerPlugin) Start() (err error) {
	if c.GetCommonConf().Quic.ListenAddr == "" {
		return pkg.ErrNotListen
	}
	c.clients.L = &sync.RWMutex{}
	if c.DB == nil {
		return pkg.ErrNoDB
	}
	c.DB.AutoMigrate(&cascade.Instance{})
	var instance []*cascade.Instance
	c.DB.Find(&instance)
	for _, instance := range instance {
		c.clients.Add(instance)
		if instance.Online {
			instance.Online = false
			c.DB.Save(instance)
		}
	}
	return
}

var _ = m7s.InstallPlugin[CascadeServerPlugin](m7s.PluginMeta{
	DefaultYaml: `quic:
  listenaddr: :44944`,
	ServiceDesc:         &pb.Server_ServiceDesc,
	RegisterGRPCHandler: pb.RegisterServerHandler,
})

type CascadeServer struct {
	task.Work
	quic.Connection
	conf   *CascadeServerPlugin
	client *cascade.Instance
}

func (c *CascadeServerPlugin) OnQUICConnect(conn quic.Connection) task.ITask {
	ret := &CascadeServer{
		Connection: conn,
		conf:       c,
	}
	ret.Logger = c.Logger.With("remoteAddr", conn.RemoteAddr().String())
	return ret
}

func (task *CascadeServer) Go() (err error) {
	remoteAddr := task.Connection.RemoteAddr().String()
	var stream quic.Stream
	if stream, err = task.AcceptStream(task); err != nil {
		task.Error("AcceptStream", "err", err)
		return
	}
	var secret string
	r := bufio.NewReader(stream)
	if secret, err = r.ReadString(0); err != nil {
		task.Error("read secret", "err", err)
		return
	}
	secret = secret[:len(secret)-1] // 去掉msg末尾的0
	child := &cascade.Instance{}
	if secret != "" {
		tx := task.conf.DB.First(child, "secret = ?", secret)
		err = tx.Error
	} else {
		tx := task.conf.DB.First(child, "ip = ?", remoteAddr)
		err = tx.Error
	}
	if err == nil {
		task.conf.clients.Set(child)
	} else if task.conf.AutoRegister {
		child.Secret = sql.NullString{String: secret, Valid: secret != ""}
		child.IP = remoteAddr
		err = task.conf.DB.First(child, "ip = ?", remoteAddr).Error
		if err != nil {
			err = task.conf.DB.Create(child).Error
		}
		task.conf.clients.Set(child)
	} else {
		task.Error("invalid secret:", "secret", secret)
		_, err = stream.Write([]byte{1, 0})
		return
	}
	child.IP = remoteAddr
	child.Online = true
	if child.Name == "" {
		child.Name = remoteAddr
	}
	err = task.conf.DB.Updates(child).Error
	child.Connection = task.Connection
	task.client = child
	_, err = stream.Write([]byte{0, 0})
	err = stream.Close()
	task.Info("client register:", "remoteAddr", remoteAddr)
	for err == nil {
		var receiveRequestTask cascade.ReceiveRequestTask
		receiveRequestTask.Connection = task.Connection
		receiveRequestTask.Plugin = &task.conf.Plugin
		receiveRequestTask.Handler = task.conf.GetGlobalCommonConf().GetHandler(task.Logger)
		if receiveRequestTask.Stream, err = task.AcceptStream(task); err == nil {
			task.AddTask(&receiveRequestTask)
		}
	}
	return
}

func (task *CascadeServer) Dispose() {
	if task.Connection != nil {
		task.Connection.CloseWithError(quic.ApplicationErrorCode(0), task.StopReason().Error())
	}
	if task.client != nil {
		task.client.Online = false
		task.conf.DB.Save(task.client)
	}
}

// API_relay_ 用于转发请求, api/relay/:instanceId/*
func (c *CascadeServerPlugin) API_relay_(w http.ResponseWriter, r *http.Request) {
	paths := strings.Split(r.URL.Path, "/")
	instanceId, err := strconv.ParseUint(paths[3], 10, 32)
	instance, ok := c.clients.Get(uint(instanceId))
	if err != nil || !ok {
		http.Error(w, "instance not found", http.StatusNotFound)
		return
	}
	if !instance.Online {
		http.Error(w, "instance offline", http.StatusServiceUnavailable)
		return
	}
	relayURL := "/" + strings.Join(paths[4:], "/")
	r.URL.Path = relayURL
	if r.URL.RawQuery != "" {
		relayURL += "?" + r.URL.RawQuery
	}
	c.Debug("relayQuic", "relayURL", relayURL)
	var relayer cascade.Http2Quic
	relayer.Stream, err = instance.OpenStreamSync(r.Context())
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	relayer.ServeHTTP(w, r)
}

func (c *CascadeServerPlugin) GetClientList(ctx context.Context, req *pb.GetClientListRequest) (clientList *pb.GetClientListResponse, err error) {
	clientList = &pb.GetClientListResponse{}

	for instance := range c.clients.Range {
		client := &pb.CascadeClient{
			Id:          uint32(instance.ID),
			Name:        instance.Name,
			Ip:          instance.IP,
			Online:      instance.Online,
			CreatedTime: timestamppb.New(instance.CreatedAt),
			UpdatedTime: timestamppb.New(instance.UpdatedAt),
		}
		clientList.Data = append(clientList.Data, client)
	}

	return
}

func (c *CascadeServerPlugin) CreateClient(ctx context.Context, req *pb.CreateClientRequest) (result *pb.CreateClientResponse, err error) {
	result = &pb.CreateClientResponse{}

	instance := &cascade.Instance{
		Name:   req.Name,
		Secret: sql.NullString{String: req.Secret, Valid: req.Secret != ""},
	}

	if err = c.DB.Create(instance).Error; err != nil {
		return
	}

	c.clients.Set(instance)

	result.Data = &pb.CascadeClient{
		Id:          uint32(instance.ID),
		Name:        instance.Name,
		CreatedTime: timestamppb.New(instance.CreatedAt),
		UpdatedTime: timestamppb.New(instance.UpdatedAt),
	}

	return
}

func (c *CascadeServerPlugin) UpdateClient(ctx context.Context, req *pb.UpdateClientRequest) (result *pb.UpdateClientResponse, err error) {
	result = &pb.UpdateClientResponse{}

	instance := &cascade.Instance{}
	if err = c.DB.First(instance, req.Id).Error; err != nil {
		return
	}

	instance.Name = req.Name
	instance.Secret = sql.NullString{String: req.Secret, Valid: req.Secret != ""}

	if err = c.DB.Save(instance).Error; err != nil {
		return
	}

	c.clients.Set(instance)

	result.Data = &pb.CascadeClient{
		Id:          uint32(instance.ID),
		Name:        instance.Name,
		UpdatedTime: timestamppb.New(instance.UpdatedAt),
	}

	return
}

func (c *CascadeServerPlugin) DeleteClient(ctx context.Context, req *pb.DeleteClientRequest) (result *pb.DeleteClientResponse, err error) {
	result = &pb.DeleteClientResponse{}

	instance := &cascade.Instance{}
	if err = c.DB.First(instance, req.Id).Error; err != nil {
		return
	}

	if err = c.DB.Delete(instance).Error; err != nil {
		return
	}

	c.clients.RemoveByKey(instance.ID)
	return
}