Skip to content

Commit f8ce884

Browse files
Fix default machine ID to respect BitsMachineID setting
Co-authored-by: YoshiyukiMineo <7577673+YoshiyukiMineo@users.noreply.github.com>
1 parent 8069979 commit f8ce884

3 files changed

Lines changed: 102 additions & 1 deletion

File tree

v2/machine_bits_test.go

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
package sonyflake
2+
3+
import (
4+
"net"
5+
"testing"
6+
7+
"github.com/sony/sonyflake/v2/mock"
8+
)
9+
10+
// TestDefaultMachineIDWithCustomBits tests that when using default machine ID
11+
// with custom BitsMachineID, the machine ID is properly masked to fit within
12+
// the specified bit length.
13+
func TestDefaultMachineIDWithCustomBits(t *testing.T) {
14+
testCases := []struct {
15+
name string
16+
bitsMachineID int
17+
mockIP net.IP
18+
expectError bool
19+
}{
20+
{
21+
name: "10 bits machine ID with IP that fits",
22+
bitsMachineID: 10,
23+
mockIP: net.IP{192, 168, 0, 1}, // lower 16 bits = 1, fits in 10 bits
24+
expectError: false,
25+
},
26+
{
27+
name: "10 bits machine ID with IP that exceeds without masking",
28+
bitsMachineID: 10,
29+
mockIP: net.IP{192, 168, 255, 255}, // lower 16 bits = 65535, needs masking to fit in 10 bits
30+
expectError: false,
31+
},
32+
{
33+
name: "8 bits machine ID",
34+
bitsMachineID: 8,
35+
mockIP: net.IP{192, 168, 100, 200}, // lower 16 bits = 25800
36+
expectError: false,
37+
},
38+
{
39+
name: "default 16 bits",
40+
bitsMachineID: 0, // will use default 16
41+
mockIP: net.IP{192, 168, 255, 255},
42+
expectError: false,
43+
},
44+
}
45+
46+
for _, tc := range testCases {
47+
t.Run(tc.name, func(t *testing.T) {
48+
// Create a mock that returns our test IP
49+
mockInterfaceAddrs := mock.NewInterfaceAddrsWithIP(tc.mockIP)
50+
51+
settings := Settings{
52+
BitsMachineID: tc.bitsMachineID,
53+
}
54+
55+
// Temporarily replace the default interface addrs function
56+
oldDefaultInterfaceAddrs := defaultInterfaceAddrs
57+
defaultInterfaceAddrs = mockInterfaceAddrs
58+
defer func() { defaultInterfaceAddrs = oldDefaultInterfaceAddrs }()
59+
60+
sf, err := New(settings)
61+
62+
if tc.expectError {
63+
if err == nil {
64+
t.Errorf("expected error but got none")
65+
}
66+
} else {
67+
if err != nil {
68+
t.Errorf("unexpected error: %v", err)
69+
}
70+
if sf == nil {
71+
t.Error("sonyflake instance should not be nil")
72+
}
73+
74+
// Verify the machine ID fits within the specified bits
75+
expectedBits := tc.bitsMachineID
76+
if expectedBits == 0 {
77+
expectedBits = defaultBitsMachine
78+
}
79+
maxMachineID := 1 << expectedBits
80+
if sf.machine >= maxMachineID {
81+
t.Errorf("machine ID %d exceeds max for %d bits (%d)", sf.machine, expectedBits, maxMachineID)
82+
}
83+
}
84+
})
85+
}
86+
}

v2/mock/mock.go

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,3 +34,13 @@ func NewNilInterfaceAddrs() types.InterfaceAddrs {
3434
return []net.Addr{}, nil
3535
}
3636
}
37+
38+
// NewInterfaceAddrsWithIP returns a private IP address with the given IP.
39+
func NewInterfaceAddrsWithIP(ip net.IP) types.InterfaceAddrs {
40+
ifat := make([]net.Addr, 0, 1)
41+
ifat = append(ifat, &net.IPNet{IP: ip, Mask: []byte{255, 0, 0, 0}})
42+
43+
return func() ([]net.Addr, error) {
44+
return ifat, nil
45+
}
46+
}

v2/sonyflake.go

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,8 @@ import (
3636
//
3737
// MachineID returns the unique ID of a Sonyflake instance.
3838
// If MachineID returns an error, the instance will not be created.
39-
// If MachineID is nil, the default MachineID is used, which returns the lower 16 bits of the private IP address.
39+
// If MachineID is nil, the default MachineID is used, which returns the lower bits
40+
// of the private IP address, masked to fit within BitsMachineID bits.
4041
//
4142
// CheckMachineID validates the uniqueness of a machine ID.
4243
// If CheckMachineID returns false, the instance will not be created.
@@ -154,6 +155,10 @@ func New(st Settings) (*Sonyflake, error) {
154155
var err error
155156
if st.MachineID == nil {
156157
sf.machine, err = lower16BitPrivateIP(defaultInterfaceAddrs)
158+
if err == nil {
159+
// Mask to use only the required number of bits
160+
sf.machine = sf.machine & (1<<sf.bitsMachine - 1)
161+
}
157162
} else {
158163
sf.machine, err = st.MachineID()
159164
}

0 commit comments

Comments
 (0)