github.com/Asutorufa/yuhaiin@v0.3.6-0.20240502055049-7984da7023a0/pkg/node/node.go

package node

import (
	"context"
	"fmt"
	"net"
	"sync"

	"github.com/Asutorufa/yuhaiin/pkg/log"
	"github.com/Asutorufa/yuhaiin/pkg/net/netapi"
	"github.com/Asutorufa/yuhaiin/pkg/protos/node"
	gn "github.com/Asutorufa/yuhaiin/pkg/protos/node/grpc"
	"github.com/Asutorufa/yuhaiin/pkg/protos/node/latency"
	"github.com/Asutorufa/yuhaiin/pkg/protos/node/point"
	"github.com/Asutorufa/yuhaiin/pkg/protos/node/subscribe"
	pt "github.com/Asutorufa/yuhaiin/pkg/protos/node/tag"
	"github.com/Asutorufa/yuhaiin/pkg/utils/jsondb"
	"google.golang.org/protobuf/types/known/durationpb"
	"google.golang.org/protobuf/types/known/emptypb"
	"google.golang.org/protobuf/types/known/wrapperspb"
)

type Nodes struct {
	gn.UnimplementedNodeServer

	ruleTags func() []string

	db       *jsondb.DB[*node.Node]
	manager  *manager
	outBound *outbound
	links    *link
}

func NewNodes(path string) *Nodes {
	f := &Nodes{
		db: load(path),
	}

	f.manager = NewManager(f.db.Data.Manager)
	f.outBound = NewOutbound(f.db, f.manager)
	f.links = NewLink(f.db, f.outBound, f.manager)

	return f
}

func (n *Nodes) SetRuleTags(f func() []string) { n.ruleTags = f }

func (n *Nodes) Now(context.Context, *emptypb.Empty) (*gn.NowResp, error) {
	return &gn.NowResp{
		Tcp: n.db.Data.Tcp,
		Udp: n.db.Data.Udp,
	}, nil
}

func (n *Nodes) Get(_ context.Context, s *wrapperspb.StringValue) (*point.Point, error) {
	p, ok := n.manager.GetNode(s.Value)
	if !ok {
		return &point.Point{}, fmt.Errorf("node not found")
	}

	return p, nil
}

func (n *Nodes) Save(c context.Context, p *point.Point) (*point.Point, error) {
	if p.Name == "" || p.Group == "" {
		return &point.Point{}, fmt.Errorf("add point name or group is empty")
	}
	n.manager.DeleteNode(p.Hash)
	n.manager.AddNode(p)
	return p, n.db.Save()
}

func (n *Nodes) Manager(context.Context, *emptypb.Empty) (*node.Manager, error) {
	m := n.manager.GetManager()

	if m.Tags == nil {
		m.Tags = map[string]*pt.Tags{}
	}

	if n.ruleTags != nil {
		for _, v := range n.ruleTags() {
			if _, ok := m.Tags[v]; !ok {
				m.Tags[v] = &pt.Tags{}
			}
		}
	}

	return m, nil
}

func (n *Nodes) Use(c context.Context, s *gn.UseReq) (*point.Point, error) {
	p, err := n.Get(c, &wrapperspb.StringValue{Value: s.Hash})
	if err != nil {
		return &point.Point{}, fmt.Errorf("get node failed: %w", err)
	}

	if s.Tcp {
		n.db.Data.Tcp = p
	}
	if s.Udp {
		n.db.Data.Udp = p
	}

	err = n.db.Save()
	if err != nil {
		return p, fmt.Errorf("save config failed: %w", err)
	}
	return p, nil
}

func (n *Nodes) Remove(_ context.Context, s *wrapperspb.StringValue) (*emptypb.Empty, error) {
	n.manager.DeleteNode(s.Value)
	return &emptypb.Empty{}, n.db.Save()
}

type latencyDialer struct {
	netapi.Proxy
	ipv6 bool
}

func (w *latencyDialer) Conn(ctx context.Context, a netapi.Address) (net.Conn, error) {
	if w.ipv6 {
		a.PreferIPv6(true)
	} else {
		a.PreferIPv4(true)
	}

	return w.Proxy.Conn(ctx, a)
}

func (w *latencyDialer) PacketConn(ctx context.Context, a netapi.Address) (net.PacketConn, error) {
	if w.ipv6 {
		a.PreferIPv6(true)
	} else {
		a.PreferIPv4(true)
	}

	return w.Proxy.PacketConn(ctx, a)
}

func (n *Nodes) Latency(c context.Context, req *latency.Requests) (*latency.Response, error) {
	resp := &latency.Response{IdLatencyMap: make(map[string]*durationpb.Duration)}
	var mu sync.Mutex

	var wg sync.WaitGroup
	for _, s := range req.Requests {
		wg.Add(1)
		go func(s *latency.Request) {
			defer wg.Done()
			p, err := n.Get(c, &wrapperspb.StringValue{Value: s.GetHash()})
			if err != nil {
				return
			}

			px, err := n.outBound.GetDialer(p)
			if err != nil {
				return
			}

			px = &latencyDialer{Proxy: px, ipv6: s.GetIpv6()}

			var t *durationpb.Duration
			z, ok := s.Protocol.Protocol.(interface {
				Latency(netapi.Proxy) (*durationpb.Duration, error)
			})
			if ok {
				t, err = z.Latency(px)
				if err != nil {
					log.Error("latency failed", "err", err)
				}
			}

			mu.Lock()
			resp.IdLatencyMap[s.Id] = t
			mu.Unlock()
		}(s)
	}

	wg.Wait()
	return resp, nil
}

func (n *Nodes) Outbound() *outbound { return n.outBound }

func load(path string) *jsondb.DB[*node.Node] {
	defaultNode := &node.Node{
		Tcp:   &point.Point{},
		Udp:   &point.Point{},
		Links: map[string]*subscribe.Link{},
		Manager: &node.Manager{
			GroupsV2: map[string]*node.Nodes{},
			Nodes:    map[string]*point.Point{},
			Tags:     map[string]*pt.Tags{},
		},
	}

	return jsondb.Open[*node.Node](path, defaultNode)
}