Refactor vision unpadding

proxy/vless/encoding/encoding.go

@@ -217,6 +217,7 @@ func XtlsRead(reader buf.Reader, writer buf.Writer, timer signal.ActivityUpdater
 		shouldSwitchToDirectCopy := false
 		var remainingContent int32 = -1
 		var remainingPadding int32 = -1
+		var remainingCommand int32 = -1
 		currentCommand := 0
 		for {
 			if shouldSwitchToDirectCopy {
@@ -258,12 +259,17 @@ func XtlsRead(reader buf.Reader, writer buf.Writer, timer signal.ActivityUpdater
 			buffer, err := reader.ReadMultiBuffer()
 			if !buffer.IsEmpty() {
 				if withinPaddingBuffers || *numberOfPacketToFilter > 0 {
-					buffer = XtlsUnpadding(ctx, buffer, userUUID, &remainingContent, &remainingPadding, &currentCommand)
+					mb2 := make(buf.MultiBuffer, 0, len(buffer))
+					for _, b := range buffer {
+						newbuffer := XtlsUnpadding(ctx, b, userUUID, &remainingContent, &remainingPadding, &remainingCommand, &currentCommand)
+						if newbuffer.Len() > 0 {
+							mb2 = append(mb2, newbuffer)
+						}
+					}
+					buffer = mb2
 					if remainingContent == 0 && remainingPadding == 0 {
 						if currentCommand == 1 {
 							withinPaddingBuffers = false
-							remainingContent = -1
-							remainingPadding = -1 // set to initial state to parse the next padding
 						} else if currentCommand == 2 {
 							withinPaddingBuffers = false
 							shouldSwitchToDirectCopy = true
@@ -512,72 +518,70 @@ func XtlsPadding(b *buf.Buffer, command byte, userUUID *[]byte, longPadding bool
 }
 
 // XtlsUnpadding remove padding and parse command
-func XtlsUnpadding(ctx context.Context, buffer buf.MultiBuffer, userUUID []byte, remainingContent *int32, remainingPadding *int32, currentCommand *int) buf.MultiBuffer {
-	posindex := 0
-	var posByte int32 = 0
-	if *remainingContent == -1 && *remainingPadding == -1 {
-		for i, b := range buffer {
-			if b.Len() >= 21 && bytes.Equal(userUUID, b.BytesTo(16)) {
-				posindex = i
-				posByte = 16
-				*remainingContent = 0
-				*remainingPadding = 0
-				*currentCommand = 0
-				break
-			}
+func XtlsUnpadding(ctx context.Context, b *buf.Buffer, userUUID []byte, remainingContent *int32, remainingPadding *int32, remainingCommand *int32, currentCommand *int) *buf.Buffer {
+	if *remainingCommand == -1 && *remainingContent == -1 && *remainingPadding == -1 { // inital state
+		if b.Len() >= 21 && bytes.Equal(userUUID, b.BytesTo(16)) {
+			b.Advance(16)
+			*remainingCommand = 5
+		} else {
+			return b
 		}
 	}
-	if *remainingContent == -1 && *remainingPadding == -1 {
-		return buffer
-	}
-	mb2 := make(buf.MultiBuffer, 0, len(buffer))
-	for i := 0; i < posindex; i++ {
-		newbuffer := buf.New()
-		newbuffer.Write(buffer[i].Bytes())
-		mb2 = append(mb2, newbuffer)
-	}
-	for i := posindex; i < len(buffer); i++ {
-		b := buffer[i]
-		for posByte < b.Len() {
-			if *remainingContent <= 0 && *remainingPadding <= 0 {
-				if *currentCommand == 1 { // possible buffer after padding, no need to worry about xtls (command 2)
-					len := b.Len() - posByte
-					newbuffer := buf.New()
-					newbuffer.Write(b.BytesRange(posByte, posByte+len))
-					mb2 = append(mb2, newbuffer)
-					posByte += len
-				} else {
-					paddingInfo := b.BytesRange(posByte, posByte+5)
-					*currentCommand = int(paddingInfo[0])
-					*remainingContent = int32(paddingInfo[1])<<8 | int32(paddingInfo[2])
-					*remainingPadding = int32(paddingInfo[3])<<8 | int32(paddingInfo[4])
-					newError("Xtls Unpadding new block", i, " ", posByte, " content ", *remainingContent, " padding ", *remainingPadding, " ", paddingInfo[0]).WriteToLog(session.ExportIDToError(ctx))
-					posByte += 5
-				}
-			} else if *remainingContent > 0 {
-				len := *remainingContent
-				if b.Len() < posByte+*remainingContent {
-					len = b.Len() - posByte
-				}
-				newbuffer := buf.New()
-				newbuffer.Write(b.BytesRange(posByte, posByte+len))
-				mb2 = append(mb2, newbuffer)
-				*remainingContent -= len
-				posByte += len
-			} else { // remainingPadding > 0
-				len := *remainingPadding
-				if b.Len() < posByte+*remainingPadding {
-					len = b.Len() - posByte
-				}
-				*remainingPadding -= len
-				posByte += len
+	newbuffer := buf.New()
+	for b.Len() > 0 {
+		if *remainingCommand > 0 {
+			data, err := b.ReadByte()
+			if err != nil {
+				return newbuffer
+			}
+			switch *remainingCommand {
+			case 5:
+				*currentCommand = int(data)
+			case 4:
+				*remainingContent = int32(data)<<8
+			case 3:
+				*remainingContent = *remainingContent | int32(data)
+			case 2:
+				*remainingPadding = int32(data)<<8
+			case 1:
+				*remainingPadding = *remainingPadding | int32(data)
+				newError("Xtls Unpadding new block, content ", *remainingContent, " padding ", *remainingPadding, " command ", *currentCommand).WriteToLog(session.ExportIDToError(ctx))
+			}
+			*remainingCommand--
+		} else if *remainingContent > 0 {
+			len := *remainingContent
+			if b.Len() < len {
+				len = b.Len()
+			}
+			data, err := b.ReadBytes(len)
+			if err != nil {
+				return newbuffer
 			}
-			if posByte == b.Len() {
-				posByte = 0
+			newbuffer.Write(data)
+			*remainingContent -= len
+		} else { // remainingPadding > 0
+			len := *remainingPadding
+			if b.Len() < len {
+				len = b.Len()
+			}
+			b.Advance(len)
+			*remainingPadding -= len
+		}
+		if *remainingCommand <= 0 && *remainingContent <= 0 && *remainingPadding <= 0 { // this block done
+			if *currentCommand == 0 {
+				*remainingCommand = 5
+			} else {
+				*remainingCommand = -1 // set to initial state
+				*remainingContent = -1
+				*remainingPadding = -1
+				if b.Len() > 0 { // shouldn't happen
+					newbuffer.Write(b.Bytes())
+				}
 				break
 			}
 		}
 	}
-	buf.ReleaseMulti(buffer)
-	return mb2
+	b.Release()
+	b = nil
+	return newbuffer
 }