From 36dcab9300bbf8b0cf1aa53e1de6561344a3dc3c Mon Sep 17 00:00:00 2001
From: cathugger <cathugger@cock.li>
Date: Mon, 30 Jul 2018 01:58:52 +0000
Subject: [PATCH] optimize wire_put_uint64; use protokey for flowlabel
 fallback.

---
 src/yggdrasil/session.go | 15 ++++++++++++---
 src/yggdrasil/wire.go    | 20 ++++++++------------
 2 files changed, 20 insertions(+), 15 deletions(-)

diff --git a/src/yggdrasil/session.go b/src/yggdrasil/session.go
index 4a44887..7a1f19d 100644
--- a/src/yggdrasil/session.go
+++ b/src/yggdrasil/session.go
@@ -440,12 +440,21 @@ func (sinfo *sessionInfo) doSend(bs []byte) {
 		coords = append(coords, sinfo.coords...)            // Start with the real coords
 		coords = append(coords, 0)                          // Then target the local switchport
 		coords = wire_put_uint64(uint64(flowlabel), coords) // Then variable-length encoded flowlabel
+	} else if len(bs) >= 48 /* min UDP len, others are bigger */ &&
+		(bs[6] == 0x06 || bs[6] == 0x11 || bs[6] == 0x84) /* TCP UDP SCTP */ {
+		// if flowlabel was unspecified (0), try to use known protocols' ports
+		// protokey: proto | sport | dport
+		pkey := uint64(bs[6])<<32 /* proto */ |
+			uint64(bs[40])<<24 | uint64(bs[41])<<16 /* sport */ |
+			uint64(bs[42])<<8 | uint64(bs[43]) /* dport */
+		coords = append(coords, sinfo.coords...) // Start with the real coords
+		coords = append(coords, 0)               // Then target the local switchport
+		coords = wire_put_uint64(pkey, coords)   // Then variable-length encoded protokey
 	} else {
-		// 0 value means that flowlabels aren't being generated by OS.
+		// flowlabel was unspecified (0) and protocol unrecognised.
 		// To save bytes, we're not including it, therefore we won't need self-port override either.
 		// So just use sinfo.coords directly to avoid golang GC allocations.
-		// Recent enough Linux kernel supports flowlabels out of the box so this will be rare.
-		// XXX Attempt to look into TCP/UDP/SCTP/DCCP headers' sport/dport fields there?
+		// Recent enough Linux and BSDs support flowlabels (auto_flowlabel) out of the box so this will be rare.
 		coords = sinfo.coords
 	}
 	payload, nonce := boxSeal(&sinfo.sharedSesKey, bs, &sinfo.myNonce)
diff --git a/src/yggdrasil/wire.go b/src/yggdrasil/wire.go
index e92b4fc..d05624e 100644
--- a/src/yggdrasil/wire.go
+++ b/src/yggdrasil/wire.go
@@ -25,19 +25,15 @@ func wire_encode_uint64(elem uint64) []byte {
 
 // Encode uint64 using a variable length scheme.
 // Similar to binary.Uvarint, but big-endian.
-func wire_put_uint64(elem uint64, out []byte) []byte {
-	bs := make([]byte, 0, 10)
-	bs = append(bs, byte(elem&0x7f))
-	for e := elem >> 7; e > 0; e >>= 7 {
-		bs = append(bs, byte(e|0x80))
+func wire_put_uint64(e uint64, out []byte) []byte {
+	var b [10]byte
+	i := len(b) - 1
+	b[i] = byte(e & 0x7f)
+	for e >>= 7; e != 0; e >>= 7 {
+		i--
+		b[i] = byte(e | 0x80)
 	}
-	// Now reverse bytes, because we set them in the wrong order
-	// TODO just put them in the right place the first time...
-	last := len(bs) - 1
-	for idx := 0; idx < len(bs)/2; idx++ {
-		bs[idx], bs[last-idx] = bs[last-idx], bs[idx]
-	}
-	return append(out, bs...)
+	return append(out, b[i:]...)
 }
 
 // Returns the length of a wire encoded uint64 of this value.