github.com/v2fly/v2ray-core/v5@v5.16.2-0.20240507031116-8191faa6e095/proxy/shadowsocks2022/udp_aes.go (about)

     1  package shadowsocks2022
     2  
     3  import (
     4  	"bytes"
     5  	"crypto/cipher"
     6  	"io"
     7  
     8  	"github.com/lunixbochs/struc"
     9  
    10  	"github.com/v2fly/v2ray-core/v5/common/buf"
    11  	"github.com/v2fly/v2ray-core/v5/common/net"
    12  )
    13  
    14  type AESUDPClientPacketProcessor struct {
    15  	requestSeparateHeaderBlockCipher  cipher.Block
    16  	responseSeparateHeaderBlockCipher cipher.Block
    17  	mainPacketAEAD                    func([]byte) cipher.AEAD
    18  	EIHGenerator                      func([]byte) ExtensibleIdentityHeaders
    19  }
    20  
    21  func NewAESUDPClientPacketProcessor(requestSeparateHeaderBlockCipher, responseSeparateHeaderBlockCipher cipher.Block, mainPacketAEAD func([]byte) cipher.AEAD, eih func([]byte) ExtensibleIdentityHeaders) *AESUDPClientPacketProcessor {
    22  	return &AESUDPClientPacketProcessor{
    23  		requestSeparateHeaderBlockCipher:  requestSeparateHeaderBlockCipher,
    24  		responseSeparateHeaderBlockCipher: responseSeparateHeaderBlockCipher,
    25  		mainPacketAEAD:                    mainPacketAEAD,
    26  		EIHGenerator:                      eih,
    27  	}
    28  }
    29  
    30  type separateHeader struct {
    31  	SessionID [8]byte
    32  	PacketID  uint64
    33  }
    34  
    35  type header struct {
    36  	Type          byte
    37  	TimeStamp     uint64
    38  	PaddingLength uint16 `struc:"sizeof=Padding"`
    39  	Padding       []byte
    40  }
    41  
    42  type respHeader struct {
    43  	Type            byte
    44  	TimeStamp       uint64
    45  	ClientSessionID [8]byte
    46  	PaddingLength   uint16 `struc:"sizeof=Padding"`
    47  	Padding         []byte
    48  }
    49  
    50  type cachedUDPState struct {
    51  	sessionAEAD     cipher.AEAD
    52  	sessionRecvAEAD cipher.AEAD
    53  }
    54  
    55  func (p *AESUDPClientPacketProcessor) EncodeUDPRequest(request *UDPRequest, out *buf.Buffer,
    56  	cache UDPClientPacketProcessorCachedStateContainer,
    57  ) error {
    58  	separateHeaderStruct := separateHeader{PacketID: request.PacketID, SessionID: request.SessionID}
    59  	separateHeaderBuffer := buf.New()
    60  	defer separateHeaderBuffer.Release()
    61  	{
    62  		err := struc.Pack(separateHeaderBuffer, &separateHeaderStruct)
    63  		if err != nil {
    64  			return newError("failed to pack separateHeader").Base(err)
    65  		}
    66  	}
    67  	separateHeaderBufferBytes := separateHeaderBuffer.Bytes()
    68  	{
    69  		encryptedDest := out.Extend(16)
    70  		p.requestSeparateHeaderBlockCipher.Encrypt(encryptedDest, separateHeaderBufferBytes)
    71  	}
    72  
    73  	if p.EIHGenerator != nil {
    74  		eih := p.EIHGenerator(separateHeaderBufferBytes[0:16])
    75  		eihHeader := struct {
    76  			EIH ExtensibleIdentityHeaders
    77  		}{
    78  			EIH: eih,
    79  		}
    80  		err := struc.Pack(out, &eihHeader)
    81  		if err != nil {
    82  			return newError("failed to pack eih").Base(err)
    83  		}
    84  	}
    85  
    86  	headerStruct := header{
    87  		Type:          UDPHeaderTypeClientToServerStream,
    88  		TimeStamp:     request.TimeStamp,
    89  		PaddingLength: 0,
    90  		Padding:       nil,
    91  	}
    92  	requestBodyBuffer := buf.New()
    93  	{
    94  		err := struc.Pack(requestBodyBuffer, &headerStruct)
    95  		if err != nil {
    96  			return newError("failed to header").Base(err)
    97  		}
    98  	}
    99  	{
   100  		err := addrParser.WriteAddressPort(requestBodyBuffer, request.Address, net.Port(request.Port))
   101  		if err != nil {
   102  			return newError("failed to write address port").Base(err)
   103  		}
   104  	}
   105  	{
   106  		_, err := io.Copy(requestBodyBuffer, bytes.NewReader(request.Payload.Bytes()))
   107  		if err != nil {
   108  			return newError("failed to copy payload").Base(err)
   109  		}
   110  	}
   111  	{
   112  		cacheKey := string(separateHeaderBufferBytes[0:8])
   113  		receivedCacheInterface := cache.GetCachedState(cacheKey)
   114  		cachedState := &cachedUDPState{}
   115  		if receivedCacheInterface != nil {
   116  			cachedState = receivedCacheInterface.(*cachedUDPState)
   117  		}
   118  		if cachedState.sessionAEAD == nil {
   119  			cachedState.sessionAEAD = p.mainPacketAEAD(separateHeaderBufferBytes[0:8])
   120  			cache.PutCachedState(cacheKey, cachedState)
   121  		}
   122  
   123  		mainPacketAEADMaterialized := cachedState.sessionAEAD
   124  
   125  		encryptedDest := out.Extend(int32(mainPacketAEADMaterialized.Overhead()) + requestBodyBuffer.Len())
   126  		mainPacketAEADMaterialized.Seal(encryptedDest[:0], separateHeaderBuffer.Bytes()[4:16], requestBodyBuffer.Bytes(), nil)
   127  	}
   128  	return nil
   129  }
   130  
   131  func (p *AESUDPClientPacketProcessor) DecodeUDPResp(input []byte, resp *UDPResponse,
   132  	cache UDPClientPacketProcessorCachedStateContainer,
   133  ) error {
   134  	separateHeaderBuffer := buf.New()
   135  	defer separateHeaderBuffer.Release()
   136  	{
   137  		encryptedDest := separateHeaderBuffer.Extend(16)
   138  		p.responseSeparateHeaderBlockCipher.Decrypt(encryptedDest, input)
   139  	}
   140  	separateHeaderStruct := separateHeader{}
   141  	{
   142  		err := struc.Unpack(separateHeaderBuffer, &separateHeaderStruct)
   143  		if err != nil {
   144  			return newError("failed to unpack separateHeader").Base(err)
   145  		}
   146  	}
   147  	resp.PacketID = separateHeaderStruct.PacketID
   148  	resp.SessionID = separateHeaderStruct.SessionID
   149  	{
   150  		cacheKey := string(separateHeaderBuffer.Bytes()[0:8])
   151  		receivedCacheInterface := cache.GetCachedServerState(cacheKey)
   152  		cachedState := &cachedUDPState{}
   153  		if receivedCacheInterface != nil {
   154  			cachedState = receivedCacheInterface.(*cachedUDPState)
   155  		}
   156  
   157  		if cachedState.sessionRecvAEAD == nil {
   158  			cachedState.sessionRecvAEAD = p.mainPacketAEAD(separateHeaderBuffer.Bytes()[0:8])
   159  			cache.PutCachedServerState(cacheKey, cachedState)
   160  		}
   161  
   162  		mainPacketAEADMaterialized := cachedState.sessionRecvAEAD
   163  		decryptedDestBuffer := buf.New()
   164  		decryptedDest := decryptedDestBuffer.Extend(int32(len(input)) - 16 - int32(mainPacketAEADMaterialized.Overhead()))
   165  		_, err := mainPacketAEADMaterialized.Open(decryptedDest[:0], separateHeaderBuffer.Bytes()[4:16], input[16:], nil)
   166  		if err != nil {
   167  			return newError("failed to open main packet").Base(err)
   168  		}
   169  		decryptedDestReader := bytes.NewReader(decryptedDest)
   170  		headerStruct := respHeader{}
   171  		{
   172  			err := struc.Unpack(decryptedDestReader, &headerStruct)
   173  			if err != nil {
   174  				return newError("failed to unpack header").Base(err)
   175  			}
   176  		}
   177  		resp.TimeStamp = headerStruct.TimeStamp
   178  		addressReaderBuf := buf.New()
   179  		defer addressReaderBuf.Release()
   180  		var port net.Port
   181  		resp.Address, port, err = addrParser.ReadAddressPort(addressReaderBuf, decryptedDestReader)
   182  		if err != nil {
   183  			return newError("failed to read address port").Base(err)
   184  		}
   185  		resp.Port = int(port)
   186  		readedLength := decryptedDestReader.Size() - int64(decryptedDestReader.Len())
   187  		decryptedDestBuffer.Advance(int32(readedLength))
   188  		resp.Payload = decryptedDestBuffer
   189  		resp.ClientSessionID = headerStruct.ClientSessionID
   190  		return nil
   191  	}
   192  }