From 412a5ed445e9c3e686f67b98920ae65191ef9a04 Mon Sep 17 00:00:00 2001
From: "Jason A. Donenfeld" <Jason@zx2c4.com>
Date: Tue, 25 Sep 2018 02:35:01 +0200
Subject: [PATCH] Keep tabs on memory usage

Signed-off-by: Jason A. Donenfeld <Jason@zx2c4.com>
---
 .../PacketTunnelProvider.swift                   | 10 ++++++----
 wireguard-go                                     |  2 +-
 wireguard-go-bridge/Makefile                     |  2 +-
 .../src/git.zx2c4.com/wireguard-go/api-ios.go    |  1 -
 .../git.zx2c4.com/wireguard-go/queueconstants.go | 16 ++++++++++++++++
 5 files changed, 24 insertions(+), 7 deletions(-)
 create mode 100644 wireguard-go-bridge/src/git.zx2c4.com/wireguard-go/queueconstants.go

diff --git a/WireGuardNetworkExtension/PacketTunnelProvider.swift b/WireGuardNetworkExtension/PacketTunnelProvider.swift
index e884b8f..cb8d092 100644
--- a/WireGuardNetworkExtension/PacketTunnelProvider.swift
+++ b/WireGuardNetworkExtension/PacketTunnelProvider.swift
@@ -136,6 +136,7 @@ class PacketTunnelProvider: NEPacketTunnelProvider {
         return withStringsAsGoStrings(interfaceName, settings) { (nameGoStr, settingsGoStr) -> Int32 in
             return withUnsafeMutablePointer(to: &wgContext) { (wgCtxPtr) -> Int32 in
                 return wgTurnOn(nameGoStr, settingsGoStr, { (wgCtxPtr, buf, len) -> Int in
+                    autoreleasepool {
                         // read_fn: Read from the TUN interface and pass it on to WireGuard
                         guard let wgCtxPtr = wgCtxPtr else { return 0 }
                         guard let buf = buf else { return 0 }
@@ -149,7 +150,9 @@ class PacketTunnelProvider: NEPacketTunnelProvider {
                             return packetData.count
                         }
                         return 0
+                    }
                 }, { (wgCtxPtr, buf, len) -> Int in
+                    autoreleasepool {
                         // write_fn: Receive packets from WireGuard and write to the TUN interface
                         guard let wgCtxPtr = wgCtxPtr else { return 0 }
                         guard let buf = buf else { return 0 }
@@ -170,6 +173,7 @@ class PacketTunnelProvider: NEPacketTunnelProvider {
                             return len
                         }
                         return 0
+                    }
                 },
                     wgCtxPtr)
             }
@@ -203,18 +207,16 @@ class WireGuardContext {
         if outboundPackets.isEmpty {
             readPacketCondition.lock()
             packetFlow.readPacketObjects(completionHandler: packetsRead)
-            // Wait till the completion handler of packetFlow.readPacketObjects() finishes
             while outboundPackets.isEmpty && !self.isTunnelClosed {
                 readPacketCondition.wait()
             }
             readPacketCondition.unlock()
         }
         isTunnelClosed = self.isTunnelClosed
-        if outboundPackets.isEmpty {
-            return nil
-        } else {
+        if !outboundPackets.isEmpty {
             return outboundPackets.removeFirst()
         }
+        return nil
     }
 
     func writePacket(packet: NEPacket, isTunnelClosed: inout Bool) -> Bool {
diff --git a/wireguard-go b/wireguard-go
index ebc7541..70bcf9e 160000
--- a/wireguard-go
+++ b/wireguard-go
@@ -1 +1 @@
-Subproject commit ebc7541953269b39cd73d703166b9b8ee7b34e37
+Subproject commit 70bcf9ecb801dadd82c68143209ca2707aa63d2b
diff --git a/wireguard-go-bridge/Makefile b/wireguard-go-bridge/Makefile
index 813feb4..95cdfec 100644
--- a/wireguard-go-bridge/Makefile
+++ b/wireguard-go-bridge/Makefile
@@ -2,7 +2,7 @@
 #
 # Copyright (C) 2018 Jason A. Donenfeld <Jason@zx2c4.com>. All Rights Reserved.
 
-FILES := $(filter-out %/main.go,$(wildcard ../wireguard-go/*/*.go) $(wildcard ../wireguard-go/*.go))
+FILES := $(filter-out %/main.go %/queueconstants.go,$(wildcard ../wireguard-go/*/*.go) $(wildcard ../wireguard-go/*.go))
 
 ARCHES := arm64 armv7 x86_64
 GOARCH_arm64 := arm64
diff --git a/wireguard-go-bridge/src/git.zx2c4.com/wireguard-go/api-ios.go b/wireguard-go-bridge/src/git.zx2c4.com/wireguard-go/api-ios.go
index 753edf4..dcc588f 100644
--- a/wireguard-go-bridge/src/git.zx2c4.com/wireguard-go/api-ios.go
+++ b/wireguard-go-bridge/src/git.zx2c4.com/wireguard-go/api-ios.go
@@ -52,7 +52,6 @@ var tunnelHandles map[int32]*Device
 
 func init() {
 	versionString = C.CString(WireGuardGoVersion)
-	preallocatedBuffers = 64
 	roamingDisabled = true
 	tunnelHandles = make(map[int32]*Device)
 	signals := make(chan os.Signal)
diff --git a/wireguard-go-bridge/src/git.zx2c4.com/wireguard-go/queueconstants.go b/wireguard-go-bridge/src/git.zx2c4.com/wireguard-go/queueconstants.go
new file mode 100644
index 0000000..410af14
--- /dev/null
+++ b/wireguard-go-bridge/src/git.zx2c4.com/wireguard-go/queueconstants.go
@@ -0,0 +1,16 @@
+/* SPDX-License-Identifier: GPL-2.0
+ *
+ * Copyright (C) 2017-2018 WireGuard LLC. All Rights Reserved.
+ */
+
+package main
+
+/* Fit within memory limits for iOS */
+
+const (
+        QueueOutboundSize          = 1024
+        QueueInboundSize           = 1024
+        QueueHandshakeSize         = 1024
+        MaxSegmentSize             = 1700
+        PreallocatedBuffersPerPool = 1024
+)