From 4986cfe78b49173b8a9d86b9ec7e3024d0a3c218 Mon Sep 17 00:00:00 2001
From: Mathias Hall-Andersen <mathias@hall-andersen.dk>
Date: Tue, 22 Aug 2017 17:22:45 +0200
Subject: [PATCH] Add support for fwmark on linux

---
 src/config.go | 27 ++++++++++++++++++++++-----
 src/conn.go   | 16 +++++++++++++++-
 src/device.go |  7 ++++---
 src/tun.go    |  6 +++---
 4 files changed, 44 insertions(+), 12 deletions(-)

diff --git a/src/config.go b/src/config.go
index 474134b..871232c 100644
--- a/src/config.go
+++ b/src/config.go
@@ -145,10 +145,10 @@ func ipcSetOperation(device *Device, socket *bufio.ReadWriter) *IPCError {
 					return &IPCError{Code: ipcErrorInvalid}
 				}
 
-				netc := &device.net
-				netc.mutex.Lock()
-				netc.addr = addr
-				netc.mutex.Unlock()
+				device.net.mutex.Lock()
+				device.net.addr = addr
+				device.net.mutex.Unlock()
+
 				err = updateUDPConn(device)
 				if err != nil {
 					logError.Println("Failed to set listen_port:", err)
@@ -158,7 +158,24 @@ func ipcSetOperation(device *Device, socket *bufio.ReadWriter) *IPCError {
 				// TODO: Clear source address of all peers
 
 			case "fwmark":
-				logError.Println("FWMark not handled yet")
+				fwmark, err := strconv.ParseInt(value, 10, 32)
+				if err != nil {
+					logError.Println("Invalid fwmark", err)
+					return &IPCError{Code: ipcErrorInvalid}
+				}
+
+				device.net.mutex.Lock()
+				device.net.fwmark = int(fwmark)
+				err = setMark(
+					device.net.conn,
+					device.net.fwmark,
+				)
+				device.net.mutex.Unlock()
+				if err != nil {
+					logError.Println("Failed to set fwmark:", err)
+					return &IPCError{Code: ipcErrorIO}
+				}
+
 				// TODO: Clear source address of all peers
 
 			case "public_key":
diff --git a/src/conn.go b/src/conn.go
index e23b350..7b35829 100644
--- a/src/conn.go
+++ b/src/conn.go
@@ -13,6 +13,7 @@ func updateUDPConn(device *Device) error {
 
 	if netc.conn != nil {
 		netc.conn.Close()
+		netc.conn = nil
 	}
 
 	// open new connection
@@ -26,11 +27,24 @@ func updateUDPConn(device *Device) error {
 			return err
 		}
 
+		// set fwmark
+
+		err = setMark(netc.conn, netc.fwmark)
+		if err != nil {
+			return err
+		}
+
 		// retrieve port (may have been chosen by kernel)
 
 		addr := conn.LocalAddr()
 		netc.conn = conn
-		netc.addr, _ = net.ResolveUDPAddr(addr.Network(), addr.String())
+		netc.addr, _ = net.ResolveUDPAddr(
+			addr.Network(),
+			addr.String(),
+		)
+
+		// notify goroutines
+
 		signalSend(device.signal.newUDPConn)
 	}
 
diff --git a/src/device.go b/src/device.go
index 2a0d0ca..2ead768 100644
--- a/src/device.go
+++ b/src/device.go
@@ -21,9 +21,10 @@ type Device struct {
 		messageBuffers sync.Pool
 	}
 	net struct {
-		mutex sync.RWMutex
-		addr  *net.UDPAddr // UDP source address
-		conn  *net.UDPConn // UDP "connection"
+		mutex  sync.RWMutex
+		addr   *net.UDPAddr // UDP source address
+		conn   *net.UDPConn // UDP "connection"
+		fwmark int
 	}
 	mutex        sync.RWMutex
 	privateKey   NoisePrivateKey
diff --git a/src/tun.go b/src/tun.go
index b4fbc62..8e8c759 100644
--- a/src/tun.go
+++ b/src/tun.go
@@ -34,28 +34,28 @@ func (device *Device) RoutineTUNEventReader() {
 			if err != nil {
 				logError.Println("Failed to load updated MTU of device:", err)
 			} else if int(old) != mtu {
-				atomic.StoreInt32(&device.tun.mtu, int32(mtu))
 				if mtu+MessageTransportSize > MaxMessageSize {
 					logInfo.Println("MTU updated:", mtu, "(too large)")
 				} else {
 					logInfo.Println("MTU updated:", mtu)
 				}
+				atomic.StoreInt32(&device.tun.mtu, int32(mtu))
 			}
 		}
 
 		if event&TUNEventUp != 0 {
 			if !device.tun.isUp.Get() {
+				logInfo.Println("Interface set up")
 				device.tun.isUp.Set(true)
 				updateUDPConn(device)
-				logInfo.Println("Interface set up")
 			}
 		}
 
 		if event&TUNEventDown != 0 {
 			if device.tun.isUp.Get() {
+				logInfo.Println("Interface set down")
 				device.tun.isUp.Set(false)
 				closeUDPConn(device)
-				logInfo.Println("Interface set down")
 			}
 		}
 	}