package quic import ( "context" "fmt" "github.com/google/uuid" "github.com/lucas-clemente/quic-go" "github.com/pkg/errors" "github.com/rs/zerolog" ) const ( sessionIDLen = len(uuid.UUID{}) ) type SessionDatagram struct { ID uuid.UUID Payload []byte } type BaseDatagramMuxer interface { // MuxSession suffix the session ID to the payload so the other end of the QUIC connection can demultiplex the // payload from multiple datagram sessions MuxSession(sessionID uuid.UUID, payload []byte) error // ServeReceive starts a loop to receive datagrams from the QUIC connection ServeReceive(ctx context.Context) error } type DatagramMuxer struct { session quic.Connection logger *zerolog.Logger demuxChan chan<- *SessionDatagram } func NewDatagramMuxer(quicSession quic.Connection, log *zerolog.Logger, demuxChan chan<- *SessionDatagram) *DatagramMuxer { logger := log.With().Uint8("datagramVersion", 1).Logger() return &DatagramMuxer{ session: quicSession, logger: &logger, demuxChan: demuxChan, } } // Maximum application payload to send to / receive from QUIC datagram frame func (dm *DatagramMuxer) mtu() int { return maxDatagramPayloadSize } func (dm *DatagramMuxer) MuxSession(sessionID uuid.UUID, payload []byte) error { if len(payload) > dm.mtu() { // TODO: TUN-5302 return ICMP packet too big message // drop packet for now, eventually reply with ICMP for PMTUD return fmt.Errorf("origin UDP payload has %d bytes, which exceeds transport MTU %d", len(payload), dm.mtu()) } payloadWithMetadata, err := suffixSessionID(sessionID, payload) if err != nil { return errors.Wrap(err, "Failed to suffix session ID to datagram, it will be dropped") } if err := dm.session.SendMessage(payloadWithMetadata); err != nil { return errors.Wrap(err, "Failed to send datagram back to edge") } return nil } func (dm *DatagramMuxer) ServeReceive(ctx context.Context) error { for { // Extracts datagram session ID, then sends the session ID and payload to receiver // which determines how to proxy to the origin. It assumes the datagram session has already been // registered with receiver through other side channel msg, err := dm.session.ReceiveMessage() if err != nil { return err } if err := dm.demux(ctx, msg); err != nil { dm.logger.Error().Err(err).Msg("Failed to demux datagram") if err == context.Canceled { return err } } } } func (dm *DatagramMuxer) demux(ctx context.Context, msg []byte) error { sessionID, payload, err := extractSessionID(msg) if err != nil { return err } sessionDatagram := SessionDatagram{ ID: sessionID, Payload: payload, } select { case dm.demuxChan <- &sessionDatagram: return nil case <-ctx.Done(): return ctx.Err() } } // Each QUIC datagram should be suffixed with session ID. // extractSessionID extracts the session ID and a slice with only the payload func extractSessionID(b []byte) (uuid.UUID, []byte, error) { msgLen := len(b) if msgLen < sessionIDLen { return uuid.Nil, nil, fmt.Errorf("session ID has %d bytes, but data only has %d", sessionIDLen, len(b)) } // Parse last 16 bytess as UUID and remove it from slice sessionID, err := uuid.FromBytes(b[len(b)-sessionIDLen:]) if err != nil { return uuid.Nil, nil, err } b = b[:len(b)-sessionIDLen] return sessionID, b, nil } // SuffixSessionID appends the session ID at the end of the payload. Suffix is more performant than prefix because // the payload slice might already have enough capacity to append the session ID at the end func suffixSessionID(sessionID uuid.UUID, b []byte) ([]byte, error) { if len(b)+len(sessionID) > MaxDatagramFrameSize { return nil, fmt.Errorf("datagram size exceed %d", MaxDatagramFrameSize) } b = append(b, sessionID[:]...) return b, nil }