MapPortHandler Callback for container support

agent.go

@@ -167,6 +167,9 @@ type Agent struct {
 	automaticRenomination bool
 	renominationInterval  time.Duration
 	lastRenominationTime  time.Time
+
+	// Port mapping support for container
+	mapPort func(candidate Candidate) int
 }
 
 // NewAgent creates a new Agent.
@@ -919,6 +922,10 @@ func (a *Agent) addCandidate(ctx context.Context, cand Candidate, candidateConn
 				return
 			}
 		}
+		// Callback for mapPort before candidate starts
+		if a.mapPort != nil {
+			cand.setMappedPort(a.mapPort(cand))
+		}
 
 		a.setCandidateExtensions(cand)
 		cand.start(a, candidateConn, a.startedCh)

agent_options.go

@@ -523,3 +523,17 @@ func WithLoggerFactory(loggerFactory logging.LoggerFactory) AgentOption {
 		return nil
 	}
 }
+
+func WithMapPortHandler(handler func(cand Candidate) int, candTyp CandidateType) AgentOption {
+	return func(a *Agent) error {
+		a.mapPort = func(candidate Candidate) int {
+			if candidate.Type() == candTyp {
+				return handler(candidate)
+			}
+
+			return 0
+		}
+
+		return nil
+	}
+}

agent_test.go

@@ -2299,3 +2299,71 @@ func TestAutomaticRenominationRelayToDirect(t *testing.T) {
 	shouldRenominate := agent.shouldRenominate(relayPair, hostPair)
 	require.True(t, shouldRenominate, "Should always renominate from relay to direct connection")
 }
+
+func TestMapPortHandler(t *testing.T) {
+	handler := func(cand Candidate) int {
+		return cand.Port() + 1000
+	}
+	agent, err := NewAgentWithOptions(WithMapPortHandler(handler, CandidateTypeHost))
+	require.NoError(t, err)
+	defer func() {
+		require.NoError(t, agent.Close())
+	}()
+
+	dummyConn := &net.UDPConn{}
+
+	for i := 0; i < 5; i++ {
+		cfg := CandidateHostConfig{
+			Network:   "udp",
+			Address:   "192.168.0.2",
+			Port:      1000 + i,
+			Component: 1,
+		}
+
+		cand, errCand := NewCandidateHost(&cfg)
+		require.NoError(t, errCand)
+		err = agent.addCandidate(context.Background(), cand, dummyConn)
+		require.NoError(t, err)
+	}
+
+	actualCandidates, err := agent.GetLocalCandidates()
+	require.NoError(t, err)
+
+	for _, candidate := range actualCandidates {
+		require.Equal(t, candidate.Port()+1000, candidate.getMappedPort())
+	}
+}
+
+func TestMapPortHandlerDifferentCandidateType(t *testing.T) {
+	handler := func(cand Candidate) int {
+		return cand.Port() + 1000
+	}
+	agent, err := NewAgentWithOptions(WithMapPortHandler(handler, CandidateTypePeerReflexive))
+	require.NoError(t, err)
+	defer func() {
+		require.NoError(t, agent.Close())
+	}()
+
+	dummyConn := &net.UDPConn{}
+
+	for i := 0; i < 5; i++ {
+		cfg := CandidateHostConfig{
+			Network:   "udp",
+			Address:   "192.168.0.2",
+			Port:      1000 + i,
+			Component: 1,
+		}
+
+		cand, errCand := NewCandidateHost(&cfg)
+		require.NoError(t, errCand)
+		err = agent.addCandidate(context.Background(), cand, dummyConn)
+		require.NoError(t, err)
+	}
+
+	actualCandidates, err := agent.GetLocalCandidates()
+	require.NoError(t, err)
+
+	for i, candidate := range actualCandidates {
+		require.Equal(t, 1000+i, candidate.Port())
+	}
+}

candidate.go

@@ -46,6 +46,10 @@ type Candidate interface {
 	Address() string
 	Port() int
 
+	// Port mapping support for containers
+	getMappedPort() int
+	setMappedPort(port int)
+
 	Priority() uint32
 
 	// A transport address related to a

candidate_base.go

@@ -27,6 +27,7 @@ type candidateBase struct {
 	component      uint16
 	address        string
 	port           int
+	mappedPort     int
 	relatedAddress *CandidateRelatedAddress
 	tcpType        TCPType
 
@@ -96,6 +97,14 @@ func (c *candidateBase) Port() int {
 	return c.port
 }
 
+func (c *candidateBase) getMappedPort() int {
+	return c.mappedPort
+}
+
+func (c *candidateBase) setMappedPort(port int) {
+	c.mappedPort = port
+}
+
 // Type returns candidate type.
 func (c *candidateBase) Type() CandidateType {
 	return c.candidateType
@@ -416,7 +425,7 @@ func (c *candidateBase) Priority() uint32 {
 }
 
 // Equal is used to compare two candidateBases.
-func (c *candidateBase) Equal(other Candidate) bool {
+func (c *candidateBase) Equal(other Candidate) bool { //nolint:cyclop
 	if c.addr() != other.addr() {
 		if c.addr() == nil || other.addr() == nil {
 			return false
@@ -430,6 +439,7 @@ func (c *candidateBase) Equal(other Candidate) bool {
 		c.Type() == other.Type() &&
 		c.Address() == other.Address() &&
 		c.Port() == other.Port() &&
+		c.getMappedPort() == other.getMappedPort() &&
 		c.TCPType() == other.TCPType() &&
 		c.RelatedAddress().Equal(other.RelatedAddress())
 }
@@ -521,14 +531,18 @@ func (c *candidateBase) Marshal() string {
 	if val == " " {
 		val = ""
 	}
+	port := c.Port()
+	if c.mappedPort != 0 {
+		port = c.mappedPort
+	}
 
 	val = fmt.Sprintf("%s %d %s %d %s %d typ %s",
 		val,
 		c.Component(),
 		c.NetworkType().NetworkShort(),
 		c.Priority(),
 		removeZoneIDFromAddress(c.Address()),
-		c.Port(),
+		port,
 		c.Type())
 
 	if r := c.RelatedAddress(); r != nil && r.Address != "" && r.Port != 0 {

candidate_test.go

@@ -352,6 +352,19 @@ func mustCandidatePeerReflexiveWithExtensions(
 	return cand
 }
 
+func mustCandidateHostWithMappedPort(
+	t *testing.T,
+	conf *CandidateHostConfig,
+	mapPort func(cand Candidate) int,
+) Candidate {
+	t.Helper()
+	cand, err := NewCandidateHost(conf)
+	require.NoError(t, err)
+	cand.setMappedPort(mapPort(cand))
+
+	return cand
+}
+
 func TestCandidateMarshal(t *testing.T) {
 	for idx, test := range []struct {
 		candidate   Candidate
@@ -580,6 +593,66 @@ func TestCandidateMarshal(t *testing.T) {
 	}
 }
 
+func TestCandidateMarshalWithMappedPort(t *testing.T) {
+	for idx, test := range []struct {
+		candidate   Candidate
+		marshaled   string
+		expectError bool
+	}{
+		{
+			mustCandidateHostWithMappedPort(t, &CandidateHostConfig{
+				Network:    NetworkTypeTCP4.String(),
+				Address:    "172.28.142.173",
+				Port:       7686,
+				Priority:   1671430143,
+				Foundation: "3359356140",
+			}, func(cand Candidate) int { return 7687 }),
+			"candidate:3359356140 0 tcp 1671430143 172.28.142.173 7687 typ host",
+			false,
+		},
+		{
+			mustCandidateHostWithMappedPort(t, &CandidateHostConfig{
+				Network:    NetworkTypeTCP4.String(),
+				Address:    "172.28.142.173",
+				Port:       7686,
+				Priority:   1671430143,
+				Foundation: "3359356140",
+			}, func(cand Candidate) int {
+				if cand.Port() != 7686 {
+					return 7687
+				}
+
+				return 7688
+			}),
+			"candidate:3359356140 0 tcp 1671430143 172.28.142.173 7688 typ host",
+			false,
+		},
+		{
+			mustCandidateHostWithMappedPort(t, &CandidateHostConfig{
+				Network:    NetworkTypeTCP4.String(),
+				Address:    "172.28.142.173",
+				Port:       7686,
+				Priority:   1671430143,
+				Foundation: "3359356140",
+			}, func(cand Candidate) int {
+				return 0
+			}),
+			"candidate:3359356140 0 tcp 1671430143 172.28.142.173 7686 typ host",
+			false,
+		},
+	} {
+		t.Run(strconv.Itoa(idx), func(t *testing.T) {
+			candidateMarshalOutput := test.candidate.Marshal()
+
+			if strings.HasPrefix(test.marshaled, "candidate:") {
+				require.Equal(t, test.marshaled[len("candidate:"):], candidateMarshalOutput)
+			} else {
+				require.Equal(t, test.marshaled, candidateMarshalOutput)
+			}
+		})
+	}
+}
+
 func TestCandidateWriteTo(t *testing.T) {
 	listener, err := net.ListenTCP("tcp", &net.TCPAddr{
 		IP:   net.IP{127, 0, 0, 1},