github.com/nxtrace/NTrace-core@v1.3.1-0.20240513132635-39169291e8c9/ipgeo/leo.go (about)

     1  package ipgeo
     2  
     3  import (
     4  	"encoding/json"
     5  	"errors"
     6  	"strconv"
     7  	"sync"
     8  	"time"
     9  
    10  	"github.com/nxtrace/NTrace-core/wshandle"
    11  	"github.com/tidwall/gjson"
    12  )
    13  
    14  /***
    15   * 原理介绍 By Leo
    16   * WebSocket 一共开启了一个发送和一个接收协程，在 New 了一个连接的实例对象后，不给予关闭，持续化连接
    17   * 当有新的IP请求时，一直在等待IP数据的发送协程接收到从 leo.go 的 sendIPRequest 函数发来的IP数据，向服务端发送数据
    18   * 由于实际使用时有大量并发，但是 ws 在同一时刻每次有且只能处理一次发送一条数据，所以必须给 ws 连接上互斥锁，保证每次只有一个协程访问
    19   * 运作模型可以理解为一个 Node 一直在等待数据，当获得一个新的任务后，转交给下一个协程，不再关注这个 Node 的下一步处理过程，并且回到空闲状态继续等待新的任务
    20  ***/
    21  
    22  // IPPool IP 查询池 map - ip - ip channel
    23  type IPPool struct {
    24  	pool    map[string]chan IPGeoData
    25  	poolMux sync.Mutex
    26  }
    27  
    28  var IPPools = IPPool{
    29  	pool: make(map[string]chan IPGeoData),
    30  }
    31  
    32  func sendIPRequest(ip string) {
    33  	wsConn := wshandle.GetWsConn()
    34  	wsConn.MsgSendCh <- ip
    35  }
    36  
    37  func receiveParse() {
    38  	// 获得连接实例
    39  	wsConn := wshandle.GetWsConn()
    40  	// 防止多协程抢夺一个ws连接，导致死锁，当一个协程获得ws的控制权后上锁
    41  	wsConn.ConnMux.Lock()
    42  	// 函数退出时解锁，给其他协程使用
    43  	defer wsConn.ConnMux.Unlock()
    44  	for {
    45  		// 接收到了一条IP信息
    46  		data := <-wsConn.MsgReceiveCh
    47  
    48  		// json解析 -> data
    49  		res := gjson.Parse(data)
    50  		// 根据返回的IP信息，发送给对应等待回复的IP通道上
    51  		var domain = res.Get("domain").String()
    52  
    53  		if res.Get("domain").String() == "" {
    54  			domain = res.Get("owner").String()
    55  		}
    56  
    57  		m := make(map[string][]string)
    58  		err := json.Unmarshal([]byte(res.Get("router").String()), &m)
    59  		if err != nil {
    60  			// 此处是正常的，因为有些IP没有路由信息
    61  		}
    62  
    63  		lat, _ := strconv.ParseFloat(res.Get("lat").String(), 32)
    64  		lng, _ := strconv.ParseFloat(res.Get("lng").String(), 32)
    65  
    66  		IPPools.pool[gjson.Parse(data).Get("ip").String()] <- IPGeoData{
    67  			Asnumber:  res.Get("asnumber").String(),
    68  			Country:   res.Get("country").String(),
    69  			CountryEn: res.Get("country_en").String(),
    70  			Prov:      res.Get("prov").String(),
    71  			ProvEn:    res.Get("prov_en").String(),
    72  			City:      res.Get("city").String(),
    73  			CityEn:    res.Get("city_en").String(),
    74  			District:  res.Get("district").String(),
    75  			Owner:     domain,
    76  			Lat:       lat,
    77  			Lng:       lng,
    78  			Isp:       res.Get("isp").String(),
    79  			Whois:     res.Get("whois").String(),
    80  			Prefix:    res.Get("prefix").String(),
    81  			Router:    m,
    82  		}
    83  	}
    84  }
    85  
    86  func LeoIP(ip string, timeout time.Duration, lang string, maptrace bool) (*IPGeoData, error) {
    87  	// TODO: 根据lang的值请求中文/英文API
    88  	// TODO: 根据maptrace的值决定是否请求经纬度信息
    89  	if timeout < 5*time.Second {
    90  		timeout = 5 * time.Second
    91  	}
    92  
    93  	// 缓存中没有找到IP信息，需要请求API获取
    94  	IPPools.poolMux.Lock()
    95  	// 如果之前已经被别的协程初始化过了就不用初始化了
    96  	if IPPools.pool[ip] == nil {
    97  		IPPools.pool[ip] = make(chan IPGeoData)
    98  	}
    99  	IPPools.poolMux.Unlock()
   100  	// 发送请求
   101  	sendIPRequest(ip)
   102  	// 同步开启监听
   103  	go receiveParse()
   104  
   105  	// 拥塞，等待数据返回
   106  	select {
   107  	case res := <-IPPools.pool[ip]:
   108  		return &res, nil
   109  	// 5秒后依旧没有接收到返回的IP数据，不再等待，超时异常处理
   110  	case <-time.After(timeout):
   111  		// 这里不可以返回一个 nil，否则在访问对象内部的键值的时候会报空指针的 Fatal Error
   112  		return &IPGeoData{}, errors.New("TimeOut")
   113  	}
   114  }