github.com/Uhtred009/v2ray-core-1@v4.31.2+incompatible/transport/internet/kcp/io.go (about)

     1  // +build !confonly
     2  
     3  package kcp
     4  
     5  import (
     6  	"crypto/cipher"
     7  	"crypto/rand"
     8  	"io"
     9  
    10  	"v2ray.com/core/common"
    11  	"v2ray.com/core/common/buf"
    12  	"v2ray.com/core/transport/internet"
    13  )
    14  
    15  type PacketReader interface {
    16  	Read([]byte) []Segment
    17  }
    18  
    19  type PacketWriter interface {
    20  	Overhead() int
    21  	io.Writer
    22  }
    23  
    24  type KCPPacketReader struct {
    25  	Security cipher.AEAD
    26  	Header   internet.PacketHeader
    27  }
    28  
    29  func (r *KCPPacketReader) Read(b []byte) []Segment {
    30  	if r.Header != nil {
    31  		if int32(len(b)) <= r.Header.Size() {
    32  			return nil
    33  		}
    34  		b = b[r.Header.Size():]
    35  	}
    36  	if r.Security != nil {
    37  		nonceSize := r.Security.NonceSize()
    38  		overhead := r.Security.Overhead()
    39  		if len(b) <= nonceSize+overhead {
    40  			return nil
    41  		}
    42  		out, err := r.Security.Open(b[nonceSize:nonceSize], b[:nonceSize], b[nonceSize:], nil)
    43  		if err != nil {
    44  			return nil
    45  		}
    46  		b = out
    47  	}
    48  	var result []Segment
    49  	for len(b) > 0 {
    50  		seg, x := ReadSegment(b)
    51  		if seg == nil {
    52  			break
    53  		}
    54  		result = append(result, seg)
    55  		b = x
    56  	}
    57  	return result
    58  }
    59  
    60  type KCPPacketWriter struct {
    61  	Header   internet.PacketHeader
    62  	Security cipher.AEAD
    63  	Writer   io.Writer
    64  }
    65  
    66  func (w *KCPPacketWriter) Overhead() int {
    67  	overhead := 0
    68  	if w.Header != nil {
    69  		overhead += int(w.Header.Size())
    70  	}
    71  	if w.Security != nil {
    72  		overhead += w.Security.Overhead()
    73  	}
    74  	return overhead
    75  }
    76  
    77  func (w *KCPPacketWriter) Write(b []byte) (int, error) {
    78  	bb := buf.StackNew()
    79  	defer bb.Release()
    80  
    81  	if w.Header != nil {
    82  		w.Header.Serialize(bb.Extend(w.Header.Size()))
    83  	}
    84  	if w.Security != nil {
    85  		nonceSize := w.Security.NonceSize()
    86  		common.Must2(bb.ReadFullFrom(rand.Reader, int32(nonceSize)))
    87  		nonce := bb.BytesFrom(int32(-nonceSize))
    88  
    89  		encrypted := bb.Extend(int32(w.Security.Overhead() + len(b)))
    90  		w.Security.Seal(encrypted[:0], nonce, b, nil)
    91  	} else {
    92  		bb.Write(b)
    93  	}
    94  
    95  	_, err := w.Writer.Write(bb.Bytes())
    96  	return len(b), err
    97  }