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

     1  package node
     2  
     3  import (
     4  	"bufio"
     5  	"bytes"
     6  	"encoding/base64"
     7  	"fmt"
     8  	"io"
     9  	"log/slog"
    10  	"net/http"
    11  	"sync"
    12  
    13  	"github.com/Asutorufa/yuhaiin/internal/version"
    14  	"github.com/Asutorufa/yuhaiin/pkg/log"
    15  	"github.com/Asutorufa/yuhaiin/pkg/node/parser"
    16  	"github.com/Asutorufa/yuhaiin/pkg/protos/node"
    17  	"github.com/Asutorufa/yuhaiin/pkg/protos/node/point"
    18  	"github.com/Asutorufa/yuhaiin/pkg/protos/node/subscribe"
    19  	"github.com/Asutorufa/yuhaiin/pkg/utils/jsondb"
    20  	"github.com/Asutorufa/yuhaiin/pkg/utils/net"
    21  )
    22  
    23  type link struct {
    24  	outbound *outbound
    25  	manager  *manager
    26  
    27  	db *jsondb.DB[*node.Node]
    28  
    29  	mu sync.RWMutex
    30  }
    31  
    32  func NewLink(db *jsondb.DB[*node.Node], outbound *outbound, manager *manager) *link {
    33  	return &link{outbound: outbound, manager: manager, db: db}
    34  }
    35  
    36  func (l *link) Save(ls []*subscribe.Link) {
    37  	l.mu.Lock()
    38  	defer l.mu.Unlock()
    39  
    40  	if l.db.Data.Links == nil {
    41  		l.db.Data.Links = make(map[string]*subscribe.Link)
    42  	}
    43  
    44  	for _, z := range ls {
    45  
    46  		node, err := parseUrl([]byte(z.Url), &subscribe.Link{Name: z.Name})
    47  		if err == nil {
    48  			l.addNode(node) // link is a node
    49  		} else {
    50  			l.db.Data.Links[z.Name] = z // link is a subscription
    51  		}
    52  
    53  	}
    54  }
    55  
    56  func (l *link) Delete(names []string) {
    57  	l.mu.Lock()
    58  	defer l.mu.Unlock()
    59  
    60  	for _, z := range names {
    61  		delete(l.db.Data.Links, z)
    62  	}
    63  }
    64  
    65  func (l *link) Links() map[string]*subscribe.Link { return l.db.Data.Links }
    66  
    67  func (l *link) Update(names []string) {
    68  	if l.db.Data.Links == nil {
    69  		l.db.Data.Links = make(map[string]*subscribe.Link)
    70  	}
    71  
    72  	wg := sync.WaitGroup{}
    73  	for _, str := range names {
    74  		link, ok := l.db.Data.Links[str]
    75  		if !ok {
    76  			continue
    77  		}
    78  
    79  		wg.Add(1)
    80  		go func(link *subscribe.Link) {
    81  			defer wg.Done()
    82  			if err := l.update(l.outbound.Do, link); err != nil {
    83  				log.Error("get one link failed", "err", err)
    84  			}
    85  		}(link)
    86  	}
    87  
    88  	wg.Wait()
    89  
    90  	oo := l.db.Data.Udp
    91  	if p, ok := l.manager.GetNodeByName(oo.Group, oo.Name); ok {
    92  		l.db.Data.Udp = p
    93  	}
    94  
    95  	oo = l.db.Data.Tcp
    96  	if p, ok := l.manager.GetNodeByName(oo.Group, oo.Name); ok {
    97  		l.db.Data.Tcp = p
    98  	}
    99  }
   100  
   101  type trimBase64Reader struct {
   102  	r io.Reader
   103  }
   104  
   105  func (t *trimBase64Reader) Read(b []byte) (int, error) {
   106  	n, err := t.r.Read(b)
   107  
   108  	if n > 0 {
   109  		if i := bytes.IndexByte(b[:n], '='); i > 0 {
   110  			n = i
   111  		}
   112  	}
   113  
   114  	return n, err
   115  }
   116  
   117  func (n *link) update(do func(*http.Request) (*http.Response, error), link *subscribe.Link) error {
   118  	req, err := http.NewRequest("GET", link.Url, nil)
   119  	if err != nil {
   120  		return fmt.Errorf("create request failed: %w", err)
   121  	}
   122  
   123  	req.Header.Set("User-Agent", fmt.Sprintf("%s/%s-%s", version.AppName, version.Version, version.GitCommit))
   124  
   125  	res, err := do(req)
   126  	if err != nil {
   127  		return fmt.Errorf("get %s failed: %w", link.Name, err)
   128  	}
   129  	defer res.Body.Close()
   130  
   131  	n.manager.DeleteRemoteNodes(link.Name)
   132  
   133  	base64r := base64.NewDecoder(base64.RawStdEncoding, &trimBase64Reader{res.Body})
   134  	scanner := bufio.NewScanner(base64r)
   135  	for scanner.Scan() {
   136  		if len(scanner.Bytes()) == 0 {
   137  			continue
   138  		}
   139  
   140  		node, err := parseUrl(scanner.Bytes(), link)
   141  		if err != nil {
   142  			log.Error("parse url failed", slog.String("url", scanner.Text()), slog.Any("err", err))
   143  		} else {
   144  			n.addNode(node)
   145  		}
   146  	}
   147  
   148  	return scanner.Err()
   149  }
   150  
   151  func (n *link) addNode(node *point.Point) {
   152  	n.manager.DeleteNode(node.Hash)
   153  	n.manager.AddNode(node)
   154  }
   155  
   156  var schemeTypeMap = map[string]subscribe.Type{
   157  	"ss":     subscribe.Type_shadowsocks,
   158  	"ssr":    subscribe.Type_shadowsocksr,
   159  	"vmess":  subscribe.Type_vmess,
   160  	"trojan": subscribe.Type_trojan,
   161  }
   162  
   163  func parseUrl(str []byte, l *subscribe.Link) (no *point.Point, err error) {
   164  	t := l.Type
   165  
   166  	if t == subscribe.Type_reserve {
   167  		scheme, _, _ := net.GetScheme(string(str))
   168  		t = schemeTypeMap[scheme]
   169  	}
   170  
   171  	no, err = parser.Parse(t, str)
   172  	if err != nil {
   173  		return nil, fmt.Errorf("parse link data failed: %w", err)
   174  	}
   175  	no.Group = l.Name
   176  	return no, nil
   177  }