// SPDX-FileCopyrightText: 2023 The Pion community <https://pion.ly>
// SPDX-License-Identifier: MIT

package ice

import (
	
	
	

	
	
	
)

// UniversalUDPMux allows multiple connections to go over a single UDP port for
// host, server reflexive and relayed candidates.
// Actual connection muxing is happening in the UDPMux.
type UniversalUDPMux interface {
	UDPMux
	GetXORMappedAddr(stunAddr net.Addr, deadline time.Duration) (*stun.XORMappedAddress, error)
	GetRelayedAddr(turnAddr net.Addr, deadline time.Duration) (*net.Addr, error)
	GetConnForURL(ufrag string, url string, addr net.Addr) (net.PacketConn, error)
}

// UniversalUDPMuxDefault handles STUN and TURN servers packets by wrapping the original UDPConn overriding ReadFrom.
// It the passes packets to the UDPMux that does the actual connection muxing.
type UniversalUDPMuxDefault struct {
	*UDPMuxDefault
	params UniversalUDPMuxParams

	// Since we have a shared socket, for srflx candidates it makes sense
	// to have a shared mapped address across all the agents
	// stun.XORMappedAddress indexed by the STUN server addr
	xorMappedMap map[string]*xorMapped
}

// UniversalUDPMuxParams are parameters for UniversalUDPMux server reflexive.
type UniversalUDPMuxParams struct {
	Logger                logging.LeveledLogger
	UDPConn               net.PacketConn
	XORMappedAddrCacheTTL time.Duration
	Net                   transport.Net
}

// NewUniversalUDPMuxDefault creates an implementation of UniversalUDPMux embedding UDPMux.
func ( UniversalUDPMuxParams) *UniversalUDPMuxDefault {
	if .Logger == nil {
		.Logger = logging.NewDefaultLoggerFactory().NewLogger("ice")
	}
	if .XORMappedAddrCacheTTL == 0 {
		.XORMappedAddrCacheTTL = time.Second * 25
	}

	 := &UniversalUDPMuxDefault{
		params:       ,
		xorMappedMap: make(map[string]*xorMapped),
	}

	// Wrap UDP connection, process server reflexive messages
	// before they are passed to the UDPMux connection handler (connWorker)
	.params.UDPConn = &udpConn{
		PacketConn: .UDPConn,
		mux:        ,
		logger:     .Logger,
	}

	// Embed UDPMux
	 := UDPMuxParams{
		Logger:  .Logger,
		UDPConn: .params.UDPConn,
		Net:     .params.Net,
	}
	.UDPMuxDefault = NewUDPMuxDefault()

	return 
}

// udpConn is a wrapper around UDPMux conn that overrides ReadFrom and handles STUN/TURN packets.
type udpConn struct {
	net.PacketConn
	mux    *UniversalUDPMuxDefault
	logger logging.LeveledLogger
}

// GetRelayedAddr creates relayed connection to the given TURN service and returns the relayed addr.
// Not implemented yet.
func ( *UniversalUDPMuxDefault) (net.Addr, time.Duration) (*net.Addr, error) {
	return nil, errNotImplemented
}

// GetConnForURL add uniques to the muxed connection by concatenating ufrag and URL
// (e.g. STUN URL) to be able to support multiple STUN/TURN servers
// and return a unique connection per server.
func ( *UniversalUDPMuxDefault) ( string,  string,  net.Addr) (net.PacketConn, error) {
	return .UDPMuxDefault.GetConn(fmt.Sprintf("%s%s", , ), )
}

// ReadFrom is called by UDPMux connWorker and handles packets coming from the STUN server discovering a mapped address.
// It passes processed packets further to the UDPMux (maybe this is not really necessary).
func ( *udpConn) ( []byte) ( int,  net.Addr,  error) {
	, ,  = .PacketConn.ReadFrom()
	if  != nil {
		return , , 
	}

	if stun.IsMessage([:]) { //nolint:nestif
		 := &stun.Message{
			Raw: append([]byte{}, [:]...),
		}

		if  = .Decode();  != nil {
			.logger.Warnf("Failed to handle decode ICE from %s: %v", .String(), )

			return , , nil
		}

		,  := .(*net.UDPAddr)
		if ! {
			// Message about this err will be logged in the UDPMux
			return , , 
		}

		if .mux.isXORMappedResponse(, .String()) {
			 = .mux.handleXORMappedResponse(, )
			if  != nil {
				.logger.Debugf("%w: %v", errGetXorMappedAddrResponse, )
				 = nil
			}

			return , , 
		}
	}

	return , , 
}

// isXORMappedResponse indicates whether the message is a XORMappedAddress and is coming from the known STUN server.
func ( *UniversalUDPMuxDefault) ( *stun.Message,  string) bool {
	.mu.Lock()
	defer .mu.Unlock()
	// Check first if it is a STUN server address,
	// because remote peer can also send similar messages but as a BindingSuccess.
	,  := .xorMappedMap[]
	,  := .Get(stun.AttrXORMappedAddress)

	return  == nil && 
}

// handleXORMappedResponse parses response from the STUN server, extracts XORMappedAddress attribute.
// and set the mapped address for the server.
func ( *UniversalUDPMuxDefault) ( *net.UDPAddr,  *stun.Message) error {
	.mu.Lock()
	defer .mu.Unlock()

	,  := .xorMappedMap[.String()]
	if ! {
		return errNoXorAddrMapping
	}

	var  stun.XORMappedAddress
	if  := .GetFrom();  != nil {
		return 
	}

	.xorMappedMap[.String()] = 
	.SetAddr(&)

	return nil
}

// GetXORMappedAddr returns *stun.XORMappedAddress if already present for a given STUN server.
// Makes a STUN binding request to discover mapped address otherwise.
// Blocks until the stun.XORMappedAddress has been discovered or deadline.
// Method is safe for concurrent use.
func ( *UniversalUDPMuxDefault) (
	 net.Addr,
	 time.Duration,
) (*stun.XORMappedAddress, error) {
	.mu.Lock()
	,  := .xorMappedMap[.String()]
	// If we already have a mapping for this STUN server (address already received)
	// and if it is not too old we return it without making a new request to STUN server
	if  {
		if .expired() {
			.closeWaiters()
			delete(.xorMappedMap, .String())
			 = false
		} else if .pending() {
			 = false
		}
	}
	.mu.Unlock()
	if  {
		return .addr, nil
	}

	// Otherwise, make a STUN request to discover the address
	// or wait for already sent request to complete
	,  := .writeSTUN()
	if  != nil {
		return nil, fmt.Errorf("%w: %s", errWriteSTUNMessage, ) //nolint:errorlint
	}

	// Block until response was handled by the connWorker routine and XORMappedAddress was updated
	select {
	case <-:
		// When channel closed, addr was obtained
		.mu.Lock()
		 := *.xorMappedMap[.String()]
		.mu.Unlock()
		if .addr == nil {
			return nil, errNoXorAddrMapping
		}

		return .addr, nil
	case <-time.After():
		return nil, errXORMappedAddrTimeout
	}
}

// writeSTUN sends a STUN request via UDP conn.
//
// The returned channel is closed when the STUN response has been received.
// Method is safe for concurrent use.
func ( *UniversalUDPMuxDefault) ( net.Addr) (chan struct{}, error) {
	.mu.Lock()
	defer .mu.Unlock()

	// If record present in the map, we already sent a STUN request,
	// just wait when waitAddrReceived will be closed
	,  := .xorMappedMap[.String()]
	if ! {
		 = &xorMapped{
			expiresAt:        time.Now().Add(.params.XORMappedAddrCacheTTL),
			waitAddrReceived: make(chan struct{}),
		}
		.xorMappedMap[.String()] = 
	}

	,  := stun.Build(stun.BindingRequest, stun.TransactionID)
	if  != nil {
		return nil, 
	}

	if _,  = .params.UDPConn.WriteTo(.Raw, );  != nil {
		return nil, 
	}

	return .waitAddrReceived, nil
}

type xorMapped struct {
	addr             *stun.XORMappedAddress
	waitAddrReceived chan struct{}
	expiresAt        time.Time
}

func ( *xorMapped) () {
	select {
	case <-.waitAddrReceived:
		// Notify was close, ok, that means we received duplicate response just exit
		break
	default:
		// Notify tha twe have a new addr
		close(.waitAddrReceived)
	}
}

func ( *xorMapped) () bool {
	return .addr == nil
}

func ( *xorMapped) () bool {
	return .expiresAt.Before(time.Now())
}

func ( *xorMapped) ( *stun.XORMappedAddress) {
	.addr = 
	.closeWaiters()
}