fix: snowflake remove network dependency (#947)

fix #938
apache · Dec 2, 2020 · 3123c43 · 3123c43
1 parent 9e4c617
commit 3123c43
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 14 deletions.
diff --git a/api/go.sum b/api/go.sum
@@ -118,6 +118,7 @@ github.com/shiningrush/droplet v0.2.1 h1:p2utttTbCfgiL+x0Zrb2jFeWspB7/o+v3e+R94G
 github.com/shiningrush/droplet v0.2.1/go.mod h1:akW2vIeamvMD6zj6wIBfzYn6StGXBxwlW3gA+hcHu5M=
 github.com/shiningrush/droplet v0.2.2 h1:jEqSGoJXlybt1yQdauu+waE+l7KYlguNs8VayMfQ96Q=
 github.com/shiningrush/droplet v0.2.2/go.mod h1:akW2vIeamvMD6zj6wIBfzYn6StGXBxwlW3gA+hcHu5M=
+github.com/shiningrush/droplet v0.2.3/go.mod h1:akW2vIeamvMD6zj6wIBfzYn6StGXBxwlW3gA+hcHu5M=
 github.com/shiningrush/droplet/wrapper/gin v0.2.0 h1:LHkU+TbSkpePgXrTg3hJoSZlCMS03GeWMl0t+oLkd44=
 github.com/shiningrush/droplet/wrapper/gin v0.2.0/go.mod h1:ZJu+sCRrVXn5Pg618c1KK3Ob2UiXGuPM1ROx5uMM9YQ=
 github.com/sirupsen/logrus v1.7.0 h1:ShrD1U9pZB12TX0cVy0DtePoCH97K8EtX+mg7ZARUtM=

diff --git a/api/internal/utils/utils.go b/api/internal/utils/utils.go
@@ -37,35 +37,42 @@ func init() {
 		}
 		salt = uint16(i)
 	}
-
+	ips, err := getLocalIPs()
+	if err != nil {
+		panic(err)
+	}
 	_sf = sonyflake.NewSonyflake(sonyflake.Settings{
 		MachineID: func() (u uint16, e error) {
-			return sumIP(GetOutboundIP()) + salt, nil
+			return sumIPs(ips) + salt, nil
 		},
 	})
 	if _sf == nil {
 		panic("sonyflake init failed")
 	}
 }
 
-func sumIP(ip net.IP) uint16 {
+func sumIPs(ips []net.IP) uint16 {
 	total := 0
-	for i := range ip {
-		total += int(ip[i])
+	for _, ip := range ips {
+		for i := range ip {
+			total += int(ip[i])
+		}
 	}
 	return uint16(total)
 }
 
-// Get preferred outbound ip of this machine
-func GetOutboundIP() net.IP {
-	conn, err := net.Dial("udp", "8.8.8.8:80")
+func getLocalIPs() ([]net.IP, error) {
+	var ips []net.IP
+	addrs, err := net.InterfaceAddrs()
 	if err != nil {
-		panic(err)
+		return ips, err
 	}
-	defer conn.Close()
-
-	localAddr := conn.LocalAddr().(*net.UDPAddr)
-	return localAddr.IP
+	for _, a := range addrs {
+		if ipNet, ok := a.(*net.IPNet); ok && !ipNet.IP.IsLoopback() && ipNet.IP.To4() != nil {
+			ips = append(ips, ipNet.IP)
+		}
+	}
+	return ips, nil
 }
 
 func GetFlakeUid() uint64 {

diff --git a/api/internal/utils/utils_test.go b/api/internal/utils/utils_test.go
@@ -17,8 +17,9 @@
 package utils
 
 import (
-	"github.com/stretchr/testify/assert"
 	"testing"
+
+	"github.com/stretchr/testify/assert"
 )
 
 func TestGetFlakeUid(t *testing.T) {
@@ -29,4 +30,15 @@ func TestGetFlakeUid(t *testing.T) {
 func TestGetFlakeUidStr(t *testing.T) {
 	id := GetFlakeUidStr()
 	assert.NotEqual(t, "", id)
+	assert.Equal(t, 18, len(id))
+}
+
+func TestGetLocalIPs(t *testing.T) {
+	_, err := getLocalIPs()
+	assert.Equal(t, nil, err)
+}
+
+func TestSumIPs_with_nil(t *testing.T) {
+	total := sumIPs(nil)
+	assert.Equal(t, uint16(0), total)
 }