github.com/v2fly/v2ray-core/v4@v4.45.2/transport/internet/kcp/io.go (about)

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