Check compose args (#74)

* v2: take only 'bitsMachine' least significant bits from generated machine IDs so it doesn't corrupt the other ID parts (#72)

Co-authored-by: ok32 <e.starikoff@gmail.com>
Co-authored-by: Yoshiyuki Mineo <Yoshiyuki.Mineo@jp.sony.com>

* Refactor Sonyflake Compose method to return errors for invalid parameters and update related tests. Consolidate error handling for start time, sequence, and machine ID validations. Remove unused variable in tests for clarity.

---------

Co-authored-by: ok32 <artuh.gubanova@gmail.com>
Co-authored-by: ok32 <e.starikoff@gmail.com>

v2/sonyflake.go

@@ -74,8 +74,9 @@ var (
 	ErrInvalidBitsSequence  = errors.New("invalid bit length for sequence number")
 	ErrInvalidBitsMachineID = errors.New("invalid bit length for machine id")
 	ErrInvalidTimeUnit      = errors.New("invalid time unit")
+	ErrInvalidSequence      = errors.New("invalid sequence number")
 	ErrInvalidMachineID     = errors.New("invalid machine id")
-	ErrStartTimeAhead       = errors.New("start time is ahead of now")
+	ErrStartTimeAhead       = errors.New("start time is ahead")
 	ErrOverTimeLimit        = errors.New("over the time limit")
 	ErrNoPrivateAddress     = errors.New("no private ip address")
 )
@@ -157,6 +158,10 @@ func New(st Settings) (*Sonyflake, error) {
 		return nil, err
 	}
 
+	if sf.machine < 0 || sf.machine >= 1<<sf.bitsMachine {
+		return nil, ErrInvalidMachineID
+	}
+
 	if st.CheckMachineID != nil && !st.CheckMachineID(sf.machine) {
 		return nil, ErrInvalidMachineID
 	}
@@ -252,6 +257,32 @@ func (sf *Sonyflake) ToTime(id int64) time.Time {
 	return time.Unix(0, (sf.startTime+sf.timePart(id))*sf.timeUnit)
 }
 
+// Compose creates a Sonyflake ID from its components.
+// The time parameter should be the time when the ID was generated.
+// The sequence parameter should be between 0 and 2^BitsSequence-1 (inclusive).
+// The machineID parameter should be between 0 and 2^BitsMachineID-1 (inclusive).
+func (sf *Sonyflake) Compose(t time.Time, sequence, machineID int) (int64, error) {
+	elapsedTime := sf.toInternalTime(t.UTC()) - sf.startTime
+	if elapsedTime < 0 {
+		return 0, ErrStartTimeAhead
+	}
+	if elapsedTime >= 1<<sf.bitsTime {
+		return 0, ErrOverTimeLimit
+	}
+
+	if sequence < 0 || sequence >= 1<<sf.bitsSequence {
+		return 0, ErrInvalidSequence
+	}
+
+	if machineID < 0 || machineID >= 1<<sf.bitsMachine {
+		return 0, ErrInvalidMachineID
+	}
+
+	return elapsedTime<<(sf.bitsSequence+sf.bitsMachine) |
+		int64(sequence)<<sf.bitsMachine |
+		int64(machineID), nil
+}
+
 // Decompose returns a set of Sonyflake ID parts.
 func (sf *Sonyflake) Decompose(id int64) map[string]int64 {
 	time := sf.timePart(id)

v2/sonyflake_test.go

@@ -65,6 +65,24 @@ func TestNew(t *testing.T) {
 			},
 			err: errGetMachineID,
 		},
+		{
+			name: "too large machine id",
+			settings: Settings{
+				MachineID: func() (int, error) {
+					return 1 << defaultBitsMachine, nil
+				},
+			},
+			err: ErrInvalidMachineID,
+		},
+		{
+			name: "negative machine id",
+			settings: Settings{
+				MachineID: func() (int, error) {
+					return -1, nil
+				},
+			},
+			err: ErrInvalidMachineID,
+		},
 		{
 			name: "invalid machine id",
 			settings: Settings{
@@ -239,10 +257,11 @@ func pseudoSleep(sf *Sonyflake, period time.Duration) {
 	sf.startTime -= int64(period) / sf.timeUnit
 }
 
+const year = time.Duration(365*24) * time.Hour
+
 func TestNextID_ReturnsError(t *testing.T) {
 	sf := newSonyflake(t, Settings{StartTime: time.Now()})
 
-	year := time.Duration(365*24) * time.Hour
 	pseudoSleep(sf, time.Duration(174)*year)
 	nextID(t, sf)
 
@@ -253,22 +272,6 @@ func TestNextID_ReturnsError(t *testing.T) {
 	}
 }
 
-func TestToTime(t *testing.T) {
-	start := time.Now()
-	sf := newSonyflake(t, Settings{
-		TimeUnit:  time.Millisecond,
-		StartTime: start,
-	})
-
-	id := nextID(t, sf)
-
-	tm := sf.ToTime(id)
-	diff := tm.Sub(start)
-	if diff < 0 || diff >= time.Duration(sf.timeUnit) {
-		t.Errorf("unexpected time: %v", tm)
-	}
-}
-
 func TestPrivateIPv4(t *testing.T) {
 	testCases := []struct {
 		description    string
@@ -351,3 +354,142 @@ func TestLower16BitPrivateIP(t *testing.T) {
 		})
 	}
 }
+
+func TestToTime(t *testing.T) {
+	start := time.Now()
+	sf := newSonyflake(t, Settings{
+		TimeUnit:  time.Millisecond,
+		StartTime: start,
+	})
+
+	id := nextID(t, sf)
+
+	tm := sf.ToTime(id)
+	diff := tm.Sub(start)
+	if diff < 0 || diff >= time.Duration(sf.timeUnit) {
+		t.Errorf("unexpected time: %v", tm)
+	}
+}
+
+func TestComposeAndDecompose(t *testing.T) {
+	now := time.Now()
+	sf := newSonyflake(t, Settings{
+		TimeUnit:  time.Millisecond,
+		StartTime: now,
+	})
+
+	testCases := []struct {
+		name      string
+		time      time.Time
+		sequence  int
+		machineID int
+	}{
+		{
+			name:      "zero values",
+			time:      now,
+			sequence:  0,
+			machineID: 0,
+		},
+		{
+			name:      "max sequence",
+			time:      now,
+			sequence:  1<<sf.bitsSequence - 1,
+			machineID: 0,
+		},
+		{
+			name:      "max machine id",
+			time:      now,
+			sequence:  0,
+			machineID: 1<<sf.bitsMachine - 1,
+		},
+		{
+			name:      "future time",
+			time:      now.Add(time.Hour),
+			sequence:  0,
+			machineID: 0,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			id, err := sf.Compose(tc.time, tc.sequence, tc.machineID)
+			if err != nil {
+				t.Fatalf("unexpected error: %v", err)
+			}
+
+			parts := sf.Decompose(id)
+
+			// Verify time part
+			expectedTime := sf.toInternalTime(tc.time.UTC()) - sf.startTime
+			if parts["time"] != expectedTime {
+				t.Errorf("time mismatch: got %d, want %d", parts["time"], expectedTime)
+			}
+
+			// Verify sequence part
+			if parts["sequence"] != int64(tc.sequence) {
+				t.Errorf("sequence mismatch: got %d, want %d", parts["sequence"], tc.sequence)
+			}
+
+			// Verify machine id part
+			if parts["machine"] != int64(tc.machineID) {
+				t.Errorf("machine id mismatch: got %d, want %d", parts["machine"], tc.machineID)
+			}
+
+			// Verify id part
+			if parts["id"] != id {
+				t.Errorf("id mismatch: got %d, want %d", parts["id"], id)
+			}
+		})
+	}
+}
+
+func TestCompose_ReturnsError(t *testing.T) {
+	start := time.Now()
+	sf := newSonyflake(t, Settings{StartTime: start})
+
+	testCases := []struct {
+		name      string
+		time      time.Time
+		sequence  int
+		machineID int
+		err       error
+	}{
+		{
+			name:      "start time ahead",
+			time:      start.Add(-time.Second),
+			sequence:  0,
+			machineID: 0,
+			err:       ErrStartTimeAhead,
+		},
+		{
+			name:      "over time limit",
+			time:      start.Add(time.Duration(175) * year),
+			sequence:  0,
+			machineID: 0,
+			err:       ErrOverTimeLimit,
+		},
+		{
+			name:      "invalid sequence",
+			time:      start,
+			sequence:  1 << sf.bitsSequence,
+			machineID: 0,
+			err:       ErrInvalidSequence,
+		},
+		{
+			name:      "invalid machine id",
+			time:      start,
+			sequence:  0,
+			machineID: 1 << sf.bitsMachine,
+			err:       ErrInvalidMachineID,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			_, err := sf.Compose(tc.time, tc.sequence, tc.machineID)
+			if !errors.Is(err, tc.err) {
+				t.Errorf("unexpected error: %v", err)
+			}
+		})
+	}
+}