2424import org .jetbrains .annotations .NotNull ;
2525
2626import tech .httptoolkit .android .TagKt ;
27+ import tech .httptoolkit .android .vpn .capture .CaptureController ;
28+ import tech .httptoolkit .android .vpn .capture .ProxyProtocolHandler ;
2729import tech .httptoolkit .android .vpn .socket .DataConst ;
2830import tech .httptoolkit .android .vpn .socket .ICloseSession ;
2931import tech .httptoolkit .android .vpn .socket .SocketProtector ;
3638import java .nio .channels .DatagramChannel ;
3739import java .nio .channels .SocketChannel ;
3840import java .nio .channels .spi .AbstractSelectableChannel ;
41+ import java .util .ArrayList ;
42+ import java .util .List ;
3943import java .util .Map ;
4044import java .util .concurrent .ConcurrentHashMap ;
4145
@@ -50,6 +54,12 @@ public class SessionManager implements ICloseSession {
5054 private final Map <String , Session > table = new ConcurrentHashMap <>();
5155 private SocketProtector protector = SocketProtector .getInstance ();
5256
57+ private final CaptureController captureController ;
58+
59+ public SessionManager (CaptureController captureController ) {
60+ this .captureController = captureController ;
61+ }
62+
5363 /**
5464 * keep java garbage collector from collecting a session
5565 * @param session Session
@@ -133,7 +143,7 @@ public Session createNewUDPSession(int ip, int port, int srcIp, int srcPort) thr
133143 Session existingSession = table .get (keys );
134144 if (existingSession != null ) return existingSession ;
135145
136- Session session = new Session (SessionProtocol .UDP , srcIp , srcPort , ip , port , this );
146+ Session session = new Session (SessionProtocol .UDP , srcIp , srcPort , ip , port , this , null );
137147
138148 DatagramChannel channel ;
139149
@@ -171,7 +181,17 @@ public Session createNewTCPSession(int ip, int port, int srcIp, int srcPort) thr
171181 // We return the initialized session, which will be reacked to indicate rejection.
172182 if (existingSession != null ) return existingSession ;
173183
174- Session session = new Session (SessionProtocol .TCP , srcIp , srcPort , ip , port , this );
184+ String ips = PacketUtil .intToIPAddress (ip );
185+ boolean shouldCapture = captureController .shouldCapture (ips , port );
186+ SocketAddress socketAddress = shouldCapture
187+ ? captureController .getProxyAddress ()
188+ : new InetSocketAddress (ips , port );
189+
190+ ProxyProtocolHandler proxyHandler = shouldCapture
191+ ? captureController .getProxyHandler (ips , port )
192+ : null ;
193+
194+ Session session = new Session (SessionProtocol .TCP , srcIp , srcPort , ip , port , this , proxyHandler );
175195
176196 SocketChannel channel ;
177197 channel = SocketChannel .open ();
@@ -181,7 +201,6 @@ public Session createNewTCPSession(int ip, int port, int srcIp, int srcPort) thr
181201 channel .socket ().setReceiveBufferSize (DataConst .MAX_RECEIVE_BUFFER_SIZE );
182202 channel .configureBlocking (false );
183203
184- String ips = PacketUtil .intToIPAddress (ip );
185204 Log .d (TAG ,"created new SocketChannel for " + key );
186205
187206 protector .protect (channel .socket ());
@@ -190,12 +209,6 @@ public Session createNewTCPSession(int ip, int port, int srcIp, int srcPort) thr
190209 session .setChannel (channel );
191210
192211 // Initiate connection straight away, to reduce latency
193- // We use the real address, unless tcpPortRedirection redirects us to a different
194- // target address for traffic on this port.
195- SocketAddress socketAddress = tcpPortRedirection .get (port ) != null
196- ? tcpPortRedirection .get (port )
197- : new InetSocketAddress (ips , port );
198-
199212 Log .d (TAG ,"Initiate connecting to remote tcp server: " + socketAddress .toString ());
200213 boolean connected = channel .connect (socketAddress );
201214 session .setConnected (connected );
@@ -205,9 +218,4 @@ public Session createNewTCPSession(int ip, int port, int srcIp, int srcPort) thr
205218 return session ;
206219 }
207220
208- private SparseArray <InetSocketAddress > tcpPortRedirection = new SparseArray <>();
209-
210- public void setTcpPortRedirections (SparseArray <InetSocketAddress > tcpPortRedirection ) {
211- this .tcpPortRedirection = tcpPortRedirection ;
212- }
213221}
0 commit comments