move timestamp generator to protocol

common/protocol/time.go

@@ -1,6 +1,8 @@
 package protocol
 
 import (
+	"math/rand"
+
 	"github.com/v2ray/v2ray-core/common/serial"
 )
 
@@ -9,3 +11,12 @@ type Timestamp int64
 func (this Timestamp) Bytes() []byte {
 	return serial.Int64Literal(this).Bytes()
 }
+
+type TimestampGenerator func() Timestamp
+
+func NewTimestampGenerator(base Timestamp, delta int) TimestampGenerator {
+	return func() Timestamp {
+		rangeInDelta := rand.Intn(delta*2) - delta
+		return base + Timestamp(rangeInDelta)
+	}
+}

common/protocol/time_test.go

@@ -4,8 +4,7 @@ import (
 	"testing"
 	"time"
 
-	"github.com/v2ray/v2ray-core/common/protocol"
-	. "github.com/v2ray/v2ray-core/proxy/vmess/protocol"
+	. "github.com/v2ray/v2ray-core/common/protocol"
 	v2testing "github.com/v2ray/v2ray-core/testing"
 	"github.com/v2ray/v2ray-core/testing/assert"
 )
@@ -15,10 +14,10 @@ func TestGenerateRandomInt64InRange(t *testing.T) {
 
 	base := time.Now().Unix()
 	delta := 100
-	generator := NewRandomTimestampGenerator(protocol.Timestamp(base), delta)
+	generator := NewTimestampGenerator(Timestamp(base), delta)
 
 	for i := 0; i < 100; i++ {
-		v := int64(generator.Next())
+		v := int64(generator())
 		assert.Int64(v).AtMost(base + int64(delta))
 		assert.Int64(v).AtLeast(base - int64(delta))
 	}

proxy/vmess/outbound/outbound.go

@@ -107,7 +107,7 @@ func (this *VMessOutboundHandler) handleRequest(conn net.Conn, request *protocol
 
 	buffer := alloc.NewBuffer().Clear()
 	defer buffer.Release()
-	buffer, err = request.ToBytes(protocol.NewRandomTimestampGenerator(proto.Timestamp(time.Now().Unix()), 30), buffer)
+	buffer, err = request.ToBytes(proto.NewTimestampGenerator(proto.Timestamp(time.Now().Unix()), 30), buffer)
 	if err != nil {
 		log.Error("VMessOut: Failed to serialize VMess request: ", err)
 		return

proxy/vmess/protocol/rand.go

@@ -1,28 +0,0 @@
-package protocol
-
-import (
-	"math/rand"
-
-	"github.com/v2ray/v2ray-core/common/protocol"
-)
-
-type RandomTimestampGenerator interface {
-	Next() protocol.Timestamp
-}
-
-type RealRandomTimestampGenerator struct {
-	base  protocol.Timestamp
-	delta int
-}
-
-func NewRandomTimestampGenerator(base protocol.Timestamp, delta int) RandomTimestampGenerator {
-	return &RealRandomTimestampGenerator{
-		base:  base,
-		delta: delta,
-	}
-}
-
-func (this *RealRandomTimestampGenerator) Next() protocol.Timestamp {
-	rangeInDelta := rand.Intn(this.delta*2) - this.delta
-	return this.base + protocol.Timestamp(rangeInDelta)
-}

proxy/vmess/protocol/vmess.go

@@ -189,12 +189,12 @@ func (this *VMessRequestReader) Read(reader io.Reader) (*VMessRequest, error) {
 }
 
 // ToBytes returns a VMessRequest in the form of byte array.
-func (this *VMessRequest) ToBytes(timestampGenerator RandomTimestampGenerator, buffer *alloc.Buffer) (*alloc.Buffer, error) {
+func (this *VMessRequest) ToBytes(timestampGenerator proto.TimestampGenerator, buffer *alloc.Buffer) (*alloc.Buffer, error) {
 	if buffer == nil {
 		buffer = alloc.NewSmallBuffer().Clear()
 	}
 
-	timestamp := timestampGenerator.Next()
+	timestamp := timestampGenerator()
 	idHash := IDHash(this.User.AnyValidID().Bytes())
 	idHash.Write(timestamp.Bytes())
 

proxy/vmess/protocol/vmess_test.go

@@ -16,12 +16,10 @@ import (
 	"github.com/v2ray/v2ray-core/testing/assert"
 )
 
-type FakeTimestampGenerator struct {
-	timestamp proto.Timestamp
-}
-
-func (this *FakeTimestampGenerator) Next() proto.Timestamp {
-	return this.timestamp
+func newStaticTimestampGenerator(t proto.Timestamp) proto.TimestampGenerator {
+	return func() proto.Timestamp {
+		return t
+	}
 }
 
 func TestVMessSerialization(t *testing.T) {
@@ -56,7 +54,7 @@ func TestVMessSerialization(t *testing.T) {
 
 	mockTime := proto.Timestamp(1823730)
 
-	buffer, err := request.ToBytes(&FakeTimestampGenerator{timestamp: mockTime}, nil)
+	buffer, err := request.ToBytes(newStaticTimestampGenerator(mockTime), nil)
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -114,6 +112,6 @@ func BenchmarkVMessRequestWriting(b *testing.B) {
 	request.Port = v2net.Port(80)
 
 	for i := 0; i < b.N; i++ {
-		request.ToBytes(NewRandomTimestampGenerator(proto.Timestamp(time.Now().Unix()), 30), nil)
+		request.ToBytes(proto.NewTimestampGenerator(proto.Timestamp(time.Now().Unix()), 30), nil)
 	}
 }