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")
 )
@@ -260,11 +261,26 @@ func (sf *Sonyflake) ToTime(id int64) time.Time {
 // 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 {
+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)
+		int64(machineID), nil
 }
 
 // Decompose returns a set of Sonyflake ID parts.

v2/sonyflake_test.go

@@ -257,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)
 
@@ -411,7 +412,11 @@ func TestComposeAndDecompose(t *testing.T) {
 
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
-			id := sf.Compose(tc.time, tc.sequence, tc.machineID)
+			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
@@ -437,3 +442,54 @@ func TestComposeAndDecompose(t *testing.T) {
 		})
 	}
 }
+
+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)
+			}
+		})
+	}
+}