package kcp import ( "sync" "v2ray.com/core/common/buf" "v2ray.com/core/common/serial" ) // Command is a KCP command that indicate the purpose of a Segment. type Command byte const ( // CommandACK indicates an AckSegment. CommandACK Command = 0 // CommandData indicates a DataSegment. CommandData Command = 1 // CommandTerminate indicates that peer terminates the connection. CommandTerminate Command = 2 // CommandPing indicates a ping. CommandPing Command = 3 ) type SegmentOption byte const ( SegmentOptionClose SegmentOption = 1 ) type Segment interface { Release() Conversation() uint16 Command() Command ByteSize() int32 Bytes() buf.Supplier parse(conv uint16, cmd Command, opt SegmentOption, buf []byte) (bool, []byte) } const ( DataSegmentOverhead = 18 ) var dataSegmentPool = sync.Pool{ New: func() interface{} { return new(DataSegment) }, } type DataSegment struct { Conv uint16 Option SegmentOption Timestamp uint32 Number uint32 SendingNext uint32 payload *buf.Buffer timeout uint32 transmit uint32 } func NewDataSegment() *DataSegment { seg := dataSegmentPool.Get().(*DataSegment) seg.Conv = 0 seg.timeout = 0 seg.transmit = 0 return seg } func (s *DataSegment) parse(conv uint16, cmd Command, opt SegmentOption, buf []byte) (bool, []byte) { s.Conv = conv s.Option = opt if len(buf) < 15 { return false, nil } s.Timestamp = serial.BytesToUint32(buf) buf = buf[4:] s.Number = serial.BytesToUint32(buf) buf = buf[4:] s.SendingNext = serial.BytesToUint32(buf) buf = buf[4:] dataLen := int(serial.BytesToUint16(buf)) buf = buf[2:] if len(buf) < dataLen { return false, nil } s.Data().Clear() s.Data().Write(buf[:dataLen]) buf = buf[dataLen:] return true, buf } func (s *DataSegment) Conversation() uint16 { return s.Conv } func (*DataSegment) Command() Command { return CommandData } func (s *DataSegment) Detach() *buf.Buffer { r := s.payload s.payload = nil return r } func (s *DataSegment) Data() *buf.Buffer { if s.payload == nil { s.payload = buf.New() } return s.payload } func (s *DataSegment) Bytes() buf.Supplier { return func(b []byte) (int, error) { b = serial.Uint16ToBytes(s.Conv, b[:0]) b = append(b, byte(CommandData), byte(s.Option)) b = serial.Uint32ToBytes(s.Timestamp, b) b = serial.Uint32ToBytes(s.Number, b) b = serial.Uint32ToBytes(s.SendingNext, b) b = serial.Uint16ToBytes(uint16(s.payload.Len()), b) b = append(b, s.payload.Bytes()...) return len(b), nil } } func (s *DataSegment) ByteSize() int32 { return 2 + 1 + 1 + 4 + 4 + 4 + 2 + s.payload.Len() } func (s *DataSegment) Release() { s.payload.Release() s.payload = nil dataSegmentPool.Put(s) } var ackSegmentPool = sync.Pool{ New: func() interface{} { return &AckSegment{ NumberList: make([]uint32, 0, 16), } }, } type AckSegment struct { Conv uint16 Option SegmentOption ReceivingWindow uint32 ReceivingNext uint32 Timestamp uint32 NumberList []uint32 } const ackNumberLimit = 128 func NewAckSegment() *AckSegment { seg := ackSegmentPool.Get().(*AckSegment) seg.NumberList = seg.NumberList[:0] return seg } func (s *AckSegment) parse(conv uint16, cmd Command, opt SegmentOption, buf []byte) (bool, []byte) { s.Conv = conv s.Option = opt if len(buf) < 13 { return false, nil } s.ReceivingWindow = serial.BytesToUint32(buf) buf = buf[4:] s.ReceivingNext = serial.BytesToUint32(buf) buf = buf[4:] s.Timestamp = serial.BytesToUint32(buf) buf = buf[4:] count := int(buf[0]) buf = buf[1:] if len(buf) < count*4 { return false, nil } for i := 0; i < count; i++ { s.PutNumber(serial.BytesToUint32(buf)) buf = buf[4:] } return true, buf } func (s *AckSegment) Conversation() uint16 { return s.Conv } func (*AckSegment) Command() Command { return CommandACK } func (s *AckSegment) PutTimestamp(timestamp uint32) { if timestamp-s.Timestamp < 0x7FFFFFFF { s.Timestamp = timestamp } } func (s *AckSegment) PutNumber(number uint32) { s.NumberList = append(s.NumberList, number) } func (s *AckSegment) IsFull() bool { return len(s.NumberList) == ackNumberLimit } func (s *AckSegment) IsEmpty() bool { return len(s.NumberList) == 0 } func (s *AckSegment) ByteSize() int32 { return 2 + 1 + 1 + 4 + 4 + 4 + 1 + int32(len(s.NumberList)*4) } func (s *AckSegment) Bytes() buf.Supplier { return func(b []byte) (int, error) { b = serial.Uint16ToBytes(s.Conv, b[:0]) b = append(b, byte(CommandACK), byte(s.Option)) b = serial.Uint32ToBytes(s.ReceivingWindow, b) b = serial.Uint32ToBytes(s.ReceivingNext, b) b = serial.Uint32ToBytes(s.Timestamp, b) count := byte(len(s.NumberList)) b = append(b, count) for _, number := range s.NumberList { b = serial.Uint32ToBytes(number, b) } return int(s.ByteSize()), nil } } func (s *AckSegment) Release() { ackSegmentPool.Put(s) } type CmdOnlySegment struct { Conv uint16 Cmd Command Option SegmentOption SendingNext uint32 ReceivingNext uint32 PeerRTO uint32 } func NewCmdOnlySegment() *CmdOnlySegment { return new(CmdOnlySegment) } func (s *CmdOnlySegment) parse(conv uint16, cmd Command, opt SegmentOption, buf []byte) (bool, []byte) { s.Conv = conv s.Cmd = cmd s.Option = opt if len(buf) < 12 { return false, nil } s.SendingNext = serial.BytesToUint32(buf) buf = buf[4:] s.ReceivingNext = serial.BytesToUint32(buf) buf = buf[4:] s.PeerRTO = serial.BytesToUint32(buf) buf = buf[4:] return true, buf } func (s *CmdOnlySegment) Conversation() uint16 { return s.Conv } func (s *CmdOnlySegment) Command() Command { return s.Cmd } func (*CmdOnlySegment) ByteSize() int32 { return 2 + 1 + 1 + 4 + 4 + 4 } func (s *CmdOnlySegment) Bytes() buf.Supplier { return func(b []byte) (int, error) { b = serial.Uint16ToBytes(s.Conv, b[:0]) b = append(b, byte(s.Cmd), byte(s.Option)) b = serial.Uint32ToBytes(s.SendingNext, b) b = serial.Uint32ToBytes(s.ReceivingNext, b) b = serial.Uint32ToBytes(s.PeerRTO, b) return len(b), nil } } func (*CmdOnlySegment) Release() {} func ReadSegment(buf []byte) (Segment, []byte) { if len(buf) < 4 { return nil, nil } conv := serial.BytesToUint16(buf) buf = buf[2:] cmd := Command(buf[0]) opt := SegmentOption(buf[1]) buf = buf[2:] var seg Segment switch cmd { case CommandData: seg = NewDataSegment() case CommandACK: seg = NewAckSegment() default: seg = NewCmdOnlySegment() } valid, extra := seg.parse(conv, cmd, opt, buf) if !valid { return nil, nil } return seg, extra }