From 6709495ee8af2a7cc9149fb36ff492e24ac7d33b Mon Sep 17 00:00:00 2001
From: David Fifield <david@bamsoftware.com>
Date: Mon, 4 Oct 2021 14:02:48 -0600
Subject: [PATCH] Client noisePacketConn.

---
 champa-client/main.go | 94 ++++++++++++++++++++++++++++++++++++++++++++++-----
 noise/noise.go        | 20 +++++++++++
 2 files changed, 105 insertions(+), 9 deletions(-)

diff --git a/champa-client/main.go b/champa-client/main.go
index 051db90..9e5f71e 100644
--- a/champa-client/main.go
+++ b/champa-client/main.go
@@ -33,6 +33,81 @@ func readKeyFromFile(filename string) ([]byte, error) {
 	return noise.ReadKey(f)
 }
 
+// noisePacketConn implements the net.PacketConn interface. It acts as an
+// intermediary between an upper layer and an inner net.PacketConn, decrypting
+// packets on ReadFrom and encrypting them on WriteTo.
+type noisePacketConn struct {
+	sess *noise.Session
+	net.PacketConn
+}
+
+// readNoiseMessageOfTypeFrom returns the first complete Noise message whose
+// msgTime is wantedType, discarding messages of any other msgType.
+func readNoiseMessageOfTypeFrom(conn net.PacketConn, wantedType byte) ([]byte, net.Addr, error) {
+	for {
+		msgType, msg, addr, err := noise.ReadMessageFrom(conn)
+		if err != nil {
+			if err, ok := err.(net.Error); ok && err.Temporary() {
+				continue
+			}
+			return nil, nil, err
+		}
+		if msgType == wantedType {
+			return msg, addr, nil
+		}
+	}
+}
+
+// noiseDial performs a Noise handshake over the given net.PacketConn, and
+// returns a noisePacketConn with a working noise.Session.
+func noiseDial(conn net.PacketConn, addr net.Addr, pubkey []byte) (*noisePacketConn, error) {
+	p := []byte{noise.MsgTypeHandshakeInit}
+	pre, p, err := noise.InitiateHandshake(p, pubkey)
+	if err != nil {
+		return nil, err
+	}
+	// TODO: timeout or context
+	_, err = conn.WriteTo(p, addr)
+	if err != nil {
+		return nil, err
+	}
+
+	msg, _, err := readNoiseMessageOfTypeFrom(conn, noise.MsgTypeHandshakeResp)
+	if err != nil {
+		return nil, err
+	}
+
+	sess, err := pre.FinishHandshake(msg)
+	if err != nil {
+		return nil, err
+	}
+
+	return &noisePacketConn{sess, conn}, nil
+}
+
+// ReadFrom implements the net.PacketConn interface for noisePacketConn.
+func (c *noisePacketConn) ReadFrom(p []byte) (int, net.Addr, error) {
+	msg, addr, err := readNoiseMessageOfTypeFrom(c.PacketConn, noise.MsgTypeTransport)
+	if err != nil {
+		return 0, nil, err
+	}
+	dec, err := c.sess.Decrypt(nil, msg)
+	if err != nil {
+		return 0, nil, err
+	}
+	return copy(p, dec), addr, nil
+}
+
+// WriteTo implements the net.PacketConn interface for noisePacketConn.
+func (c *noisePacketConn) WriteTo(p []byte, addr net.Addr) (int, error) {
+	buf := []byte{noise.MsgTypeTransport}
+	buf, err := c.sess.Encrypt(buf, p)
+	if err != nil {
+		return 0, err
+	}
+	return c.PacketConn.WriteTo(buf, addr)
+}
+
 func handle(local *net.TCPConn, sess *smux.Session, conv uint32) error {
 	stream, err := sess.OpenStream()
 	if err != nil {
@@ -91,8 +166,15 @@ func run(serverURL, cacheURL *url.URL, front, localAddr string, pubkey []byte) e
 	pconn := NewPollingPacketConn(turbotunnel.DummyAddr{}, poll)
 	defer pconn.Close()
 
-	// Open a KCP conn on the PacketConn.
-	conn, err := kcp.NewConn2(turbotunnel.DummyAddr{}, nil, 0, 0, pconn)
+	// Add a Noise layer over the AMP polling to encrypt and authenticate
+	// each KCP packet.
+	nconn, err := noiseDial(pconn, turbotunnel.DummyAddr{}, pubkey)
+	if err != nil {
+		return err
+	}
+
+	// Open a KCP conn over the Noise layer.
+	conn, err := kcp.NewConn2(turbotunnel.DummyAddr{}, nil, 0, 0, nconn)
 	if err != nil {
 		return fmt.Errorf("opening KCP conn: %v", err)
 	}
@@ -121,19 +203,13 @@ func run(serverURL, cacheURL *url.URL, front, localAddr string, pubkey []byte) e
 	// to permit one more packet per request, we should do it.
 	// E.g. 1400*5 = 7000, but 1320*6 = 7920.
 
-	// Put a Noise channel on top of the KCP conn.
-	rw, err := noise.NewClient(conn, pubkey)
-	if err != nil {
-		return err
-	}
-
 	// Start a smux session on the Noise channel.
 	smuxConfig := smux.DefaultConfig()
 	smuxConfig.Version = 2
 	smuxConfig.KeepAliveTimeout = idleTimeout
 	smuxConfig.MaxReceiveBuffer = 4 * 1024 * 1024 // default is 4 * 1024 * 1024
 	smuxConfig.MaxStreamBuffer = 1 * 1024 * 1024  // default is 65536
-	sess, err := smux.Client(rw, smuxConfig)
+	sess, err := smux.Client(conn, smuxConfig)
 	if err != nil {
 		return fmt.Errorf("opening smux session: %v", err)
 	}
diff --git a/noise/noise.go b/noise/noise.go
index 429d1d4..056903e 100644
--- a/noise/noise.go
+++ b/noise/noise.go
@@ -13,6 +13,7 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"net"
 	"strings"
 
 	"github.com/flynn/noise"
@@ -22,9 +23,28 @@ import (
 // The length of public and private keys as returned by GeneratePrivkey.
 const KeyLen = 32
 
+const (
+	MsgTypeHandshakeInit = 1
+	MsgTypeHandshakeResp = 2
+	MsgTypeTransport     = 4
+)
+
 // cipherSuite represents 25519_ChaChaPoly_BLAKE2s.
 var cipherSuite = noise.NewCipherSuite(noise.DH25519, noise.CipherChaChaPoly, noise.HashBLAKE2s)
 
+func ReadMessageFrom(conn net.PacketConn) (byte, []byte, net.Addr, error) {
+	var buf [1500]byte
+	for {
+		n, addr, err := conn.ReadFrom(buf[:])
+		if err != nil {
+			return 0, nil, nil, err
+		}
+		if n >= 1 {
+			return buf[0], buf[1:n], addr, nil
+		}
+	}
+}
+
 // readMessage reads a length-prefixed message from r. It returns a nil error
 // only when a complete message was read. It returns io.EOF only when there were
 // 0 bytes remaining to read from r. It returns io.ErrUnexpectedEOF when EOF
-- 
2.11.4.GIT