sing-box-extended/service/node_manager/server/service.go

package server

import (
	"context"
	"errors"
	"sync"

	"github.com/sagernet/sing-box/adapter"
	boxService "github.com/sagernet/sing-box/adapter/service"
	"github.com/sagernet/sing-box/common/listener"
	"github.com/sagernet/sing-box/common/tls"
	C "github.com/sagernet/sing-box/constant"
	"github.com/sagernet/sing-box/log"
	"github.com/sagernet/sing-box/option"
	CM "github.com/sagernet/sing-box/service/manager/constant"
	pb "github.com/sagernet/sing-box/service/node_manager/manager"
	"github.com/sagernet/sing/common"
	E "github.com/sagernet/sing/common/exceptions"
	N "github.com/sagernet/sing/common/network"
	aTLS "github.com/sagernet/sing/common/tls"
	"github.com/sagernet/sing/service"
	"golang.org/x/net/http2"
	"google.golang.org/grpc"
)

func RegisterService(registry *boxService.Registry) {
	boxService.Register[option.NodeManagerServerServiceOptions](registry, C.TypeNodeManagerServer, NewService)
}

type Service struct {
	pb.UnimplementedManagerServer
	boxService.Adapter

	ctx        context.Context
	logger     log.ContextLogger
	listener   *listener.Listener
	tlsConfig  tls.ServerConfig
	grpcServer *grpc.Server
	manager    CM.Manager
	options    option.NodeManagerServerServiceOptions

	mtx sync.Mutex
}

func NewService(ctx context.Context, logger log.ContextLogger, tag string, options option.NodeManagerServerServiceOptions) (adapter.Service, error) {
	return &Service{
		Adapter: boxService.NewAdapter(C.TypeManager, tag),
		ctx:     ctx,
		logger:  logger,
		listener: listener.New(listener.Options{
			Context: ctx,
			Logger:  logger,
			Network: []string{N.NetworkTCP},
			Listen:  options.ListenOptions,
		}),
		options: options,
	}, nil
}

func (s *Service) AddNode(node *pb.Node, stream grpc.ServerStreamingServer[pb.NodeData]) error {
	remoteNode, errChan := NewRemoteNode(s.ctx, s.logger, stream)
	err := s.manager.AddNode(node.Uuid, remoteNode)
	if err != nil {
		if err == CM.ErrNotFound {
			return err
		} else {
			s.logger.Error(err)
			return E.New("internal error")
		}
	}
	return <-errChan
}

func (s *Service) AcquireLock(ctx context.Context, request *pb.AcquireLockRequest) (*pb.LockData, error) {
	handleId, err := s.manager.AcquireLock(int(request.LimiterId), request.Id)
	if err != nil {
		return nil, err
	}
	return &pb.LockData{HandleId: handleId}, nil
}

func (s *Service) RefreshLock(ctx context.Context, data *pb.LockData) (*pb.Empty, error) {
	return nil, s.manager.RefreshLock(int(data.LimiterId), data.Id, data.HandleId)
}

func (s *Service) ReleaseLock(ctx context.Context, data *pb.LockData) (*pb.Empty, error) {
	return nil, s.manager.ReleaseLock(int(data.LimiterId), data.Id, data.HandleId)
}

func (s *Service) Start(stage adapter.StartStage) error {
	if stage != adapter.StartStateStart {
		return nil
	}
	boxManager := service.FromContext[adapter.ServiceManager](s.ctx)
	service, ok := boxManager.Get(s.options.Manager)
	if !ok {
		return E.New("manager ", s.options.Manager, " not found")
	}
	s.manager, ok = service.(CM.Manager)
	if !ok {
		return E.New("invalid", s.options.Manager, " manager")
	}
	if s.options.TLS != nil {
		tlsConfig, err := tls.NewServer(s.ctx, s.logger, common.PtrValueOrDefault(s.options.TLS))
		if err != nil {
			return err
		}
		s.tlsConfig = tlsConfig
	}
	if s.tlsConfig != nil {
		err := s.tlsConfig.Start()
		if err != nil {
			return E.Cause(err, "create TLS config")
		}
	}
	tcpListener, err := s.listener.ListenTCP()
	if err != nil {
		return err
	}
	if s.tlsConfig != nil {
		if !common.Contains(s.tlsConfig.NextProtos(), http2.NextProtoTLS) {
			s.tlsConfig.SetNextProtos(append([]string{"h2"}, s.tlsConfig.NextProtos()...))
		}
		tcpListener = aTLS.NewListener(tcpListener, s.tlsConfig)
	}
	s.grpcServer = grpc.NewServer()
	pb.RegisterManagerServer(s.grpcServer, s)
	go func() {
		err = s.grpcServer.Serve(tcpListener)
		if err != nil && !errors.Is(err, grpc.ErrServerStopped) {
			s.logger.Error("serve error: ", err)
		}
	}()
	return nil
}

func (s *Service) Close() error {
	return nil
}