TUN-6667: DatagramMuxerV2 provides a method to receive RawPacket

quic/datagramv2.go

@@ -16,6 +16,8 @@ type datagramV2Type byte
 const (
 	udp datagramV2Type = iota
 	ip
+	// Same as sessionDemuxChan capacity
+	packetChanCapacity = 16
 )
 
 func suffixType(b []byte, datagramType datagramV2Type) ([]byte, error) {
@@ -35,24 +37,24 @@ type DatagramMuxerV2 struct {
 	session          quic.Connection
 	logger           *zerolog.Logger
 	sessionDemuxChan chan<- *packet.Session
-	packetDemuxChan  chan<- []byte
+	packetDemuxChan  chan packet.RawPacket
 }
 
 func NewDatagramMuxerV2(
 	quicSession quic.Connection,
 	log *zerolog.Logger,
 	sessionDemuxChan chan<- *packet.Session,
-	packetDemuxChan chan<- []byte) *DatagramMuxerV2 {
+) *DatagramMuxerV2 {
 	logger := log.With().Uint8("datagramVersion", 2).Logger()
 	return &DatagramMuxerV2{
 		session:          quicSession,
 		logger:           &logger,
 		sessionDemuxChan: sessionDemuxChan,
-		packetDemuxChan:  packetDemuxChan,
+		packetDemuxChan:  make(chan packet.RawPacket, packetChanCapacity),
 	}
 }
 
-// MuxSession suffix the session ID and datagram version to the payload so the other end of the QUIC connection can
+// SendToSession suffix the session ID and datagram version to the payload so the other end of the QUIC connection can
 // demultiplex the payload from multiple datagram sessions
 func (dm *DatagramMuxerV2) SendToSession(session *packet.Session) error {
 	if len(session.Payload) > dm.mtu() {
@@ -73,10 +75,10 @@ func (dm *DatagramMuxerV2) SendToSession(session *packet.Session) error {
 	return nil
 }
 
-// MuxPacket suffix the datagram type to the packet. The other end of the QUIC connection can demultiplex by parsing
+// SendPacket suffix the datagram type to the packet. The other end of the QUIC connection can demultiplex by parsing
 // the payload as IP and look at the source and destination.
-func (dm *DatagramMuxerV2) MuxPacket(packet []byte) error {
-	payloadWithVersion, err := suffixType(packet, ip)
+func (dm *DatagramMuxerV2) SendPacket(pk packet.RawPacket) error {
+	payloadWithVersion, err := suffixType(pk.Data, ip)
 	if err != nil {
 		return errors.Wrap(err, "Failed to suffix datagram type, it will be dropped")
 	}
@@ -102,6 +104,15 @@ func (dm *DatagramMuxerV2) ServeReceive(ctx context.Context) error {
 	}
 }
 
+func (dm *DatagramMuxerV2) ReceivePacket(ctx context.Context) (packet.RawPacket, error) {
+	select {
+	case <-ctx.Done():
+		return packet.RawPacket{}, ctx.Err()
+	case pk := <-dm.packetDemuxChan:
+		return pk, nil
+	}
+}
+
 func (dm *DatagramMuxerV2) demux(ctx context.Context, msgWithType []byte) error {
 	if len(msgWithType) < 1 {
 		return fmt.Errorf("QUIC datagram should have at least 1 byte")
@@ -110,28 +121,38 @@ func (dm *DatagramMuxerV2) demux(ctx context.Context, msgWithType []byte) error
 	msg := msgWithType[0 : len(msgWithType)-1]
 	switch msgType {
 	case udp:
-		sessionID, payload, err := extractSessionID(msg)
-		if err != nil {
-			return err
-		}
-		sessionDatagram := packet.Session{
-			ID:      sessionID,
-			Payload: payload,
-		}
-		select {
-		case dm.sessionDemuxChan <- &sessionDatagram:
-			return nil
-		case <-ctx.Done():
-			return ctx.Err()
-		}
+		return dm.handleSession(ctx, msg)
 	case ip:
-		select {
-		case dm.packetDemuxChan <- msg:
-			return nil
-		case <-ctx.Done():
-			return ctx.Err()
-		}
+		return dm.handlePacket(ctx, msg)
 	default:
 		return fmt.Errorf("Unexpected datagram type %d", msgType)
 	}
 }
+
+func (dm *DatagramMuxerV2) handleSession(ctx context.Context, session []byte) error {
+	sessionID, payload, err := extractSessionID(session)
+	if err != nil {
+		return err
+	}
+	sessionDatagram := packet.Session{
+		ID:      sessionID,
+		Payload: payload,
+	}
+	select {
+	case dm.sessionDemuxChan <- &sessionDatagram:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
+func (dm *DatagramMuxerV2) handlePacket(ctx context.Context, pk []byte) error {
+	select {
+	case <-ctx.Done():
+		return ctx.Err()
+	case dm.packetDemuxChan <- packet.RawPacket{
+		Data: pk,
+	}:
+		return nil
+	}
+}