Skip to content

Commit

Permalink
[client] Fix inbound tracking in userspace firewall (#3111)
Browse files Browse the repository at this point in the history
* Don't create state for inbound SYN

* Allow final ack in some cases

* Relax state machine test a little
  • Loading branch information
lixmal authored Dec 25, 2024
1 parent 0dbaddc commit b3c87cb
Show file tree
Hide file tree
Showing 3 changed files with 10 additions and 38 deletions.
1 change: 0 additions & 1 deletion client/firewall/uspfilter/conntrack/common.go
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@ import (

// BaseConnTrack provides common fields and locking for all connection types
type BaseConnTrack struct {
sync.RWMutex
SourceIP net.IP
DestIP net.IP
SourcePort uint16
Expand Down
40 changes: 8 additions & 32 deletions client/firewall/uspfilter/conntrack/tcp.go
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ type TCPConnKey struct {
type TCPConnTrack struct {
BaseConnTrack
State TCPState
sync.RWMutex
}

// TCPTracker manages TCP connection states
Expand Down Expand Up @@ -131,36 +132,8 @@ func (t *TCPTracker) IsValidInbound(srcIP net.IP, dstIP net.IP, srcPort uint16,
return false
}

// Handle new SYN packets
if flags&TCPSyn != 0 && flags&TCPAck == 0 {
key := makeConnKey(dstIP, srcIP, dstPort, srcPort)
t.mutex.Lock()
if _, exists := t.connections[key]; !exists {
// Use preallocated IPs
srcIPCopy := t.ipPool.Get()
dstIPCopy := t.ipPool.Get()
copyIP(srcIPCopy, dstIP)
copyIP(dstIPCopy, srcIP)

conn := &TCPConnTrack{
BaseConnTrack: BaseConnTrack{
SourceIP: srcIPCopy,
DestIP: dstIPCopy,
SourcePort: dstPort,
DestPort: srcPort,
},
State: TCPStateSynReceived,
}
conn.lastSeen.Store(time.Now().UnixNano())
conn.established.Store(false)
t.connections[key] = conn
}
t.mutex.Unlock()
return true
}

// Look up existing connection
key := makeConnKey(dstIP, srcIP, dstPort, srcPort)

t.mutex.RLock()
conn, exists := t.connections[key]
t.mutex.RUnlock()
Expand All @@ -172,8 +145,7 @@ func (t *TCPTracker) IsValidInbound(srcIP net.IP, dstIP net.IP, srcPort uint16,
// Handle RST packets
if flags&TCPRst != 0 {
conn.Lock()
isEstablished := conn.IsEstablished()
if isEstablished || conn.State == TCPStateSynSent || conn.State == TCPStateSynReceived {
if conn.IsEstablished() || conn.State == TCPStateSynSent || conn.State == TCPStateSynReceived {
conn.State = TCPStateClosed
conn.SetEstablished(false)
conn.Unlock()
Expand All @@ -183,7 +155,6 @@ func (t *TCPTracker) IsValidInbound(srcIP net.IP, dstIP net.IP, srcPort uint16,
return false
}

// Update state
conn.Lock()
t.updateState(conn, flags, false)
conn.UpdateLastSeen()
Expand Down Expand Up @@ -306,6 +277,11 @@ func (t *TCPTracker) isValidStateForFlags(state TCPState, flags uint8) bool {
return flags&TCPFin != 0 || flags&TCPAck != 0
case TCPStateLastAck:
return flags&TCPAck != 0
case TCPStateClosed:
// Accept retransmitted ACKs in closed state
// This is important because the final ACK might be lost
// and the peer will retransmit their FIN-ACK
return flags&TCPAck != 0
}
return false
}
Expand Down
7 changes: 2 additions & 5 deletions client/firewall/uspfilter/conntrack/tcp_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -125,11 +125,8 @@ func TestTCPStateMachine(t *testing.T) {
valid := tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPRst)
require.True(t, valid, "RST should be allowed for established connection")

// Verify connection is closed
valid = tracker.IsValidInbound(dstIP, srcIP, dstPort, srcPort, TCPPush|TCPAck)
t.Helper()

require.False(t, valid, "Data should be blocked after RST")
// Connection is logically dead but we don't enforce blocking subsequent packets
// The connection will be cleaned up by timeout
},
},
{
Expand Down

0 comments on commit b3c87cb

Please sign in to comment.