@@ -157,7 +157,7 @@ func NewTCPTracker(timeout time.Duration, logger *nblog.Logger, flowLogger nftyp
157157 return tracker
158158}
159159
160- func (t * TCPTracker ) updateIfExists (srcIP , dstIP netip.Addr , srcPort , dstPort uint16 , flags uint8 , direction nftypes.Direction , size int ) (ConnKey , bool ) {
160+ func (t * TCPTracker ) updateIfExists (srcIP , dstIP netip.Addr , srcPort , dstPort uint16 , flags uint8 , direction nftypes.Direction , size int ) (ConnKey , uint16 , bool ) {
161161 key := ConnKey {
162162 SrcIP : srcIP ,
163163 DstIP : dstIP ,
@@ -171,28 +171,30 @@ func (t *TCPTracker) updateIfExists(srcIP, dstIP netip.Addr, srcPort, dstPort ui
171171
172172 if exists {
173173 t .updateState (key , conn , flags , direction , size )
174- return key , true
174+ return key , uint16 ( conn . DNATOrigPort . Load ()), true
175175 }
176176
177- return key , false
177+ return key , 0 , false
178178}
179179
180- // TrackOutbound records an outbound TCP connection
181- func (t * TCPTracker ) TrackOutbound (srcIP , dstIP netip.Addr , srcPort , dstPort uint16 , flags uint8 , size int ) {
182- if _ , exists := t .updateIfExists (dstIP , srcIP , dstPort , srcPort , flags , nftypes .Egress , size ); ! exists {
183- // if (inverted direction) conn is not tracked, track this direction
184- t .track (srcIP , dstIP , srcPort , dstPort , flags , nftypes .Egress , nil , size )
180+ // TrackOutbound records an outbound TCP connection and returns the original port if DNAT reversal is needed
181+ func (t * TCPTracker ) TrackOutbound (srcIP , dstIP netip.Addr , srcPort , dstPort uint16 , flags uint8 , size int ) uint16 {
182+ if _ , origPort , exists := t .updateIfExists (dstIP , srcIP , dstPort , srcPort , flags , nftypes .Egress , size ); exists {
183+ return origPort
185184 }
185+ // if (inverted direction) conn is not tracked, track this direction
186+ t .track (srcIP , dstIP , srcPort , dstPort , flags , nftypes .Egress , nil , size , 0 )
187+ return 0
186188}
187189
188190// TrackInbound processes an inbound TCP packet and updates connection state
189- func (t * TCPTracker ) TrackInbound (srcIP , dstIP netip.Addr , srcPort , dstPort uint16 , flags uint8 , ruleID []byte , size int ) {
190- t .track (srcIP , dstIP , srcPort , dstPort , flags , nftypes .Ingress , ruleID , size )
191+ func (t * TCPTracker ) TrackInbound (srcIP , dstIP netip.Addr , srcPort , dstPort uint16 , flags uint8 , ruleID []byte , size int , dnatOrigPort uint16 ) {
192+ t .track (srcIP , dstIP , srcPort , dstPort , flags , nftypes .Ingress , ruleID , size , dnatOrigPort )
191193}
192194
193195// track is the common implementation for tracking both inbound and outbound connections
194- func (t * TCPTracker ) track (srcIP , dstIP netip.Addr , srcPort , dstPort uint16 , flags uint8 , direction nftypes.Direction , ruleID []byte , size int ) {
195- key , exists := t .updateIfExists (srcIP , dstIP , srcPort , dstPort , flags , direction , size )
196+ func (t * TCPTracker ) track (srcIP , dstIP netip.Addr , srcPort , dstPort uint16 , flags uint8 , direction nftypes.Direction , ruleID []byte , size int , origPort uint16 ) {
197+ key , _ , exists := t .updateIfExists (srcIP , dstIP , srcPort , dstPort , flags , direction , size )
196198 if exists || flags & TCPSyn == 0 {
197199 return
198200 }
@@ -210,8 +212,13 @@ func (t *TCPTracker) track(srcIP, dstIP netip.Addr, srcPort, dstPort uint16, fla
210212
211213 conn .tombstone .Store (false )
212214 conn .state .Store (int32 (TCPStateNew ))
215+ conn .DNATOrigPort .Store (uint32 (origPort ))
213216
214- t .logger .Trace2 ("New %s TCP connection: %s" , direction , key )
217+ if origPort != 0 {
218+ t .logger .Trace4 ("New %s TCP connection: %s (port DNAT %d -> %d)" , direction , key , origPort , dstPort )
219+ } else {
220+ t .logger .Trace2 ("New %s TCP connection: %s" , direction , key )
221+ }
215222 t .updateState (key , conn , flags , direction , size )
216223
217224 t .mutex .Lock ()
@@ -449,6 +456,21 @@ func (t *TCPTracker) cleanup() {
449456 }
450457}
451458
459+ // GetConnection safely retrieves a connection state
460+ func (t * TCPTracker ) GetConnection (srcIP netip.Addr , srcPort uint16 , dstIP netip.Addr , dstPort uint16 ) (* TCPConnTrack , bool ) {
461+ t .mutex .RLock ()
462+ defer t .mutex .RUnlock ()
463+
464+ key := ConnKey {
465+ SrcIP : srcIP ,
466+ DstIP : dstIP ,
467+ SrcPort : srcPort ,
468+ DstPort : dstPort ,
469+ }
470+ conn , exists := t .connections [key ]
471+ return conn , exists
472+ }
473+
452474// Close stops the cleanup routine and releases resources
453475func (t * TCPTracker ) Close () {
454476 t .tickerCancel ()
0 commit comments